"""Recompute the Dory-Beam-4B results tables and paired bootstrap CIs from the published eval outputs.

  pip install huggingface_hub pyarrow
  python audit_eval.py                                  # accuracy · mean tokens for every published configuration
  python audit_eval.py dory-beam-4b:2 dory1.0-4b:2      # paired difference A − B per benchmark, 95% bootstrap CI
  python audit_eval.py dory-beam-4b:2-effort-high dory1.0-4b:2

Reads eval/benchmarks/*.parquet from the ProCreations/Beam-Reasoning-Traces dataset (or --dir with local copies).
Accuracy averages every sampled answer; the CI resamples problems (each problem's samples stay together).
"""

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

import pyarrow.parquet as pq

BENCHES = ["gpqa", "math500", "aime24", "aime25", "amc23", "heldout"]


def load(directory):
    runs = defaultdict(list)
    for f in sorted(Path(directory).glob("*.parquet")):
        for r in pq.read_table(f).to_pylist():
            mode = "" if r["reasoning_mode"] == "normal" else f"-{r['reasoning_mode']}"
            runs[f"{r['model']}:{r['loops']}{mode}"].append(r)
    return runs


def per_problem(rows, bench):
    d = defaultdict(list)
    for r in rows:
        if r["bench"] == bench:
            d[r["problem_id"]].append(r["correct"])
    return {k: sum(v) / len(v) for k, v in d.items()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pair", nargs="*", help="two configurations to compare, e.g. dory-beam-4b:2 dory1.0-4b:2")
    ap.add_argument("--dir", help="local folder with the benchmark parquet files")
    ap.add_argument("--boot", type=int, default=4000)
    args = ap.parse_args()
    directory = args.dir
    if not directory:
        from huggingface_hub import snapshot_download
        directory = Path(snapshot_download("ProCreations/Beam-Reasoning-Traces", repo_type="dataset",
                                           allow_patterns=["eval/benchmarks/*.parquet"])) / "eval" / "benchmarks"
    runs = load(directory)

    if not args.pair:
        print(f"{'configuration':32s}" + "".join(f"{b:>16s}" for b in BENCHES))
        for name, rows in sorted(runs.items()):
            cells = []
            for b in BENCHES:
                rs = [r for r in rows if r["bench"] == b]
                acc = 100 * sum(r["correct"] for r in rs) / len(rs)
                cells.append(f"{acc:6.1f} · {sum(r['output_tokens'] for r in rs) / len(rs) / 1000:4.1f}k")
            print(f"{name:32s}" + "".join(f"{c:>16s}" for c in cells))
        return

    a, b = args.pair
    print(f"{a} minus {b}: accuracy difference in points, 95% paired bootstrap CI over problems")
    rng = random.Random(0)
    for bench in BENCHES:
        pa, pb = per_problem(runs[a], bench), per_problem(runs[b], bench)
        ids = sorted(pa)
        diff = sum(pa[i] - pb[i] for i in ids) / len(ids)
        boots = sorted(sum(pa[i] - pb[i] for i in s) / len(s)
                       for s in ([rng.choice(ids) for _ in ids] for _ in range(args.boot)))
        lo, hi = boots[int(0.025 * args.boot)], boots[int(0.975 * args.boot)]
        flag = "" if lo < 0 < hi else "  *"
        print(f"  {bench:8s} {100 * diff:+6.1f}  [{100 * lo:+6.1f}, {100 * hi:+6.1f}]{flag}")


if __name__ == "__main__":
    main()