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

import argparse
import csv
import os
import sys
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.identities import (
    audit_reconstruction_residual,
    factorization_residual,
    kl_regret_identity_residual,
    max_gap_recursion_residual,
    trajectory_decomposition_residual,
)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--seeds", type=int, default=20)
    parser.add_argument("--output", default=str(ROOT / "results" / "exact_factorial.csv"))
    args = parser.parse_args()

    rows = []
    for seed in range(args.seeds):
        for entropy in ("low", "high"):
            for dispersion in ("low", "high"):
                world = build_world(seed, entropy, dispersion)
                ctrl = solve(world, "control")
                gen = solve(world, "generative")
                gen_value = evaluate_policy(world, gen.policy)[world.start]
                factorization, normalization = factorization_residual(world, gen)
                relative, counting = audit_reconstruction_residual(world)
                path, event_mean = trajectory_decomposition_residual(world, ctrl)
                behavior = {
                    state: world.behavior_policy(state)
                    for state in world.decision_states(include_terminal=False)
                }
                rows.append(
                    {
                        "seed": seed,
                        "entropy": entropy,
                        "dispersion": dispersion,
                        "mean_root_event_entropy": float(
                            np.mean(
                                [
                                    -np.sum(
                                        world.event_kernel(world.start, action)
                                        * np.log(world.event_kernel(world.start, action))
                                    )
                                    for action in range(world.n_actions)
                                ]
                            )
                        ),
                        "v_control": ctrl.V[world.start],
                        "v_generative": gen.V[world.start],
                        "root_gap": gen.V[world.start] - ctrl.V[world.start],
                        "generative_policy_regret": ctrl.V[world.start] - gen_value,
                        "max_gap_recursion_residual": max_gap_recursion_residual(world, gen, ctrl),
                        "factorization_residual": factorization,
                        "normalization_residual": normalization,
                        "relative_audit_residual": relative,
                        "counting_audit_residual": counting,
                        "path_decomposition_residual": path,
                        "event_mean_residual": event_mean,
                        "kl_regret_identity_residual": kl_regret_identity_residual(
                            world, ctrl, behavior
                        ),
                    }
                )

    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()

