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

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

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.neural import (
    extract_conditionals,
    neural_conditional_metrics,
    neural_exact_trajectory_kl,
    sample_token_sequences,
    train_transformer,
)


def run_model(
    env_seed: int,
    entropy: str,
    dispersions: list[str],
    train_seed: int,
    n: int,
    steps: int,
    batch_size: int,
    device: str,
    torch_threads: int,
) -> tuple[list[dict[str, float | int | str]], list[dict[str, float | int | str]]]:
    training_world = build_world(env_seed, entropy, dispersions[0])
    sequences = sample_token_sequences(
        training_world, n, seed=700000 + 1000 * env_seed
    )
    model, _pad, history = train_transformer(
        training_world,
        sequences,
        seed=900000 + 1000 * env_seed + train_seed,
        steps=steps,
        batch_size=batch_size,
        device=device,
        num_threads=torch_threads,
    )
    conditionals = extract_conditionals(training_world, model, device)
    trajectory_kl = neural_exact_trajectory_kl(training_world, conditionals)
    metrics = neural_conditional_metrics(training_world, conditionals)
    rows: list[dict[str, float | int | str]] = []
    for dispersion in dispersions:
        world = build_world(env_seed, entropy, dispersion)
        true_ctrl = solve(world, "control")
        learned_ctrl = solve(world, "control", event_kernels=conditionals.mu_hat)
        learned_gen = solve(world, "generative", event_kernels=conditionals.mu_hat)
        ctrl_true = evaluate_policy(world, learned_ctrl.policy)[world.start]
        gen_true = evaluate_policy(world, learned_gen.policy)[world.start]
        rows.append(
            {
                "env_seed": env_seed,
                "train_seed": train_seed,
                "entropy": entropy,
                "dispersion": dispersion,
                "n": n,
                "steps": steps,
                "final_train_loss": history[-1]["loss"],
                "trajectory_kl": trajectory_kl,
                **metrics,
                "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,
                "generative_policy_regret": true_ctrl.V[world.start] - gen_true,
            }
        )
    curves = [
        {
            "env_seed": env_seed,
            "train_seed": train_seed,
            "entropy": entropy,
            "dispersion": "shared",
            **item,
        }
        for item in history
    ]
    return rows, curves


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--env-seeds", type=int, default=1)
    parser.add_argument("--train-seeds", type=int, default=1)
    parser.add_argument("--n", type=int, default=5000)
    parser.add_argument("--steps", type=int, default=300)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--cells", default="high-high")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--torch-threads", type=int, default=8)
    parser.add_argument("--output", default=str(ROOT / "results" / "neural.csv"))
    args = parser.parse_args()
    cells = [tuple(cell.split("-")) for cell in args.cells.split(",")]

    rows = []
    curves = []
    dispersions_by_entropy = {}
    for entropy, dispersion in cells:
        dispersions_by_entropy.setdefault(entropy, []).append(dispersion)

    tasks = [
        (env_seed, entropy, dispersions, train_seed)
        for env_seed in range(args.env_seeds)
        for entropy, dispersions in dispersions_by_entropy.items()
        for train_seed in range(args.train_seeds)
    ]
    if args.workers > 1 and args.device != "cpu":
        raise ValueError("parallel neural execution currently supports CPU only")
    if args.workers <= 1:
        for env_seed, entropy, dispersions, train_seed in tasks:
            model_rows, model_curves = run_model(
                env_seed, entropy, dispersions, train_seed, args.n, args.steps,
                args.batch_size, args.device, args.torch_threads
            )
            rows.extend(model_rows)
            curves.extend(model_curves)
            print(
                f"completed env={env_seed} entropy={entropy} train_seed={train_seed}",
                flush=True,
            )
    else:
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            futures = {
                executor.submit(
                    run_model, env_seed, entropy, dispersions, train_seed,
                    args.n, args.steps, args.batch_size, args.device,
                    args.torch_threads
                ): (env_seed, entropy, train_seed)
                for env_seed, entropy, dispersions, train_seed in tasks
            }
            for future in as_completed(futures):
                env_seed, entropy, train_seed = futures[future]
                model_rows, model_curves = future.result()
                rows.extend(model_rows)
                curves.extend(model_curves)
                print(
                    f"completed env={env_seed} entropy={entropy} train_seed={train_seed}",
                    flush=True,
                )

    rows.sort(key=lambda row: (row["env_seed"], row["train_seed"], row["entropy"], row["dispersion"]))
    curves.sort(key=lambda row: (row["env_seed"], row["train_seed"], row["entropy"], row["step"]))

    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)
    curve_output = output.with_name(output.stem + "_curves.csv")
    with curve_output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(curves[0]))
        writer.writeheader()
        writer.writerows(curves)
    print(f"wrote {len(rows)} rows to {output}")
    print(f"wrote {len(curves)} curve rows to {curve_output}")


if __name__ == "__main__":
    main()
