#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from butterfly.dp import evaluate_policy, solve
from butterfly.environment import build_world
from butterfly.estimation import (
    build_support_gate,
    exact_trajectory_kl,
    model_errors_on_gate,
    sample_tabular,
    theorem_e_constant,
    transition_metrics,
)


def run_replica(
    env_seed: int,
    data_seed: int,
    sample_sizes: list[int],
    thresholds: list[int],
) -> list[dict[str, float | int]]:
    rows: list[dict[str, float | int]] = []
    world = build_world(env_seed, entropy="high", dispersion="high")
    true_ctrl = solve(world, "control")
    for n in sample_sizes:
        estimate = sample_tabular(
            world,
            n_trajectories=n,
            seed=100000 * env_seed + 1000 * data_seed + n,
        )
        learned_ctrl = solve(
            world,
            "control",
            event_kernels=estimate.mu_hat,
            action_rewards=estimate.action_rewards,
            event_rewards=estimate.event_rewards,
        )
        learned_gen = solve(
            world,
            "generative",
            event_kernels=estimate.mu_hat,
            action_rewards=estimate.action_rewards,
            event_rewards=estimate.event_rewards,
        )
        ctrl_true_value = evaluate_policy(world, learned_ctrl.policy)[world.start]
        gen_true_value = evaluate_policy(world, learned_gen.policy)[world.start]
        trajectory_kl = exact_trajectory_kl(world, estimate)

        for threshold in thresholds:
            weighted_tv, max_tv, chance_coverage = transition_metrics(
                world, estimate, threshold
            )
            gate = build_support_gate(world, estimate, threshold)
            true_gated = solve(world, "control", gate=gate)
            learned_gated = solve(
                world,
                "control",
                event_kernels=estimate.mu_hat,
                action_rewards=estimate.action_rewards,
                event_rewards=estimate.event_rewards,
                gate=gate,
            )
            learned_gated_true_value = evaluate_policy(world, learned_gated.policy)[
                world.start
            ]
            epsilon_r, epsilon_mu, r_bound = model_errors_on_gate(
                world, estimate, gate
            )
            ck = theorem_e_constant(world, epsilon_r, epsilon_mu, r_bound)
            rows.append(
                {
                    "env_seed": env_seed,
                    "data_seed": data_seed,
                    "n": n,
                    "threshold": threshold,
                    "trajectory_kl": trajectory_kl,
                    "weighted_tv": weighted_tv,
                    "max_supported_tv": max_tv,
                    "chance_coverage": chance_coverage,
                    "control_value_error": abs(
                        learned_ctrl.V[world.start] - true_ctrl.V[world.start]
                    ),
                    "generative_value_error": abs(
                        learned_gen.V[world.start] - true_ctrl.V[world.start]
                    ),
                    "control_policy_regret": true_ctrl.V[world.start]
                    - ctrl_true_value,
                    "generative_policy_regret": true_ctrl.V[world.start]
                    - gen_true_value,
                    "coverage_restriction_cost": true_ctrl.V[world.start]
                    - true_gated.V[world.start],
                    "within_gate_regret": true_gated.V[world.start]
                    - learned_gated_true_value,
                    "total_gated_regret": true_ctrl.V[world.start]
                    - learned_gated_true_value,
                    "epsilon_r": epsilon_r,
                    "epsilon_mu": epsilon_mu,
                    "c_k": ck,
                    "regret_bound_2c_k": 2.0 * ck,
                }
            )
    return rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--env-seeds", type=int, default=2)
    parser.add_argument("--data-seeds", type=int, default=2)
    parser.add_argument("--sample-sizes", default="100,300,1000,3000,10000,30000")
    parser.add_argument("--thresholds", default="1,5,10,20,50")
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--output", default=str(ROOT / "results" / "tabular.csv"))
    args = parser.parse_args()
    sample_sizes = [int(x) for x in args.sample_sizes.split(",")]
    thresholds = [int(x) for x in args.thresholds.split(",")]

    tasks = [
        (env_seed, data_seed)
        for env_seed in range(args.env_seeds)
        for data_seed in range(args.data_seeds)
    ]
    rows: list[dict[str, float | int]] = []
    if args.workers <= 1:
        for env_seed, data_seed in tasks:
            rows.extend(run_replica(env_seed, data_seed, sample_sizes, thresholds))
            print(f"completed env={env_seed} data_seed={data_seed}", flush=True)
    else:
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            futures = {
                executor.submit(
                    run_replica, env_seed, data_seed, sample_sizes, thresholds
                ): (env_seed, data_seed)
                for env_seed, data_seed in tasks
            }
            for future in as_completed(futures):
                env_seed, data_seed = futures[future]
                rows.extend(future.result())
                print(f"completed env={env_seed} data_seed={data_seed}", flush=True)

    rows.sort(key=lambda row: (row["env_seed"], row["data_seed"], row["n"], row["threshold"]))

    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    print(f"wrote {len(rows)} rows to {output}")


if __name__ == "__main__":
    main()
