#!/usr/bin/env python3
"""D-59b ship rule for a release candidate against the file it would replace: per test group, the score difference
on the same items with its paired standard error (both models answered identical questions), a verdict per group
(drop / gain beyond 2 SE, or within noise), and the KL difference. The D-59b fixed thresholds are printed beside the
paired ones; the verdict uses the paired SE (the exact test of the same rule).

    python scripts/ship_rule.py --base marco-3bit-iq3s --cand marco-3bit-b3-20 [--kl-base 0.02854 --kl-cand 0.02773]

Rule: replace only if no group drops beyond 2 SE, KL is not worse by more than 0.0003, and there is a real
improvement: a group gaining beyond 2 SE, or KL better by more than 0.0003 (28 Oct 05:50 reading of the D-59b rule,
for the human to confirm: B3-20 is -0.0008 KL with every benchmark within noise).
"""
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
CMP = REPO / "reports" / "compare"
LANGS, FULL_LANGS = ["zh", "ar", "ko", "hu", "bn"], ["de", "es", "fr", "ja", "ru"]
GROUPS = {"MMLU-Redux": ["mmlu_redux"], "Global-MMLU-Lite (5)": [f"gmmlu_{l}" for l in LANGS],
          "Belebele (5)": [f"belebele_{l}" for l in LANGS], "GSM8K-250": ["gsm8k"], "MMLU-Pro": ["mmlu_pro"],
          "Global-MMLU (5 more)": [f"gmmlux_{l}" for l in FULL_LANGS], "MATH-500": ["math500"],
          "IFEval (prompt)": ["ifeval"], "HumanEval+": ["humanevalplus"], "MBPP+": ["mbppplus"]}
FIXED_2SE = {"MMLU-Redux": 3.1, "Global-MMLU-Lite (5)": 2.2, "Belebele (5)": 2.8, "GSM8K-250": 4.3, "MMLU-Pro": 2.1,
             "Global-MMLU (5 more)": 1.4, "MATH-500": 4.1, "IFEval (prompt)": 4.2, "HumanEval+": 7.5, "MBPP+": 4.8}
KL_TOL = 0.0003


def items(name: str) -> dict:
    out = {}
    for suffix in ("", ".code"):
        p = CMP / f"{name}{suffix}.jsonl"
        if p.exists():
            for line in p.open():
                r = json.loads(line)
                out[(r["task"], r["id"])] = r
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", required=True)
    ap.add_argument("--cand", required=True)
    ap.add_argument("--kl-base", type=float)
    ap.add_argument("--kl-cand", type=float)
    a = ap.parse_args()
    b, c = items(a.base), items(a.cand)
    rows, drops, gains, missing = [], [], [], []
    for g, tasks in GROUPS.items():
        keys = [k for k in b if k[0] in tasks and k in c]
        if not keys:
            missing.append(g)
            continue
        d = [int(c[k]["correct"]) - int(b[k]["correct"]) for k in keys]
        n = len(d)
        mean = sum(d) / n
        se = math.sqrt(sum((x - mean) ** 2 for x in d) / (n - 1) / n) if n > 1 else float("nan")
        delta, two_se = 100 * mean, 200 * se
        verdict = "drop" if delta < -two_se else "gain" if delta > two_se else "noise"
        (drops if verdict == "drop" else gains if verdict == "gain" else []).append(g)
        rows.append((g, n, 100 * sum(int(b[k]["correct"]) for k in keys) / n,
                     100 * sum(int(c[k]["correct"]) for k in keys) / n, delta, two_se, FIXED_2SE[g], verdict))
    print(f"{a.cand} vs {a.base}")
    print(f"{'test':22s} {'n':>5s} {'base':>6s} {'cand':>6s} {'delta':>6s} {'2SE':>5s} {'fixed':>5s}  verdict")
    for g, n, sb, sc, dl, ts, fx, v in rows:
        print(f"{g:22s} {n:5d} {sb:6.1f} {sc:6.1f} {dl:+6.1f} {ts:5.1f} {fx:5.1f}  {v}")
    kl_ok = kl_gain = None
    if a.kl_base is not None and a.kl_cand is not None:
        kl_ok = a.kl_cand <= a.kl_base + KL_TOL
        kl_gain = a.kl_cand < a.kl_base - KL_TOL
        print(f"KL {a.kl_base:.5f} -> {a.kl_cand:.5f} ({a.kl_cand - a.kl_base:+.5f}): {'ok' if kl_ok else 'WORSE'}")
    if missing:
        print("not measured yet:", ", ".join(missing))
    ok = not drops and (bool(gains) or bool(kl_gain)) and kl_ok is not False
    print(json.dumps({"base": a.base, "cand": a.cand, "drops": drops, "gains": gains, "kl_ok": kl_ok, "kl_gain": kl_gain,
                      "missing": missing, "replace": ok and not missing}))


if __name__ == "__main__":
    main()