"""
sanity_dense.py -- REF (dense) gate: does our pipeline reproduce the OFFICIAL
LLaDA-8B-Base EVAL.md numbers? (README section 0: compare to LLaDA official,
NOT the Sink-Aware Dense row, which is not self-consistent under one protocol.)

Reads the per-item converter summaries written by analysis/lmeval_to_items.py
(field acc_primary = the pre-registered primary metric per task) and diffs them
against official EVAL.md. If a task is off beyond sampling error, investigate the
pipeline BEFORE running any compression eval.

Official LLaDA-8B-Base (evaluation/EVAL.md, lm-eval):
  arc_challenge 47.9 (acc, w/ CFG) | hellaswag 72.5 (acc_norm, w/ CFG)
  winogrande 74.8 (acc, 5-shot)    | piqa 74.4 (acc_norm, w/ CFG)
  mmlu 65.9 (acc, w/o CFG, 5-shot) | gsm8k 70.0 (exact_match, 256/256/256)

Usage:
  python sanity_dense.py           # reads results/eval/ref/*_ref_*_summary.json
"""
import os
import sys
import glob
import math
import argparse

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import common as C

TARGET = {  # official EVAL.md, primary metric per task -- LLaDA-8B-BASE
    "arc_challenge": 47.9,
    "hellaswag": 72.5,
    "winogrande": 74.8,
    "piqa": 74.4,
    "mmlu": 65.9,
    "gsm8k": 70.0,
    # arc_easy: not tabulated in EVAL.md -> no target
}

# LLaDA-8B-INSTRUCT, evaluation/EVAL.md Instruct table, "Internal toolkit" row --
# that is the column the paper's Tab.1/Tab.2 report. (The OpenCompass row of the same
# table runs 0.2-2.4pp lower; we gate against the reported numbers, not the reproduction.)
# Only the four tasks our Instruct table shares with that table get a target.
# bbh_llada and ifeval are absent from it entirely -> no target, no gate.
TARGET_INSTRUCT = {
    "mmlu_generative": 65.5,
    "gsm8k": 69.4,
    "humaneval_instruct": 49.4,
    "mbpp_instruct": 41.0,
}

TARGETS = {"base": TARGET, "instruct": TARGET_INSTRUCT}


def wilson_hw(p, n, z=1.96):
    return z * math.sqrt(max(p * (1 - p), 1e-12) / n) if n else float("nan")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--glob", default=None)
    ap.add_argument("--tol_pp", type=float, default=3.0)
    ap.add_argument("--family", default="base", choices=list(TARGETS))
    ap.add_argument("--out_root", default=None,
                    help="eval results root (default results/eval, i.e. the Base table)")
    args = ap.parse_args()
    root = os.path.dirname(os.path.abspath(__file__))
    target = TARGETS[args.family]
    out_root = args.out_root or os.path.join(root, "results", "eval")
    pat = args.glob or os.path.join(out_root, "ref", "*_ref_*_summary.json")
    files = sorted(glob.glob(pat))
    if not files:
        sys.exit(f"no REF summaries at {pat}; run eval/run_lmeval.py + "
                 "analysis/lmeval_to_items.py (arm=ref) first")
    rows = {}
    for fp in files:
        s = C.load_json(fp)
        t = s.get("task")
        if not t or "acc_primary" not in s:
            continue
        if t not in rows or (s.get("n") or 0) > (rows[t].get("n") or 0):
            rows[t] = s

    print(f"{'task':<15} {'n':>6} {'metric':>10} {'ours':>7} {'EVAL.md':>8} "
          f"{'delta':>7} {'95%CI':>7}  flag")
    print("-" * 72)
    flags = []
    for t in sorted(rows):
        s = rows[t]
        n = s["n"]
        acc = 100.0 * (s.get("acc_primary") or 0)
        ci = 100.0 * wilson_hw(s.get("acc_primary") or 0, n)
        tgt = target.get(t)
        if tgt is None:
            print(f"{t:<15} {n:>6} {s['primary_metric']:>10} {acc:>7.2f} "
                  f"{'--':>8} {'--':>7} {ci:>7.2f}  (no target)")
            continue
        d = abs(acc - tgt)
        bad = d > (args.tol_pp + ci)
        flags.append((t, round(d, 2))) if bad else None
        print(f"{t:<15} {n:>6} {s['primary_metric']:>10} {acc:>7.2f} {tgt:>8.2f} "
              f"{d:>7.2f} {ci:>7.2f}  {'FLAG' if bad else 'ok'}")
    out = {"family": args.family, "target_evalmd": target,
           "ours": {t: {"n": rows[t]["n"], "acc_primary": rows[t].get("acc_primary"),
                        "metric": rows[t].get("primary_metric")} for t in rows},
           "flags": flags, "git_hash": C.git_hash()}
    C.dump_json(out, os.path.join(out_root,
                                  f"{C.git_hash()}_sanity_dense_{args.family}.json"))
    if flags:
        print(f"\n!! FLAGGED {flags} -- investigate pipeline before compression eval.")
    else:
        print("\nAll tasks within tolerance of official EVAL.md -> pipeline sane.")


if __name__ == "__main__":
    main()