# E20 analysis: base vs tuned on MATH-500, exact McNemar, per subject and level, tokens. Writes logs/e20_result.json.
import json, sys
from math import comb
from collections import defaultdict
sys.path.insert(0, "."); from grade import load
b = {r["item"]: r for r in load("out/base_math500.jsonl") if r["sample"] == 0}
t = {r["item"]: r for r in load("out/tuned_math500.jsonl") if r["sample"] == 0}
ks = sorted(set(b) & set(t)); n = len(ks)
cb = sum(b[k]["ok"] for k in ks); ct = sum(t[k]["ok"] for k in ks)
w = sum(t[k]["ok"] and not b[k]["ok"] for k in ks); l = sum(b[k]["ok"] and not t[k]["ok"] for k in ks)
m = w + l; p = min(1.0, 2 * sum(comb(m, i) for i in range(min(w, l) + 1)) / 2 ** m) if m else 1.0
def tab(key):
    d = defaultdict(lambda: [0, 0, 0])
    for k in ks: d[b[k][key]][0] += b[k]["ok"]; d[b[k][key]][1] += t[k]["ok"]; d[b[k][key]][2] += 1
    return {str(x): d[x] for x in sorted(d)}
tok = lambda R: sum(R[k]["predicted_n"] or 0 for k in ks)
cap = lambda R: sum((R[k]["predicted_n"] or 0) >= 2048 for k in ks)
res = {"n": n, "base_correct": cb, "tuned_correct": ct, "base_acc": round(100 * cb / n, 1), "tuned_acc": round(100 * ct / n, 1),
       "delta_pts": round(100 * (ct - cb) / n, 1), "tuned_wins": w, "tuned_losses": l, "mcnemar_exact_p": round(p, 4),
       "base_tokens": tok(b), "tuned_tokens": tok(t), "base_capped": cap(b), "tuned_capped": cap(t),
       "by_subject_base_tuned_n": tab("subject"), "by_level_base_tuned_n": tab("level"),
       "pass": (ct - cb) / n * 100 >= 3.0 and p < 0.1}
json.dump(res, open("logs/e20_result.json", "w"), indent=1); print(json.dumps(res, indent=1))