"""Final statistics pass over all eval per_item logs (CPU only).

Produces results-loop/STATS.md + stats_final.json:
  1. Wilson 95% CIs for every headline number (overall + hard bucket).
  2. Exact McNemar tests for the key paired comparisons (same items).
  3. Pooled hard bucket across MBPP + HumanEval + Rust (paired k>0 vs k0
     within each item, pooled counts).
  4. Label-robustness check: hard bucket redefined via consensus k0
     failure across seeds instead of the single greedy labeling run.
"""

import json
import math
from pathlib import Path

OUT = Path(__file__).resolve().parent.parent / "results-loop"


def wilson(c, n, z=1.96):
    if n == 0:
        return (float("nan"), float("nan"))
    p = c / n
    d = 1 + z * z / n
    ctr = (p + z * z / (2 * n)) / d
    hw = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return (ctr - hw, ctr + hw)


def binom_two_sided(k, n):
    """Exact two-sided binomial test p-value, p0=0.5 (for McNemar)."""
    if n == 0:
        return 1.0
    def pmf(i):
        return math.comb(n, i) * 0.5 ** n
    pk = pmf(k)
    return min(1.0, sum(pmf(i) for i in range(n + 1) if pmf(i) <= pk + 1e-12))


def mcnemar(pairs):
    """pairs: list of (a_ok, b_ok). Returns dict with discordants + p."""
    b01 = sum(1 for a, b in pairs if not a and b)  # b wins
    b10 = sum(1 for a, b in pairs if a and not b)  # a wins
    return {"n": len(pairs), "a_only": b10, "b_only": b01,
            "p": binom_two_sided(min(b01, b10), b01 + b10)}


def load_labels(path, key, lab_key="label"):
    data = json.load(open(OUT / path))
    return {it[key]: it[lab_key] for it in data
            if it.get("split", "test") == "test"}


def per_item(fname, k):
    d = json.load(open(OUT / fname))
    v = d["ks"][str(k)]
    key = "task_id" if "task_id" in v["per_item"][0] else "idx"
    return {it[key]: it["ok"] for it in v["per_item"]}


def acc_ci(ok_map, subset=None):
    ids = [i for i in ok_map if subset is None or i in subset]
    c = sum(ok_map[i] for i in ids)
    lo, hi = wilson(c, len(ids))
    return {"acc": c / len(ids) if ids else float("nan"), "n": len(ids),
            "ci": [round(lo, 3), round(hi, 3)]}


def fmt(r):
    return f"{r['acc']:.3f} [{r['ci'][0]:.3f}, {r['ci'][1]:.3f}] (n={r['n']})"


def main():
    mbpp_lab = load_labels("mbpp_data.json", "task_id")
    he_lab = {k: ("hard" if v else "easy")   # plan-reachable flag; hard needs k0-fail
              for k, v in json.load(open(OUT / "humaneval_labels.json")).items()}
    rust_lab = load_labels("rust_data.json", "task_id")

    mbpp_hard = {t for t, l in mbpp_lab.items() if l == "hard"}
    rust_hard = {t for t, l in rust_lab.items() if l == "hard"}

    report = {}
    lines = ["# Final statistics pass", ""]

    # ---------- 1. headline numbers with Wilson CIs ----------
    lines += ["## Headline numbers (Wilson 95% CIs)", ""]
    ARMS = [
        # (label, file, k, hard-subset)
        ("MBPP loop s0 k=0 (base)", "eval_code_code_s0_full.json", 0, mbpp_hard),
        ("MBPP loop s0 k=2", "eval_code_code_s0_full.json", 2, mbpp_hard),
        ("MBPP loop s0 k=4", "eval_code_code_s0_full.json", 4, mbpp_hard),
        ("MBPP distill s1 k=1 (FF)", "eval_code_distill_s1.json", 1, mbpp_hard),
        ("MBPP pause16 k=1", "eval_code_pause16.json", 1, mbpp_hard),
        ("MBPP stack-train k=4", "eval_code_stack_train.json", 4, mbpp_hard),
        ("MBPP distill-in-loopmode k=2", "eval_code_distill_loopmode.json", 2, mbpp_hard),
        ("Rust transfer k=0", "eval_rust_py_transfer.json", 0, rust_hard),
        ("Rust transfer k=4", "eval_rust_py_transfer.json", 4, rust_hard),
    ]
    arm_maps = {}
    for label, f, k, hard in ARMS:
        m = per_item(f, k)
        arm_maps[label] = (m, hard)
        o, h = acc_ci(m), acc_ci(m, hard)
        report[label] = {"overall": o, "hard": h}
        lines.append(f"- **{label}**: overall {fmt(o)}; hard {fmt(h)}")

    # HumanEval: hard = plan-reachable AND k0-fail (per its eval definition)
    he_tr0 = per_item("eval_humaneval_trained.json", 0)
    he_hard = {t for t in he_tr0 if he_lab.get(t) == "hard" and not he_tr0[t]}
    for label, f, k in [("HumanEval loop k=4", "eval_humaneval_trained.json", 4),
                        ("HumanEval distill k=1", "eval_humaneval_distill_transfer.json", 1)]:
        m = per_item(f, k)
        o, h = acc_ci(m), acc_ci(m, he_hard)
        arm_maps[label] = (m, he_hard)
        report[label] = {"overall": o, "hard": h}
        lines.append(f"- **{label}**: overall {fmt(o)}; hard {fmt(h)}")

    # best-of-3 / budget-CoT (no per_item; CI from counts)
    for label, f in [("MBPP best-of-3 (compute-matched)", "eval_bestof3.json"),
                     ("MBPP budget-CoT-50", "eval_budgetcot.json")]:
        d = json.load(open(OUT / f))
        n, nh = 500, len(mbpp_hard)
        o = {"acc": d["acc"], "n": n,
             "ci": [round(x, 3) for x in wilson(round(d["acc"] * n), n)]}
        hacc = d["by_label"]["hard"]
        h = {"acc": hacc, "n": nh,
             "ci": [round(x, 3) for x in wilson(round(hacc * nh), nh)]}
        report[label] = {"overall": o, "hard": h}
        lines.append(f"- **{label}**: overall {fmt(o)}; hard {fmt(h)}")

    # distill seed spread
    hs, os_ = [], []
    files = ["eval_code_code_distill.json"] + [
        f"eval_code_distill_s{s}.json" for s in range(1, 8)]
    for f in files:
        try:
            m = per_item(f, 1)
        except FileNotFoundError:
            continue
        hs.append(acc_ci(m, mbpp_hard)["acc"])
        os_.append(acc_ci(m)["acc"])
    mean = sum(hs) / len(hs)
    sd = (sum((x - mean) ** 2 for x in hs) / (len(hs) - 1)) ** 0.5
    lines.append(f"- **MBPP distill, {len(hs)} runs (hard)**: mean {mean:.3f} "
                 f"± {sd:.3f} sd  (range {min(hs):.3f}-{max(hs):.3f}); "
                 f"overall mean {sum(os_)/len(os_):.3f}")
    report["distill_seed_spread"] = {"hard": hs, "overall": os_}

    # loop seed spread (s0..s4 k=4)
    lh = []
    for tag in ["s0_full", "s1", "s2", "s3", "s4"]:
        try:
            m = per_item(f"eval_code_code_{tag}.json", 4)
            lh.append(acc_ci(m, mbpp_hard)["acc"])
        except (FileNotFoundError, KeyError):
            pass
    if lh:
        mean = sum(lh) / len(lh)
        sd = (sum((x - mean) ** 2 for x in lh) / max(1, len(lh) - 1)) ** 0.5
        lines.append(f"- **MBPP loop seeds k=4 (hard)**: mean {mean:.3f} "
                     f"± {sd:.3f} sd (n_seeds={len(lh)})")
        report["loop_seed_spread_hard"] = lh

    # ---------- 2. McNemar paired tests ----------
    lines += ["", "## McNemar exact tests (paired on items)", ""]

    def pair(m_a, m_b, subset=None):
        ids = [i for i in m_a if i in m_b
               and (subset is None or i in subset)]
        return [(m_a[i], m_b[i]) for i in ids]

    loop0, _ = arm_maps["MBPP loop s0 k=0 (base)"]
    loop4, _ = arm_maps["MBPP loop s0 k=4"]
    dist1, _ = arm_maps["MBPP distill s1 k=1 (FF)"]
    stack4, _ = arm_maps["MBPP stack-train k=4"]

    TESTS = [
        ("loop k=4 vs k=0, overall", loop0, loop4, None),
        ("loop k=4 vs k=0, hard", loop0, loop4, mbpp_hard),
        ("distill k=1 vs loop k=4, overall", loop4, dist1, None),
        ("distill k=1 vs loop k=4, hard", loop4, dist1, mbpp_hard),
        ("stack-train k=4 vs distill k=1, hard", dist1, stack4, mbpp_hard),
        ("HumanEval loop k=4 vs k=0, overall", he_tr0,
         arm_maps["HumanEval loop k=4"][0], None),
        ("HumanEval distill k=1 vs k=0, overall", he_tr0,
         arm_maps["HumanEval distill k=1"][0], None),
        ("HumanEval trained vs UNTRAINED merge (k=2), overall",
         per_item("eval_humaneval_untrained.json", 2),
         per_item("eval_humaneval_trained.json", 2), None),
    ]
    rust0 = arm_maps["Rust transfer k=0"][0]
    rust4 = arm_maps["Rust transfer k=4"][0]
    TESTS += [("Rust loop k=4 vs k=0, overall", rust0, rust4, None),
              ("Rust loop k=4 vs k=0, hard", rust0, rust4, rust_hard)]

    report["mcnemar"] = {}
    for name, a, b, subset in TESTS:
        r = mcnemar(pair(a, b, subset))
        report["mcnemar"][name] = r
        sig = "**significant**" if r["p"] < 0.05 else "n.s."
        lines.append(f"- {name}: A-only {r['a_only']}, B-only {r['b_only']}, "
                     f"n={r['n']}, p={r['p']:.4g} ({sig})")

    # ---------- 3. pooled hard bucket across benchmarks ----------
    lines += ["", "## Pooled hard bucket (MBPP + HumanEval + Rust)",
              "", "Paired within-item k>0 vs k=0, counts pooled across "
              "benchmarks (loop arm; distill pooled where available).", ""]
    pooled_loop = (pair(loop0, loop4, mbpp_hard)
                   + pair(he_tr0, arm_maps["HumanEval loop k=4"][0], he_hard)
                   + pair(rust0, rust4, rust_hard))
    r = mcnemar(pooled_loop)
    c_base = sum(a for a, _ in pooled_loop)
    c_loop = sum(b for _, b in pooled_loop)
    n = len(pooled_loop)
    lines.append(f"- **loop**: base {c_base}/{n} -> loop {c_loop}/{n} "
                 f"({c_base/n:.3f} -> {c_loop/n:.3f}, CI "
                 f"{[round(x,3) for x in wilson(c_loop, n)]}), "
                 f"McNemar p={r['p']:.3g}")
    report["pooled_hard_loop"] = {"base": c_base, "arm": c_loop, "n": n,
                                  "mcnemar": r}
    pooled_dist = (pair(loop0, dist1, mbpp_hard)
                   + pair(he_tr0, arm_maps["HumanEval distill k=1"][0], he_hard))
    r = mcnemar(pooled_dist)
    c_base = sum(a for a, _ in pooled_dist)
    c_d = sum(b for _, b in pooled_dist)
    n = len(pooled_dist)
    lines.append(f"- **distill**: base {c_base}/{n} -> distill {c_d}/{n} "
                 f"({c_base/n:.3f} -> {c_d/n:.3f}, CI "
                 f"{[round(x,3) for x in wilson(c_d, n)]}), "
                 f"McNemar p={r['p']:.3g}")
    report["pooled_hard_distill"] = {"base": c_base, "arm": c_d, "n": n,
                                     "mcnemar": r}

    # ---------- 4. label robustness: consensus-k0 hard set ----------
    lines += ["", "## Label robustness (consensus-k0 hard set)", "",
              "Hard bucket redefined as: labeled hard AND k=0 fails in "
              "every seed's own eval run (removes single-greedy-run "
              "selection noise).", ""]
    k0maps = []
    for tag in ["s0_full", "s1", "s2", "s3", "s4"]:
        try:
            k0maps.append(per_item(f"eval_code_code_{tag}.json", 0))
        except (FileNotFoundError, KeyError):
            pass
    consensus = {t for t in mbpp_hard
                 if all(not m.get(t, False) for m in k0maps)}
    lines.append(f"- consensus hard set: {len(consensus)} of "
                 f"{len(mbpp_hard)} labeled-hard items")
    for label in ["MBPP loop s0 k=4", "MBPP distill s1 k=1 (FF)",
                  "MBPP stack-train k=4"]:
        m, _ = arm_maps[label]
        r0 = acc_ci(m, mbpp_hard)
        rc = acc_ci(m, consensus)
        lines.append(f"- {label}: labeled-hard {r0['acc']:.3f} -> "
                     f"consensus-hard {fmt(rc)}")
        report.setdefault("consensus_hard", {})[label] = rc

    (OUT / "STATS.md").write_text("\n".join(lines) + "\n")
    json.dump(report, open(OUT / "stats_final.json", "w"), indent=1)
    print("\n".join(lines))
    print("\nwrote", OUT / "STATS.md", "and stats_final.json")


if __name__ == "__main__":
    main()