#!/usr/bin/env bash
# kld-gate-qwen38.sh — RUNS ON DANKSERVER (or anywhere the reports are). The campaign's
# acceptance gate: one verdict from the KLD reports of the containers that ride one reference
# recording.
#
# The reference is the BF16 container built from the vendor checkpoint (run-qwen38-bf16.sh
# kldref, in the same drained window that records the calibration grams). Every scored
# container is measured against that one recording, so the KL scale runs from 0 = the
# checkpoint itself downward in faithfulness.
#
# TWO MODES, because the disk decides how many anchors can exist at once:
#
#   gate IQ4R.json CAND.json                    candidate vs the served payload
#   gate FP8.json IQ4R.json CAND.json           the same, inside the absolute interval
#
# iq4r.json is the payload frank-archon serves — the bar to beat. When the FP8-native release
# (172.8 GiB, fetched last because of the peak-disk law) is on disk, pass its report first:
# "half the gap" is then half of the distance between the 4-bit payload and the near-lossless
# release, not half of an invented distance to the checkpoint. Without it the gate still
# answers the campaign's question — does the candidate beat what is served — and says so.
#
# The law, on the summary of one slice (--slice all | a corpus source):
#
#   candidate vs served (2 reports)
#   1. mean KL(cand) < mean KL(iq4r)         the served payload is beaten on the corpus
#   2. p99.9(cand) <= p99.9(iq4r)            and the tail is no worse — a mean can hide a few
#                                            positions taking the whole hit
#   3. top-1(cand) >= top-1(iq4r)            the agreement the harness keys on does not fall
#   4. every source slice passes 1-3         a win bought by regressing one distribution
#                                            (code, prose, tools) is not a win
#
#   inside the interval (3 reports), laws 1-4 above (1-3 renumbered to the interval and the
#   recovery fraction, as agreed 2026-10-27) plus:
#   5. p99.9(cand) <= p99.9(iq4r)            the tail law
#   6. every source slice passes            the slice law
#
# Booting the containers is NOT this script's job: `run-qwen38-bf16.sh score RAD OUT.json`
# writes the reports and owns the single engine argv (same pin, --tp 2, --kv-cache-dtype fp8,
# same step size, --deterministic, no speculation) that makes the numbers comparable.
# This script is arithmetic, so it re-runs against old reports and is testable without a GPU.
#
# Usage: kld-gate-qwen38.sh IQ4R.json CAND.json [--slice NAME]
#        kld-gate-qwen38.sh FP8.json IQ4R.json CAND.json [--slice NAME] [--recover F]
set -uo pipefail

if [ "$#" -lt 2 ]; then
    echo "usage: kld-gate-qwen38.sh IQ4R.json CAND.json [--slice NAME]" >&2
    echo "       kld-gate-qwen38.sh FP8.json IQ4R.json CAND.json [--slice NAME] [--recover F]" >&2
    exit 2
fi
python3 - "$@" <<'PY'
import json, sys

a = sys.argv[1:]
opts, files = [], []
i = 0
while i < len(a):
    if a[i].startswith("--"):
        opts += [a[i], a[i + 1]]
        i += 2
    else:
        files.append(a[i])
        i += 1
if len(files) not in (2, 3):
    sys.exit("usage: kld-gate-qwen38.sh IQ4R.json CAND.json [--slice NAME]\n"
             "       kld-gate-qwen38.sh FP8.json IQ4R.json CAND.json [--slice NAME] [--recover F]")
o = dict(zip(opts[0::2], opts[1::2]))
slice_ = o.get("--slice", "all")
recover = float(o.get("--recover", 0.5))
if not 0.0 < recover <= 1.0:
    sys.exit(f"--recover must be in (0,1], got {recover}")
interval = len(files) == 3


def die(msg):
    print(msg, file=sys.stderr)
    raise SystemExit(2)


def load(path):
    try:
        return json.load(open(path))
    except Exception as e:
        die(f"cannot read the KLD report {path}: {e}")


def summary(path, d):
    """The slice this gate grades: 'all', or one corpus source."""
    if slice_ == "all":
        g = d.get("all")
    else:
        g = (d.get("by_source") or {}).get(slice_)
        if g is None:
            die(f"{path} has no slice {slice_!r}; it has: "
                f"{', '.join(sorted(d.get('by_source') or {})) or '(none)'}")
    if g is None or "kld" not in g or "top1_agreement" not in g:
        die(f"{path} has no {slice_!r} summary — wrong report version?")
    return g


docs = {label: load(path) for label, path in
        zip(("fp8", "iq4r", "cand") if interval else ("iq4r", "cand"), files)}
rows = []
for label, path in zip(("fp8", "iq4r", "cand") if interval else ("iq4r", "cand"), files):
    g = summary(path, docs[label])
    rows.append((label, path, docs[label].get("model", "?"), g["kld"]["mean"],
                 g["kld"]["p99_9"], g["top1_agreement"], g["positions"]))

print(f"reference: BF16 container; slice: {slice_}; "
      + ("inside the FP8 interval" if interval else "candidate vs served (no FP8 anchor)"))
print(f"{'':6} {'model':38} {'KL mean':>10} {'KL p99.9':>10} {'top-1':>8}  positions")
for label, path, model, kl, p999, top1, n in rows:
    print(f"{label:6} {str(model)[:38]:38} {kl:10.6f} {p999:10.4f} {100*top1:7.2f}%  {n}")
r = {label: (kl, p999, top1) for label, _p, _m, kl, p999, top1, _n in rows}
kl_cand, p_cand, t_cand = r["cand"]
kl_iq4r, p_iq4r, t_iq4r = r["iq4r"]
ok = True


def law(n, text, passed):
    global ok
    ok = ok and passed
    print(f"  [{'PASS' if passed else 'FAIL'}] {n}. {text}")


def tail_law(n):
    law(n, f"candidate KL p99.9 {p_cand:.4f} <= served iq4r {p_iq4r:.4f}",
        p_cand <= p_iq4r)


# --- the slice law: every corpus source must pass the metric laws -------------
def slice_law(n):
    """A win that regresses one source class is not a win: check them all."""
    mine = docs["cand"].get("by_source") or {}
    theirs = docs["iq4r"].get("by_source") or {}
    if not mine or not theirs:
        law(n, "no per-source summaries to check", True)
        return
    missing = sorted(set(theirs) - set(mine))
    extra = sorted(set(mine) - set(theirs))
    if missing or extra:
        law(n, "same source slices in both reports"
            + (f" (missing from cand: {', '.join(missing)})" if missing else "")
            + (f" (only in cand: {', '.join(extra)})" if extra else ""), False)
        return
    bad = []
    for src in sorted(theirs):
        c, q = mine[src], theirs[src]
        worse_mean = c["kld"]["mean"] >= q["kld"]["mean"]
        worse_tail = c["kld"]["p99_9"] > q["kld"]["p99_9"]
        worse_top1 = c["top1_agreement"] < q["top1_agreement"]
        if worse_mean or worse_tail or worse_top1:
            why = ",".join(m for m, w in (("mean", worse_mean), ("p99.9", worse_tail),
                                          ("top-1", worse_top1)) if w)
            bad.append(f"{src} ({why})")
    law(n, "every source slice passes (mean, p99.9, top-1)"
        + (f" — regressed: {'; '.join(bad)}" if bad else ""), not bad)


if interval:
    kl_fp8 = r["fp8"][0]
    gap = kl_iq4r - kl_fp8
    law(1, f"anchors ordered: fp8 {kl_fp8:.6f} < iq4r {kl_iq4r:.6f} (gap {gap:.6f})",
        gap > 0)
    law(2, f"candidate inside the interval: {kl_fp8:.6f} < {kl_cand:.6f} < {kl_iq4r:.6f}",
        kl_fp8 < kl_cand < kl_iq4r)
    need = kl_iq4r - recover * gap
    law(3, f"candidate closes {100*max((kl_iq4r - kl_cand)/gap, 0):.1f}% of the "
           f"{gap:.6f} gap (needs {100*recover:.0f}%, i.e. KL <= {need:.6f})"
           if gap > 0 else "gap is not positive; cannot score recovery",
        gap > 0 and kl_cand <= need)
    law(4, f"candidate top-1 {100*t_cand:.2f}% >= served iq4r top-1 {100*t_iq4r:.2f}%",
        t_cand >= t_iq4r)
    tail_law(5)
    slice_law(6)
else:
    law(1, f"candidate beats the served payload: KL {kl_cand:.6f} < iq4r {kl_iq4r:.6f}",
        kl_cand < kl_iq4r)
    tail_law(2)
    law(3, f"candidate top-1 {100*t_cand:.2f}% >= served iq4r top-1 {100*t_iq4r:.2f}%",
        t_cand >= t_iq4r)
    slice_law(4)

print("KLD GATE " + ("PASS" if ok else "FAIL"))
sys.exit(0 if ok else 1)
PY