"""Compare the coding-fix blend (rC-t050) with Dory-Beam-4B v2 (rE-t050) and base Dory, on paired samples.

  python3 scripts/compare_rc.py [--efforts default,low,high]   # -> data/eval/rc_report.md

Math: data/eval/results/<model>_loops2-<effort>-64k.jsonl (scripts/summarize_effort.py), per benchmark.
Code: data/eval/code/rc-20261031/rC-t050-<effort>/results.jsonl against full-20261030 (same per-problem seeds).
Each gap has a 95% paired bootstrap CI over problems (4,000 resamples). The ship gate: coding plus-test pass@1
better than v2 on HumanEval+ and MBPP+ combined, and no math benchmark significantly worse than v2.
"""

import argparse
import json
import random
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
RES = ROOT / "data" / "eval" / "results"
CODE = ROOT / "data" / "eval" / "code"
BENCHES = ["gpqa", "math500", "aime24", "aime25", "amc23", "heldout"]
BOOTS = 4000


def per_problem(path, key, value):
    d = defaultdict(list)
    for r in map(json.loads, open(path)):
        d[key(r)].append(value(r))
    return {k: sum(v) / len(v) for k, v in d.items()}


def gap(a, b, rng):
    keys = sorted(set(a) & set(b))
    diffs = [a[k] - b[k] for k in keys]
    mean = 100 * sum(diffs) / len(diffs)
    boots = sorted(100 * sum(rng.choices(diffs, k=len(diffs))) / len(diffs) for _ in range(BOOTS))
    return mean, boots[BOOTS // 40], boots[BOOTS - BOOTS // 40 - 1], len(keys)


def fmt(m, lo, hi):
    flag = " ▲" if lo > 0 else " ▼" if hi < 0 else ""
    return f"{m:+.1f} [{lo:+.1f}, {hi:+.1f}]{flag}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--efforts", default="default")
    args = ap.parse_args()
    rng = random.Random(0)
    lines = ["# rC-t050 (coding fix) vs Dory-Beam-4B v2 (rE-t050), 2 loops", "",
             "▲/▼ = 95% CI excludes zero.", ""]
    regress = []
    for eff in args.efforts.split(","):
        new, old = RES / f"rC-t050_loops2-{eff}-64k.jsonl", RES / f"rE-t050_loops2-{eff}-64k.jsonl"
        if new.exists():
            lines += [f"## Math, effort {eff}", "", "| Benchmark | v2 | rC | rC − v2 (95% CI) | tokens v2 → rC |",
                      "|---|---|---|---|---|"]
            rows_new = [json.loads(l) for l in open(new)]
            rows_old = [json.loads(l) for l in open(old)]
            for b in BENCHES + ["all"]:
                sel = (lambda r: True) if b == "all" else (lambda r, b=b: r["bench"] == b)
                pn = defaultdict(list); po = defaultdict(list)
                for r in rows_new:
                    if sel(r):
                        pn[(r["bench"], r["id"])].append(r["correct"])
                for r in rows_old:
                    if sel(r):
                        po[(r["bench"], r["id"])].append(r["correct"])
                a = {k: sum(v) / len(v) for k, v in pn.items()}
                o = {k: sum(v) / len(v) for k, v in po.items()}
                m, lo, hi, n = gap(a, o, rng)
                tn = sum(r["tokens"] for r in rows_new if sel(r)) / max(1, sum(1 for r in rows_new if sel(r)))
                to = sum(r["tokens"] for r in rows_old if sel(r)) / max(1, sum(1 for r in rows_old if sel(r)))
                acc_n = 100 * sum(r["correct"] for r in rows_new if sel(r)) / sum(1 for r in rows_new if sel(r))
                acc_o = 100 * sum(r["correct"] for r in rows_old if sel(r)) / sum(1 for r in rows_old if sel(r))
                lines.append(f"| {b} | {acc_o:.1f} | {acc_n:.1f} | {fmt(m, lo, hi)} | {to / 1000:.1f}k → {tn / 1000:.1f}k |")
                if hi < 0 and b != "all":
                    regress.append(f"math {b} ({eff})")
            lines.append("")
        cnew = CODE / "rc-20261031" / f"rC-t050-{eff}" / "results.jsonl"
        if cnew.exists():
            lines += [f"## Code, effort {eff} (pass@1, plus tests)", "",
                      "| Benchmark | base | v2 | rC | rC − v2 (95% CI) | rC − base (95% CI) |", "|---|---|---|---|---|---|"]
            for bench in ("humaneval", "mbpp"):
                def get(path, bench=bench):
                    d = per_problem(path, lambda r: (r["bench"], r["task_id"]), lambda r: r["plus_pass"])
                    return {k: v for k, v in d.items() if k[0] == bench}
                n_ = get(cnew)
                v2 = get(CODE / "full-20261030" / f"rE-t050-{eff}" / "results.jsonl")
                bs = get(CODE / "full-20261030" / "base" / "results.jsonl")
                m1 = gap(n_, v2, rng)
                m2 = gap(n_, bs, rng)
                avg = lambda d: 100 * sum(d.values()) / len(d)
                lines.append(f"| {bench} | {avg(bs):.1f} | {avg(v2):.1f} | {avg(n_):.1f} | {fmt(*m1[:3])} | {fmt(*m2[:3])} |")
            lines.append("")
    lines += ["## Gate", "", "Significant math regressions vs v2: " + (", ".join(regress) if regress else "none"), ""]
    out = ROOT / "data" / "eval" / "rc_report.md"
    out.write_text("\n".join(lines))
    print("\n".join(lines))


if __name__ == "__main__":
    main()