#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Rounds curve: accuracy and dispersion of the committee by deliberation round, from the debate transcripts
(round 0 = independent estimates, round r = revisions). Writes results/evaluation/rounds.csv
(columns: round, value, ci_lo, ci_hi, arm, metric, n) consumed by scripts/21_figures.py.

Usage: python3 scripts/35_rounds_curve.py --split validation [--runs validation_G_debate_r4_mixed_random_mid ...]
"""
from __future__ import annotations

import argparse
import json
import math
import re
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from agentic_valuation.config import PROCESSED, RESULTS  # noqa: E402


def per_round_points(tr: dict) -> dict[int, dict[str, float]]:
    out: dict[int, dict[str, float]] = {}
    for turn in tr["turns"]:
        if turn["kind"] in ("estimate", "revision") and turn.get("estimate"):
            out.setdefault(int(turn["round"]), {})[turn["role"]] = float(turn["estimate"]["point_estimate"])
    return out


def block_ci(vals: np.ndarray, blocks: np.ndarray, B: int = 2000, seed: int = 20261027) -> tuple[float, float]:
    rng = np.random.default_rng(seed)
    ub = np.unique(blocks)
    idx = {b: np.where(blocks == b)[0] for b in ub}
    stats = []
    for _ in range(B):
        draw = rng.choice(ub, size=len(ub), replace=True)
        sel = np.concatenate([idx[b] for b in draw])
        stats.append(vals[sel].mean())
    return float(np.percentile(stats, 2.5)), float(np.percentile(stats, 97.5))


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", default="validation")
    ap.add_argument("--runs", nargs="*", default=None)
    ap.add_argument("--out", default=None)
    a = ap.parse_args(argv)
    truth = pl.read_parquet(PROCESSED / "analytic.parquet", columns=["tx_id", "price", "mun_code", "date"]).with_columns(pl.col("date").dt.strftime("%Y-%m").alias("ym"))
    tmap = {r["tx_id"]: (math.log(r["price"]), f"{r['mun_code']}_{r['ym']}") for r in truth.iter_rows(named=True)}
    runs = a.runs or [d.name for d in sorted((RESULTS / "runs").glob(f"{a.split}_G_*")) + sorted((RESULTS / "runs").glob(f"final_G_*")) if a.split in d.name]
    rows = []
    for run in runs:
        d = RESULTS / "runs" / run
        recs = []
        for f in sorted((d / "transcripts").glob("*.json")):
            tr = json.loads(f.read_text())
            if tr["tx_id"] not in tmap:
                continue
            lp, blk = tmap[tr["tx_id"]]
            prp = per_round_points(tr)
            m_r = re.search(r"_r(\d+)_", run)
            max_r = int(m_r.group(1)) if m_r else int(tr.get("rounds", max(prp)))  # nominal rounds of the arm
            # carry forward: a committee that stopped early (IQR rule) keeps its last estimates in later rounds,
            # so every round averages the SAME set of properties
            last: dict[str, float] = {}
            for r in range(0, max(max_r, max(prp)) + 1):
                if r in prp:
                    last = {**last, **prp[r]}
                if not last:
                    continue
                v = np.log(np.array(list(last.values())))
                recs.append({"tx_id": tr["tx_id"], "round": r, "ae_median": abs(float(np.median(v)) - lp), "ae_mean": abs(float(v.mean()) - lp), "iqr": float(np.subtract(*np.percentile(v, [75, 25]))), "block": blk, "n_roles": len(v), "carried": r not in prp})
        if not recs:
            continue
        df = pd.DataFrame(recs)
        man = json.loads((d / "manifest.json").read_text()) if (d / "manifest.json").exists() else {}
        arm = man.get("arm_label") or "_".join(run.split("_")[1:4])
        committee = man.get("committee") or ("mixed_random" if "mixed_random" in run else "all_anthropic" if "all_anthropic" in run else "all_openai")
        for r, g in df.groupby("round"):
            for metric in ("ae_median", "ae_mean", "iqr"):
                lo, hi = block_ci(g[metric].to_numpy(), g["block"].to_numpy())
                rows.append({"round": int(r), "value": float(g[metric].mean()), "ci_lo": lo, "ci_hi": hi, "arm": arm, "committee": committee, "metric": metric, "n": int(len(g)), "run_id": run})
    out = Path(a.out) if a.out else RESULTS / "evaluation" / f"rounds_{a.split}.csv"
    out.parent.mkdir(parents=True, exist_ok=True)
    res = pd.DataFrame(rows)
    res.to_csv(out, index=False)
    if a.split == "test":
        res[res.metric == "ae_median"].to_csv(RESULTS / "evaluation" / "rounds.csv", index=False)
    print(res[res.metric.isin(["ae_median", "iqr"])].pivot_table(index=["arm", "round"], columns="metric", values="value").round(4))


if __name__ == "__main__":
    main()