"""Compare two graded runs over the same items: per-task accuracy (mean over samples), item-level paired difference
with a bootstrap 95% CI, plus mean generated tokens.  usage: cmp.py <tagA> <tagB> [--file graded.jsonl]"""
import collections
import json
import random
import sys


def load(tag):
    d = collections.defaultdict(list)
    for g in map(json.loads, open(f"results/{tag}/graded.jsonl", encoding="utf-8")):
        t = "code" if g["task"].startswith("code") else g["task"]
        d[(t, g["id"])].append((g["correct"], g["n_tokens"]))
    return d


a, b = load(sys.argv[1]), load(sys.argv[2])
keys = [k for k in a if k in b]
rng = random.Random(0)
for task in sorted({k[0] for k in keys}) + ["ALL"]:
    ks = [k for k in keys if task in ("ALL", k[0])]
    da = [sum(c for c, _ in a[k]) / len(a[k]) for k in ks]
    db = [sum(c for c, _ in b[k]) / len(b[k]) for k in ks]
    diff = [y - x for x, y in zip(da, db)]
    boots = sorted(sum(rng.choice(diff) for _ in diff) / len(diff) for _ in range(2000))
    ta = sum(n for k in ks for _, n in a[k]) / sum(len(a[k]) for k in ks)
    tb = sum(n for k in ks for _, n in b[k]) / sum(len(b[k]) for k in ks)
    print(f"{task:8s} n={len(ks):3d}  {100 * sum(da) / len(ks):5.1f} -> {100 * sum(db) / len(ks):5.1f}  "
          f"d={100 * sum(diff) / len(ks):+5.1f} [{100 * boots[50]:+.1f}, {100 * boots[1949]:+.1f}]  "
          f"tok {ta:.0f} -> {tb:.0f}")