#!/usr/bin/env python3
"""
compare_nll_vs_generative.py -- put the NLL proxy and the generative benchmark
side by side and judge three PRE-REGISTERED predictions against numeric
thresholds fixed before the numbers were seen.

  usage: compare_nll_vs_generative.py <longeval_out_dir> <nll_ref_dir> [> EVAL_COMPARISON.md]

Why thresholds up front: at 100 prompts per (model, depth) cell the 95% Wilson
half-width is ~+-10 points, so "arm C beats arm B" is only meaningful relative to
that resolution. Judging after the fact would be post-hoc.

The two metrics are NOT on a common scale -- one is nats of NLL margin, the other
a hit rate -- so only ORDERINGS are compared, never absolute values.
"""

import csv
import math
import os
import sys

# NLL probe dir -> generative checkpoint key
NLL_MODEL = {
    "probe_run7_yarn4": "base",
    "probe_arm_a": "arm_a",
    "probe_arm_b": "arm_b",
    "probe_arm_c": "arm_c",
}
PFRAC_DEPTH = {"p0": 0.0, "p0.25": 0.25, "p0.5": 0.5, "p0.75": 0.75, "p1": 1.0}


def wilson(h, n):
    if not n:
        return float("nan")
    z, p = 1.96, h / n
    return 100.0 * z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)


def load_nll(ref):
    """{model: {depth: margin_nats}}  margin = nll_null - nll (higher = better)."""
    out = {}
    for d, key in NLL_MODEL.items():
        p = os.path.join(ref, "logs", d, "reach.csv")
        if not os.path.exists(p):
            continue
        out[key] = {}
        for r in csv.DictReader(open(p)):
            dep = PFRAC_DEPTH.get(r["p_frac"])
            if dep is None:
                continue
            out[key][dep] = float(r["nll_null"]) - float(r["nll"])
    return out


def load_gen(out_dir):
    """{model: {depth: (hits, n)}} from depth_summary.csv."""
    p = os.path.join(out_dir, "depth_summary.csv")
    if not os.path.exists(p):
        sys.exit(f"missing {p} -- run run_longeval.sh first")
    g = {}
    for r in csv.DictReader(open(p)):
        g.setdefault(r["model"], {})[float(r["depth"])] = (int(r["hits"]), int(r["n"]))
    return g


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    out_dir, ref = sys.argv[1], sys.argv[2]
    nll, gen = load_nll(ref), load_gen(out_dir)
    models = [m for m in ("base", "arm_a", "arm_b", "arm_c") if m in gen]
    depths = sorted({d for m in gen.values() for d in m})
    pct = lambda hn: 100.0 * hn[0] / hn[1] if hn[1] else float("nan")
    G = lambda m, d: pct(gen[m][d]) if m in gen and d in gen[m] else float("nan")

    print("# NLL proxy vs generative benchmark\n")
    print("Two different measurements: NLL margin in nats (teacher-forced, answer")
    print("already in the input) vs hit rate on generated text. Only the ORDERINGS")
    print("are comparable.\n")

    print("## Generative accuracy %, pooled by depth across tasks (+-95% Wilson)\n")
    print("| depth | " + " | ".join(models) + " |")
    print("|---" * (len(models) + 1) + "|")
    for d in depths:
        cells = []
        for m in models:
            if m in gen and d in gen[m]:
                h, n = gen[m][d]
                cells.append(f"{pct((h,n)):.1f} ±{wilson(h,n):.1f} (n={n})")
            else:
                cells.append("-")
        print(f"| {d:.2f} | " + " | ".join(cells) + " |")

    print("\n## NLL margin (nats), same models and depths\n")
    print("| depth | " + " | ".join(models) + " |")
    print("|---" * (len(models) + 1) + "|")
    for d in depths:
        print(f"| {d:.2f} | " + " | ".join(
            (f"{nll[m][d]:.2f}" if m in nll and d in nll[m] else "-") for m in models) + " |")

    # ---- pre-registered predictions -------------------------------------------
    print("\n## Pre-registered predictions\n")
    verdicts = []

    # 1. depth 0.0 ~ 0% everywhere; conditional on the positive control
    d0 = {m: G(m, 0.0) for m in models}
    worst = max((v for v in d0.values() if v == v), default=float("nan"))
    ctrl = {m: G(m, 1.0) for m in models}
    ctrl_ok = any(v > 0 for v in ctrl.values() if v == v)
    v1 = ("CONFIRMED" if worst <= 5 else "REFUTED" if worst >= 15 else "INCONCLUSIVE")
    if not ctrl_ok:
        v1 = "UNINTERPRETABLE (positive control failed)"
    verdicts.append(("1. depth 0.0 ~= 0% for every model", v1,
                     f"max over models = {worst:.1f}%; positive control (any model >0% at depth 1.0) = "
                     f"{'PASS' if ctrl_ok else 'FAIL'}"))

    # 2. arm C >> arm B at depth 0.25
    if "arm_c" in gen and "arm_b" in gen:
        diff2 = G("arm_c", 0.25) - G("arm_b", 0.25)
        v2 = ("CONFIRMED" if diff2 >= 20 else "REFUTED" if abs(diff2) <= 5 else "INCONCLUSIVE")
        verdicts.append(("2. arm C >> arm B at depth 0.25", v2,
                         f"armC {G('arm_c',0.25):.1f}% - armB {G('arm_b',0.25):.1f}% = {diff2:+.1f}pp "
                         f"(NLL said 2.80 vs 0.53 nats, ~9.5 sigma)"))

    # 3. arm C non-monotonicity: depth 0.25 > depth 0.5
    if "arm_c" in gen:
        diff3 = G("arm_c", 0.25) - G("arm_c", 0.5)
        v3 = ("CONFIRMED" if diff3 >= 20 else "REFUTED" if abs(diff3) <= 5 else "INCONCLUSIVE")
        verdicts.append(("3. arm C depth 0.25 > depth 0.5", v3,
                         f"{G('arm_c',0.25):.1f}% - {G('arm_c',0.5):.1f}% = {diff3:+.1f}pp "
                         f"(NLL said 1.0865 nats at 6.87 sigma)"))

    print("| prediction | verdict | evidence |")
    print("|---|---|---|")
    for name, v, ev in verdicts:
        print(f"| {name} | **{v}** | {ev} |")

    # ---- ordering agreement ----------------------------------------------------
    print("\n## Ordering agreement (the only comparable thing)\n")
    for m in models:
        if m not in nll:
            continue
        gd = [(d, G(m, d)) for d in depths if G(m, d) == G(m, d)]
        nd = [(d, nll[m][d]) for d in depths if d in nll[m]]
        go = [d for d, _ in sorted(gd, key=lambda x: -x[1])]
        no = [d for d, _ in sorted(nd, key=lambda x: -x[1])]
        same = "AGREE" if go == no else "DIFFER"
        print(f"- **{m}**: generative best->worst {go}; NLL best->worst {no} -> {same}")
        if all(v == 0.0 for _, v in gd):
            print(f"  (all-zero generative: ordering is undefined, not disagreeing)")


if __name__ == "__main__":
    main()