"""Summarize budget forcing (modal_app/eval_dory.py::budget_force) against the hard cap, at 2 loops.

  python3 scripts/summarize_budget.py   # reads data/eval/budget/<tag>.jsonl, writes data/eval/budget/REPORT.md

Hard cap: a sample counts only if it finished on its own within the cap and was correct. Forced: samples still
thinking at the cap get "\n</think>" and up to 1,024 answer tokens; correct only if that answer ends and checks out.
Overall accuracy is over all 1,836 samples; CIs are a paired bootstrap over the 1,098 problems (4,000 resamples).
"""

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

ROOT = Path(__file__).resolve().parent.parent
DIR = ROOT / "data" / "eval" / "budget"
CONFIGS = [("base_loops2-64k", "Base Dory"), ("base_loops2eff", "Base, efficient mode"),
           ("rE-t050_loops2-low-64k", "Dory-Beam low"), ("rE-t050_loops2-default-64k", "Dory-Beam default"),
           ("rE-t050_loops2-high-64k", "Dory-Beam high"), ("mix030_loops2-64k", "Dory-Beam-Accurate")]
BENCHES = ["gpqa", "math500", "aime24", "aime25", "amc23", "heldout"]
CAPS = [2048, 4096, 8192, 16384]
BOOTS = 4000


def load(tag):
    """{cap: {(bench, id): [(hard, forced, tokens) per sample]}}"""
    d = defaultdict(lambda: defaultdict(list))
    for line in open(DIR / f"{tag}.jsonl"):
        r = json.loads(line)
        tokens = r["think_tokens"] + (r["answer_tokens"] or 0)
        d[r["cap"]][(r["bench"], r["id"])].append((not r["forced"] and r["correct"], r["correct"], tokens))
    return d


def acc(per_problem, keys, i):
    return 100 * sum(x[i] for k in keys for x in per_problem[k]) / sum(len(per_problem[k]) for k in keys)


def diff_ci(a, b, keys, i, rng):
    """Paired bootstrap over problems of acc(a) - acc(b), samples pooled within each resampled problem."""
    diffs = []
    for _ in range(BOOTS):
        ks = rng.choices(keys, k=len(keys))
        diffs.append(acc(a, ks, i) - acc(b, ks, i))
    diffs.sort()
    return diffs[BOOTS // 40], diffs[BOOTS - BOOTS // 40 - 1]


def main():
    data = {tag: load(tag) for tag, _ in CONFIGS}
    keys = sorted(data[CONFIGS[0][0]][CAPS[0]])
    rng = random.Random(0)
    lines = ["# Budget forcing vs hard cap (2 loops, 1,836 samples over 1,098 problems)", ""]

    lines += ["## Overall accuracy: hard cap → forced (mean tokens actually used when forced)", "",
              "| Cap | " + " | ".join(n for _, n in CONFIGS) + " |", "|---" * (len(CONFIGS) + 1) + "|"]
    for cap in CAPS:
        row = [f"{cap // 1024}k"]
        for tag, _ in CONFIGS:
            p = data[tag][cap]
            tok = sum(x[2] for k in keys for x in p[k]) / sum(len(p[k]) for k in keys)
            row.append(f"{acc(p, keys, 0):.1f} → **{acc(p, keys, 1):.1f}** ({tok / 1000:.1f}k)")
        lines.append("| " + " | ".join(row) + " |")

    lines += ["", "## Forced accuracy minus base Dory (forced), paired 95% CI", "",
              "| Cap | " + " | ".join(n for _, n in CONFIGS[1:]) + " |", "|---" * len(CONFIGS) + "|"]
    base = data[CONFIGS[0][0]]
    for cap in CAPS:
        row = [f"{cap // 1024}k"]
        for tag, _ in CONFIGS[1:]:
            p = data[tag][cap]
            d = acc(p, keys, 1) - acc(base[cap], keys, 1)
            lo, hi = diff_ci(p, base[cap], keys, 1, rng)
            row.append(f"{d:+.1f} [{lo:+.1f}, {hi:+.1f}]")
        lines.append("| " + " | ".join(row) + " |")

    lines += ["", "## Forced minus hard cap, same config, paired 95% CI", "",
              "| Cap | " + " | ".join(n for _, n in CONFIGS) + " |", "|---" * (len(CONFIGS) + 1) + "|"]
    for cap in CAPS:
        row = [f"{cap // 1024}k"]
        for tag, _ in CONFIGS:
            p = data[tag][cap]
            hard = {k: [(x[1], x[0]) for x in v] for k, v in p.items()}  # swap so index 0 = forced
            d = acc(p, keys, 1) - acc(p, keys, 0)
            lo, hi = diff_ci(hard, p, keys, 0, rng)
            row.append(f"{d:+.1f} [{lo:+.1f}, {hi:+.1f}]")
        lines.append("| " + " | ".join(row) + " |")

    lines += ["", "## Forced accuracy per benchmark", ""]
    for cap in CAPS:
        lines += [f"**{cap // 1024}k**", "", "| Config | " + " | ".join(BENCHES) + " |", "|---" * (len(BENCHES) + 1) + "|"]
        for tag, name in CONFIGS:
            p = data[tag][cap]
            lines.append(f"| {name} | " + " | ".join(f"{acc(p, [k for k in keys if k[0] == b], 1):.1f}"
                                                     for b in BENCHES) + " |")
        lines.append("")

    (DIR / "REPORT.md").write_text("\n".join(lines) + "\n")
    print("\n".join(lines))


if __name__ == "__main__":
    main()