"""Run the implemented exact and neural comparisons from a YAML config."""

import argparse, json, time, platform
from pathlib import Path
import numpy as np, pandas as pd, torch, yaml
from dooable.graph import Graph, toy_graph, grid_graph, string_graph
from dooable.exact import (
    solve,
    uniform_policy,
    tilted_reference,
    backward_policy,
    forward_from_backward,
    sample,
)
from dooable.learning import train
from dooable.metrics import graph_metrics


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--config", required=True)
    a = p.parse_args()
    cfg = yaml.safe_load(Path(a.config).read_text())
    out = Path(cfg["output"])
    out.mkdir(parents=True, exist_ok=True)
    if cfg.get("graph"):
        g = Graph.load(cfg["graph"])
    elif cfg.get("kind") == "grid":
        g = grid_graph(cfg.get("width", 5), cfg.get("budget", 6))
    elif cfg.get("kind") == "strings":
        g = string_graph(cfg.get("length", 4), cfg.get("budget", 2))
    else:
        g = toy_graph(cfg.get("multiplicity", 8))
    rewards = (
        json.loads(Path(cfg["rewards"]).read_text())
        if cfg.get("rewards")
        else {y: 0.0 for y in g.terminals}
    )
    temperature = cfg.get("temperature", 0.7)
    rows = []
    g.save(out / "graph.json")
    for seed in cfg.get("seeds", [0, 1, 2, 3, 4]):
        for name in cfg["methods"]:
            start = time.perf_counter()
            if name == "exact":
                forward = solve(g, rewards, temperature).forward
            elif name == "uniform":
                forward = uniform_policy(g)
            elif name == "reference_tilt":
                forward = tilted_reference(g, rewards)
            elif name == "zero_cost":
                from dooable.ablations import zero_cost_policy

                forward = zero_cost_policy(g, rewards)
            elif name == "duplicate_endpoints":
                from dooable.ablations import duplicate_endpoint_policy

                forward = duplicate_endpoint_policy(g, rewards, temperature)
            elif name in ["dooable", "tb_uniform", "tb_exact", "unnormalized"]:
                backward = {
                    "dooable": "learned",
                    "tb_uniform": "uniform",
                    "tb_exact": "exact",
                    "unnormalized": "unnormalized",
                }[name]
                model, _ = train(
                    g,
                    rewards,
                    temperature,
                    steps=cfg.get("steps", 2000),
                    batch_size=cfg.get("batch_size", 64),
                    seed=seed,
                    output=out / f"{name}_seed{seed}",
                    backward=backward,
                )
                forward = model.probabilities()
            else:
                raise ValueError(f"Unimplemented comparator {name}")
            row = {
                "method": name,
                "seed": seed,
                "seconds": time.perf_counter() - start,
                "nodes": len(g.nodes),
                "edges": len(g.edges),
                "outcomes": len(g.terminals),
                **graph_metrics(g, forward, rewards, temperature),
            }
            if g.metadata.get("kind") == "reaction":
                from dooable.chemistry import replay

                paths = sample(g, forward, cfg.get("samples", 1000), seed)
                row["replay_fraction"] = np.mean(
                    [replay(r, g.metadata["budget"]) for r in paths]
                )
                row["unique_outcomes"] = len({r["outcome"] for r in paths})
                (out / f"{name}_seed{seed}_samples.jsonl").write_text(
                    "".join(json.dumps(r) + "\n" for r in paths)
                )
            rows.append(row)
            pd.DataFrame(rows).to_csv(out / "metrics.csv", index=False)
            print(json.dumps(row), flush=True)
    df = pd.DataFrame(rows)
    numeric = [c for c in df.select_dtypes("number").columns if c != "seed"]
    df.groupby("method")[numeric].agg(["mean", "sem"]).to_csv(out / "summary.csv")
    (out / "run.json").write_text(
        json.dumps(
            {
                "config": cfg,
                "python": platform.python_version(),
                "torch": torch.__version__,
                "numpy": np.__version__,
                "platform": platform.platform(),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()