"""Beam's reasoning-effort dial on the evaluation set, next to Dory-Beam-4B's.

  .venv/bin/python -I scripts/beam_effort_study.py   # -> data/beam_effort/STUDY.md, data/beam_effort/summary.json

Beam (self-hosted NVFP4, temperature 0.7, top-p 0.99, up to 28k tokens) answered the same 1,836 eval samples at
every reasoning_effort: minimal, low, medium, high, xhigh and max. Medium is data/beam_eval/run-evalbeam2-*, the
others data/beam_eval/run-evaleffort-*. Nothing is selected on correctness. Accuracy gaps against medium are paired by
problem and sample, with 95% bootstrap CIs over problems (4,000 resamples).
"""

import glob
import json
import random
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
from verify import check  # noqa: E402

LEVELS = ["minimal", "low", "medium", "high", "xhigh", "max"]
BENCHES = ["gpqa", "math500", "aime24", "aime25", "amc23", "heldout"]
NAMES = {"gpqa": "GPQA-D", "math500": "MATH-500", "aime24": "AIME 2024", "aime25": "AIME 2025", "amc23": "AMC 2023",
         "heldout": "Held-out", "all": "All"}
DORY = [("low", "rE-t050_loops2-low-64k"), ("medium", "rE-t050_loops2-medium-64k"),
        ("default", "rE-t050_loops2-default-64k"), ("high", "rE-t050_loops2-high-64k")]
BOOTS = 4000


def load_beam():
    bench = {json.loads(l)["id"]: json.loads(l)["bench"] for l in open(ROOT / "data" / "evalsets.jsonl")}
    rows = {}
    for pattern in ("run-evalbeam2-*.jsonl", "run-evaleffort-*.jsonl"):
        for path in sorted(glob.glob(str(ROOT / "data" / "beam_eval" / pattern))):
            for line in open(path):
                r = json.loads(line)
                if "error" in r:
                    continue
                ok, _ = check(r["content"], r["gold"], r["answer_type"])
                rows[(r["effort"], r["id"], r["sample_idx"])] = {
                    "bench": bench[r["id"]], "correct": bool(ok and r["finish_reason"] == "stop"),
                    "tokens": r["usage"]["completion_tokens"], "truncated": r["finish_reason"] != "stop"}
    return rows


def summarize(rows, level, bench):
    rs = [v for (e, _, _), v in rows.items() if e == level and (bench == "all" or v["bench"] == bench)]
    if not rs:
        return None
    return {"n": len(rs), "acc": 100 * sum(r["correct"] for r in rs) / len(rs),
            "tokens": sum(r["tokens"] for r in rs) / len(rs), "truncated": 100 * sum(r["truncated"] for r in rs) / len(rs)}


def paired_gap(rows, level, ref, bench, rng):
    """Accuracy of level minus ref, on the samples both have; CI by resampling problems."""
    per = defaultdict(lambda: [0, 0, 0])
    for (e, pid, s), v in rows.items():
        if e != level or (bench != "all" and v["bench"] != bench):
            continue
        o = rows.get((ref, pid, s))
        if o is None:
            continue
        p = per[pid]
        p[0] += v["correct"]; p[1] += o["correct"]; p[2] += 1
    # sample-weighted like the accuracy column (AIME problems carry 8 samples), resampling whole problems
    probs = list(per.values())
    if not probs:
        return None
    gap = lambda ps: 100 * sum(a - b for a, b, _ in ps) / sum(n for _, _, n in ps)
    boots = sorted(gap(rng.choices(probs, k=len(probs))) for _ in range(BOOTS))
    return gap(probs), boots[BOOTS // 40], boots[BOOTS - BOOTS // 40 - 1]


def dory(tag, bench):
    rs = [json.loads(l) for l in open(ROOT / "data" / "eval" / "results" / f"{tag}.jsonl")]
    rs = [r for r in rs if bench == "all" or r["bench"] == bench]
    return {"acc": 100 * sum(r["correct"] for r in rs) / len(rs), "tokens": sum(r["tokens"] for r in rs) / len(rs)}


def main():
    rows = load_beam()
    rng = random.Random(0)
    summary = {"beam": {}, "dory_beam": {}, "base": {}}
    lines = ["# Beam's reasoning-effort dial on the evaluation set", "",
             "Accuracy % · mean output tokens. Gap = level minus medium, paired, 95% CI (▲/▼ excludes zero).", "",
             "| Level | " + " | ".join(NAMES[b] for b in BENCHES + ["all"]) + " | Gap vs medium (all) | Cut off at 28k |",
             "|---" * (len(BENCHES) + 4) + "|"]
    for lv in LEVELS:
        cells, summary["beam"][lv] = [], {}
        for b in BENCHES + ["all"]:
            s = summarize(rows, lv, b)
            summary["beam"][lv][b] = s
            cells.append("—" if s is None else f"{s['acc']:.1f} · {s['tokens'] / 1000:.1f}k")
        g = None if lv == "medium" else paired_gap(rows, lv, "medium", "all", rng)
        gap = "—" if g is None else f"{g[0]:+.1f} [{g[1]:+.1f}, {g[2]:+.1f}]" + (" ▲" if g[1] > 0 else " ▼" if g[2] < 0 else "")
        s = summary["beam"][lv]["all"]
        cut = "—" if s is None else f"{s['truncated']:.1f}%"
        lines.append(f"| **{lv}** | " + " | ".join(cells) + f" | {gap} | {cut} |")
    lines += ["", "Per-benchmark gaps vs medium (paired, 95% CI):", "", "| Level | " + " | ".join(NAMES[b] for b in BENCHES) + " |",
              "|---" * (len(BENCHES) + 1) + "|"]
    for lv in LEVELS:
        if lv == "medium":
            continue
        cells = []
        for b in BENCHES:
            g = paired_gap(rows, lv, "medium", b, rng)
            cells.append("—" if g is None else f"{g[0]:+.1f} [{g[1]:+.1f}, {g[2]:+.1f}]" + (" ▲" if g[1] > 0 else " ▼" if g[2] < 0 else ""))
        lines.append(f"| {lv} | " + " | ".join(cells) + " |")
    lines += ["", "Dory-Beam-4B (2 loops, temperature 1.0, 64k) on the same samples:", "",
              "| Effort | " + " | ".join(NAMES[b] for b in BENCHES + ["all"]) + " |", "|---" * (len(BENCHES) + 2) + "|"]
    for name, tag in DORY:
        summary["dory_beam"][name] = {b: dory(tag, b) for b in BENCHES + ["all"]}
        lines.append(f"| {name} | " + " | ".join(f"{summary['dory_beam'][name][b]['acc']:.1f} · "
                                                 f"{summary['dory_beam'][name][b]['tokens'] / 1000:.1f}k" for b in BENCHES + ["all"]) + " |")
    summary["base"] = {b: dory("base_loops2-64k", b) for b in BENCHES + ["all"]}
    # grouped points for scripts/make_effort_chart.py: [mean tokens, accuracy]
    def pts(rs):
        return [sum(r["tokens"] for r in rs) / len(rs), 100 * sum(r["correct"] for r in rs) / len(rs)]
    groups = {"all": None, "aime": {"aime24", "aime25"}, "routine": {"math500", "amc23"}}
    res = lambda tag: [json.loads(l) for l in open(ROOT / "data" / "eval" / "results" / f"{tag}.jsonl")]
    summary["chart"] = {}
    for g, bs in groups.items():
        keep = lambda r: bs is None or r["bench"] in bs
        summary["chart"][g] = {
            "beam": {lv: pts([v for (e, _, _), v in rows.items() if e == lv and keep(v)]) for lv in LEVELS},
            "dory_beam": {name: pts([r for r in res(tag) if keep(r)]) for name, tag in DORY if name != "default"},
            "base": pts([r for r in res("base_loops2-64k") if keep(r)])}
    out = ROOT / "data" / "beam_effort"
    out.mkdir(parents=True, exist_ok=True)
    (out / "STUDY.md").write_text("\n".join(lines) + "\n")
    (out / "summary.json").write_text(json.dumps(summary, indent=1))
    print("\n".join(lines))


if __name__ == "__main__":
    main()