#!/usr/bin/env python3
"""First-token KL divergence between a stock checkpoint and its abliterated twin.

WHAT THIS MEASURES, PRECISELY
-----------------------------
For each evaluation prompt we render the chat template, run both models, and read
the next-token distribution at the generation position — the distribution over the
**first token the model would emit**. KL(ref || cand) over that distribution is the
standard abliteration-damage metric: it is what Heretic minimizes as its objective,
and it is the number quoted for our gen-seat candidates (`absolute-heresy` 0.0759,
`JonathanColetti/Qwen3.8-27B-Uncensored` 0.1191).

Two classes of prompt are scored separately, and the split is the whole point:

  * **harmless** — held-out benign prompts. This is the *damage* number. Divergence
    here is collateral: the model moving on inputs the surgery had no business
    touching. Lower is better; this is the headline.
  * **harmful** — the reserved AdvBench test split. This is the *signal* number.
    Divergence here is the intended effect. Higher is better.

Their ratio is the interesting quantity — a surgical abliteration moves a lot on
refusal-triggering prompts and almost nothing elsewhere. A single averaged KL over
a mixed corpus hides exactly that, which is why this script never reports one.

WHAT THIS DOES *NOT* MEASURE
----------------------------
- **The MTP head.** `AutoModelForCausalLM` resolves to the text-only
  `Qwen3_5ForCausalLM`, which does not instantiate `mtp.*` (the same fact that
  makes `save_pretrained` wrong on the write path — see abliterate.py). So this is
  the main head's distribution shift only. MTP health is measured directly by
  acceptance rate, which is the gate that matters there (≳40%), not KL.
- **Quantization damage.** Both sides are bf16. Point this at the NVFP4 build and
  you would be measuring quant + abliteration together, which answers a different
  question.
- **Anything beyond the first token.** Divergence compounds over a rollout; a
  first-token number is a lower bound on trajectory divergence, not a summary of
  it. It is used here because it is the metric the reference figures use.

COMPARABILITY CAVEAT — read before quoting this against Heretic's numbers.
The reference figures come from Heretic's own optimizer on a *different base model*
with *its own* harmless prompt set and template. Same metric, different measurement
conditions. Treat a comparison as order-of-magnitude, not head-to-head.

BATCH SIZE IS 1, DELIBERATELY
-----------------------------
Two reasons, both learned upstream in this harness:
  1. **VRAM.** The window this runs in has the co-tenant seat still resident, so
     free memory after the 50 GB model lands is single-digit GB. A batched forward
     materializes [B, seq, vocab] logits — at B=8 that is ~0.6 GB in fp32, on top
     of activations, for no correctness benefit. With B=1 and `logits_to_keep=1`
     the lm_head runs on one position and the logits tensor is ~0.6 MB.
  2. **No padding semantics to get wrong.** `last_token_hidden` in abliterate.py
     needed a load-bearing argument about padding side and the DeltaNet linear-
     attention recurrence. At B=1 there is no padding, so that entire class of
     error is absent rather than reasoned about.
The cost is ~700 forwards per model instead of ~90 batches — roughly a minute.

GATES (this script refuses to produce a misleading number)
  exit 6 — tokenizer or vocab mismatch between ref and cand
  exit 7 — the two stages tokenized a prompt differently
  exit 8 — model sharded across GPUs or offloaded (residual stream corruption)
  exit 9 — PYTORCH_CUDA_ALLOC_CONF=expandable_segments (corrupts retained tensors)
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import statistics
import sys
import time
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))
from calibration import load_evaluation  # noqa: E402
# `render` is imported rather than copied deliberately: the KL must be measured on
# the exact rendering the direction was captured against. Two copies of eight lines
# would be the cheapest possible thing to let drift, and the drift would be silent
# — a different template changes what "the first token" even is.
from abliterate import render  # noqa: E402

# Rendering modes. A thinking model's first token means different things depending
# on which one you use, and refusal lives in the answer channel:
#   answer — enable_thinking=False, prompt ends "</think>\n\n"; the next token is
#            the first token of the ANSWER. This is where "I'm sorry" / "I can't"
#            actually appears, so this is the headline mode.
#   think  — enable_thinking=True, prompt ends "<think>\n"; the next token opens
#            the reasoning trace. Reported because the direction was captured
#            against both renderings and a divergence that shows up in only one of
#            them is a finding, not noise.
MODES = {"answer": False, "think": True}


def file_md5(path: Path) -> str | None:
    if not path.exists():
        return None
    h = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def free_vram_mib(device_index: int = 0) -> float:
    free, _total = torch.cuda.mem_get_info(device_index)
    return free / 2**20


def weights_mib(model_dir: Path) -> float:
    """Text-only weight bytes, straight from the safetensors headers, in MiB.

    Sized from the checkpoint rather than from a remembered number because the
    remembered number was wrong: the runbook carried "bf16 is 50 GB", which was
    50.10 **GiB** mislabelled, and the 3.7 GB gap is exactly the difference
    between "one seat must stop" and "both seats must stop". Vision and MTP are
    subtracted because `AutoModelForCausalLM` resolves to the text-only
    `Qwen3_5ForCausalLM` and never instantiates them.

    Reads only the 8-byte length prefix and the JSON header of each shard — no
    mmap, no tensor data. (`safe_open` would mmap the whole shard, which ENOMEMs
    on ZFS — see the quant playbook.)
    """
    import struct
    from glob import glob

    total = 0
    for shard in sorted(glob(str(model_dir / "*.safetensors"))):
        with open(shard, "rb") as f:
            n = struct.unpack("<Q", f.read(8))[0]
            header = json.loads(f.read(n))
        for key, meta in header.items():
            if key == "__metadata__":
                continue
            if ".visual." in key or key.startswith("visual.") or key.startswith("mtp."):
                continue
            lo, hi = meta["data_offsets"]
            total += hi - lo
    return total / 2**20


# --- the forward ---------------------------------------------------------------
@torch.no_grad()
def first_token_logprobs(model, tokenizer, text, device):
    """log-softmax over the vocabulary at the generation position. float32 on CPU.

    Returns (logprobs[V], input_ids tuple). The ids come back so the two stages can
    be proven to have tokenized the same thing — a KL between distributions
    conditioned on different prefixes is a number with no meaning, and nothing else
    in the pipeline would catch it.
    """
    enc = tokenizer(text, return_tensors="pt")
    ids = tuple(enc["input_ids"][0].tolist())
    enc = {k: v.to(device) for k, v in enc.items()}
    try:
        out = model(**enc, logits_to_keep=1)
    except TypeError:
        # older kwarg name, or a model whose forward does not accept it; the full
        # logits path is correct too, just fatter.
        out = model(**enc)
    logits = out.logits[0, -1, :].float()
    if not torch.isfinite(logits).all():
        raise RuntimeError("non-finite logits — refusing to score them")
    return torch.log_softmax(logits, dim=-1).cpu(), ids


def collect(model, tokenizer, prompts, device, label):
    """{mode: {"logprobs": [N, V] float32 cpu, "ids": [tuple, ...]}}"""
    result = {}
    for mode, thinking in MODES.items():
        rows, idlist = [], []
        t0 = time.time()
        for n, prompt in enumerate(prompts, 1):
            lp, ids = first_token_logprobs(
                model, tokenizer, render(tokenizer, prompt, thinking), device)
            rows.append(lp)
            idlist.append(ids)
            if n % 50 == 0 or n == len(prompts):
                print(f"    [{label}/{mode}] {n}/{len(prompts)} "
                      f"({n / max(time.time() - t0, 1e-6):.1f}/s)", flush=True)
        result[mode] = {"logprobs": torch.stack(rows), "ids": idlist}
    return result


# --- scoring -------------------------------------------------------------------
def divergences(ref_lp: torch.Tensor, cand_lp: torch.Tensor):
    """Per-prompt (kl_fwd, kl_rev, tv, top1_match) from two [N, V] log-prob blocks.

    float64 throughout. The summands are differences of logs on a 150k-entry
    simplex; in float32 the tail terms lose the precision that the head terms
    dominate, and the total drifts by a few percent for free. This costs ~1 GB of
    transient host RAM on a box with hundreds free.
    """
    p_log = ref_lp.double()
    q_log = cand_lp.double()
    p = p_log.exp()
    q = q_log.exp()
    kl_fwd = (p * (p_log - q_log)).sum(-1)          # KL(ref || cand) — the metric
    kl_rev = (q * (q_log - p_log)).sum(-1)          # reported so "which direction?"
    tv = 0.5 * (p - q).abs().sum(-1)                # Pinsker companion: TV <= sqrt(KL/2)
    top1 = (p_log.argmax(-1) == q_log.argmax(-1))
    return kl_fwd, kl_rev, tv, top1


def summarize(kl_fwd, kl_rev, tv, top1):
    kl = sorted(kl_fwd.tolist())
    n = len(kl)

    def pct(frac):
        return kl[min(n - 1, int(round(frac * (n - 1))))]

    return {
        "n": n,
        "kl_mean": float(kl_fwd.mean()),
        "kl_median": float(statistics.median(kl)),
        "kl_p90": pct(0.90),
        "kl_p95": pct(0.95),
        "kl_max": kl[-1],
        "kl_reverse_mean": float(kl_rev.mean()),
        "tv_mean": float(tv.mean()),
        "tv_median": float(statistics.median(tv.tolist())),
        "top1_agreement": float(top1.double().mean()),
        "per_prompt_kl": [round(v, 6) for v in kl_fwd.tolist()],
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ref", required=True, help="stock bf16 checkpoint dir")
    ap.add_argument("--cand", required=True, help="abliterated bf16 checkpoint dir")
    ap.add_argument("--out", required=True, help="results JSON")
    ap.add_argument("--n-harmless", type=int, default=256,
                    help="held-out benign prompts — the damage measurement")
    ap.add_argument("--n-harmful", type=int, default=104,
                    help="reserved harmful_behaviors[test] split — the signal measurement")
    ap.add_argument("--seed", type=int, default=1,
                    help="eval harmless sample seed; MUST differ from the capture's")
    ap.add_argument("--calib-harmless-n", type=int, default=416,
                    help="what the capture drew, so this run can exclude it")
    ap.add_argument("--calib-harmless-seed", type=int, default=0,
                    help="the capture's harmless seed, for the same reason")
    ap.add_argument("--noise-floor-n", type=int, default=32,
                    help="prompts re-run through ref to measure the numerical floor; 0 disables")
    ap.add_argument("--skip-gates", action="store_true",
                    help="escape hatch for a deliberate off-recipe comparison; prints loudly")
    ap.add_argument("--dry-run", action="store_true",
                    help="gates + corpus + rendering only; no model load, no GPU. Safe with "
                         "the seats up, and the thing to run before asking for a VRAM window.")
    ap.add_argument("--stage", choices=("all", "ref", "cand", "score"), default="all",
                    help="'all' (default) re-execs itself once per model then scores. The "
                         "single-model stages exist so each model gets its OWN PROCESS — see "
                         "the note on why in-process release does not work here.")
    ap.add_argument("--cache", default=None,
                    help="directory for the per-stage log-prob caches (default: beside --out)")
    args = ap.parse_args()

    ref_dir, cand_dir = Path(args.ref), Path(args.cand)

    # --- allocator gate: same defect as the capture path ----------------------
    alloc = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
    if "expandable_segments" in alloc:
        print(f"\n!! PYTORCH_CUDA_ALLOC_CONF={alloc!r} — expandable_segments corrupts "
              f"retained tensors on this stack. Unset it.", file=sys.stderr)
        sys.exit(9)

    # --- tokenizer identity gate ---------------------------------------------
    # Both distributions must be conditioned on byte-identical prefixes. If the
    # abliterated checkpoint carries a different tokenizer or chat template, every
    # KL below is a comparison of two different questions.
    tok_files = ["tokenizer.json", "tokenizer_config.json"]
    digests = {f: (file_md5(ref_dir / f), file_md5(cand_dir / f)) for f in tok_files}
    mismatched = {f: v for f, v in digests.items() if v[0] != v[1]}
    if mismatched and not args.skip_gates:
        print(f"\n!! tokenizer differs between ref and cand: {sorted(mismatched)} — the two "
              f"models would be conditioned on different prefixes and the KL would be "
              f"meaningless.", file=sys.stderr)
        sys.exit(6)
    print("tokenizer identity: " + ("MATCH" if not mismatched else f"MISMATCH {sorted(mismatched)}"))

    # --- corpus ---------------------------------------------------------------
    harmless, harmful, prov = load_evaluation(
        n_harmless=args.n_harmless, n_harmful=args.n_harmful, seed=args.seed,
        calib_harmless_n=args.calib_harmless_n,
        calib_harmless_seed=args.calib_harmless_seed)
    print(f"eval corpus: {len(harmless)} harmless (held out from calibration), "
          f"{len(harmful)} harmful (reserved test split)")
    if args.seed == args.calib_harmless_seed:
        print("!! --seed equals the capture's harmless seed; the exclusion still holds "
              "(indices are removed from the pool) but say so in the writeup.", file=sys.stderr)

    prompts = harmless + harmful
    split = len(harmless)

    from transformers import AutoModelForCausalLM, AutoTokenizer

    if args.dry_run:
        # Prove the rendering and the tokenizer round-trip on the real checkpoint
        # before spending a seat-down window. Tokenizer only — no weights touched.
        tok = AutoTokenizer.from_pretrained(ref_dir)
        for mode, thinking in MODES.items():
            text = render(tok, prompts[0], thinking)
            n_tok = len(tok(text)["input_ids"])
            print(f"\n  [{mode}] {n_tok} tokens, tail: {text[-60:]!r}")
        lens = [len(tok(render(tok, p, False))["input_ids"]) for p in prompts]
        print(f"\n  prompt lengths (answer mode): min {min(lens)} median "
              f"{int(statistics.median(lens))} max {max(lens)}")
        print(f"  first harmless: {harmless[0][:90]!r}")
        print(f"  first harmful:  {harmful[0][:90]!r}")
        print("\ndry-run complete — corpus, gates and rendering verified; "
              "nothing loaded, no GPU touched.")
        return

    cache_dir = Path(args.cache) if args.cache else Path(args.out).parent
    cache_dir.mkdir(parents=True, exist_ok=True)
    cache_path = {lbl: cache_dir / f"{Path(args.out).stem}.{lbl}.pt"
                  for lbl in ("ref", "cand")}

    def run_stage(model_dir: Path, label: str, measure_floor: bool):
        print(f"\n=== {label}: {model_dir}")

        # --- headroom gate, BEFORE the load --------------------------------
        # `device_map="auto"` does not fail when the card is too small; it
        # quietly spills modules to CPU, and only the residency gate below would
        # catch that — after paying for the load. Worse, the second stage runs on
        # whatever the first stage gave back, so a release regression shows up
        # here as offloading rather than as an error. Size the requirement from
        # the checkpoint's own headers so it stays true if the weights change.
        need = weights_mib(model_dir)
        free = free_vram_mib()
        print(f"  weights {need:,.0f} MiB (text-only, from shard headers); "
              f"free {free:,.0f} MiB")
        if free < need + 2048:
            print(f"\n!! insufficient VRAM: {free:,.0f} MiB free, need {need:,.0f} + 2048 "
                  f"MiB headroom. Both GPU0 seats must be stopped — freeing only one "
                  f"leaves ~50,900 MiB and this model's text weights are ~51,300 MiB. "
                  f"If this is the second stage, the first stage did not release.",
                  file=sys.stderr)
            sys.exit(10)

        collected, floor, vocab = _load_and_collect(model_dir, label, measure_floor)
        torch.save(
            {"logprobs": {m: collected[m]["logprobs"] for m in MODES},
             "ids": {m: collected[m]["ids"] for m in MODES},
             "floor": floor, "vocab": vocab, "model_dir": str(model_dir)},
            cache_path[label])
        print(f"  cached -> {cache_path[label]} "
              f"({cache_path[label].stat().st_size / 2**20:,.0f} MiB)")

    def _load_and_collect(model_dir: Path, label: str, measure_floor: bool):
        tok = AutoTokenizer.from_pretrained(model_dir)
        model = AutoModelForCausalLM.from_pretrained(
            model_dir, dtype=torch.bfloat16, device_map="auto", attn_implementation="sdpa")
        model.eval()
        device = next(model.parameters()).device

        # --- residency gate ---------------------------------------------------
        # Sharding this architecture across the two Blackwells zeroes the residual
        # stream past the device boundary; the logits would decode to garbage while
        # every gate below the boundary still passed. Diagnosed 2026-08-20.
        #
        # Read the placement off the PARAMETERS, not off `hf_device_map`. The map
        # is empty whenever transformers puts the whole model on one device, so a
        # map-based check reports "(unsharded)" both when everything is fine and
        # when there is no map to inspect — it cannot fail, which makes it not a
        # gate. Parameter devices are ground truth in every case.
        devices = {str(p.device) for p in model.parameters()}
        offloaded = sorted(d for d in devices if d.startswith(("cpu", "meta", "disk")))
        if len(devices) > 1 or offloaded:
            print(f"\n!! residency gate FAILED — parameters span {sorted(devices)}. "
                  f"Sharding this architecture zeroes the residual stream past the device "
                  f"boundary and the logits decode to garbage. Fix: CUDA_VISIBLE_DEVICES=0 "
                  f"with BOTH GPU0 seats stopped.", file=sys.stderr)
            sys.exit(8)
        print(f"  residency: all parameters on {sorted(devices)[0]}, "
              f"free VRAM after load: {free_vram_mib():,.0f} MiB")

        collected = collect(model, tok, prompts, device, label)

        floor = None
        if measure_floor and args.noise_floor_n:
            # Re-run a subset through the SAME model. This is not a formality: it
            # validates the whole scoring path end to end. A deterministic stack
            # must return exactly 0 here, so any non-zero value indicts this
            # script (or the stack's determinism) before it indicts the surgery.
            sub = prompts[:args.noise_floor_n]
            again = collect(model, tok, sub, device, f"{label}-repeat")
            floor = {}
            for mode in MODES:
                k, _, tv, top1 = divergences(
                    collected[mode]["logprobs"][:len(sub)], again[mode]["logprobs"])
                floor[mode] = {"n": len(sub), "kl_mean": float(k.mean()),
                               "kl_max": float(k.max()), "tv_max": float(tv.max()),
                               "top1_agreement": float(top1.double().mean())}
            print(f"  noise floor (self-KL): "
                  + ", ".join(f"{m} max {v['kl_max']:.3e}" for m, v in floor.items()))

        return collected, floor, collected["answer"]["logprobs"].shape[-1]

    # --- ONE PROCESS PER MODEL -----------------------------------------------
    #
    # This is not fastidiousness; it is the only mechanism that works. Both
    # in-process teardowns were tried and MEASURED on 2026-08-20:
    #
    #   `del model` + `gc.collect()` + `torch.cuda.empty_cache()`  -> 45,287 MiB free
    #   the same, with the model confined to an inner frame that exits -> 45,287 MiB free
    #
    # i.e. the ~51,300 MiB of weights were still resident both times. The first
    # (pre-gate) run only survived because PyTorch's allocator hit OOM during the
    # second load, ran a collection itself, and retried — the second model landed
    # on the card by rescue, not by design. Relying on that is how you end up
    # silently offloaded to CPU, which on this architecture does not error: it
    # zeroes the residual stream past the boundary and returns confident garbage.
    #
    # A process exit releases the CUDA context unconditionally, so each model gets
    # its own. The stages hand their first-token log-probs to disk (~715 MiB per
    # model) and the parent scores from the caches. Reusable, too: re-measuring a
    # different candidate against this same reference skips the ref stage.
    if args.stage in ("ref", "cand"):
        run_stage(ref_dir if args.stage == "ref" else cand_dir,
                  args.stage, measure_floor=(args.stage == "ref"))
        return

    if args.stage == "all":
        import subprocess
        for label in ("ref", "cand"):
            cmd = [sys.executable, str(Path(__file__).resolve()),
                   "--ref", str(ref_dir), "--cand", str(cand_dir), "--out", args.out,
                   "--n-harmless", str(args.n_harmless), "--n-harmful", str(args.n_harmful),
                   "--seed", str(args.seed),
                   "--calib-harmless-n", str(args.calib_harmless_n),
                   "--calib-harmless-seed", str(args.calib_harmless_seed),
                   "--noise-floor-n", str(args.noise_floor_n),
                   "--cache", str(cache_dir), "--stage", label]
            if args.skip_gates:
                cmd.append("--skip-gates")
            rc = subprocess.run(cmd).returncode
            if rc != 0:
                print(f"\n!! {label} stage exited {rc}; not scoring a partial run.",
                      file=sys.stderr)
                sys.exit(rc)
            print(f"  [{label} stage process exited; CUDA context released]")

    # --- load the stage caches ------------------------------------------------
    missing = [p for p in cache_path.values() if not p.exists()]
    if missing:
        print(f"\n!! missing stage cache(s): {missing} — run --stage ref and "
              f"--stage cand first.", file=sys.stderr)
        sys.exit(11)
    blobs = {lbl: torch.load(p, weights_only=False) for lbl, p in cache_path.items()}
    ref = {m: {"logprobs": blobs["ref"]["logprobs"][m], "ids": blobs["ref"]["ids"][m]}
           for m in MODES}
    cand = {m: {"logprobs": blobs["cand"]["logprobs"][m], "ids": blobs["cand"]["ids"][m]}
            for m in MODES}
    noise_floor = blobs["ref"]["floor"]
    ref_vocab, cand_vocab = blobs["ref"]["vocab"], blobs["cand"]["vocab"]

    if ref_vocab != cand_vocab and not args.skip_gates:
        print(f"\n!! vocab size differs: ref {ref_vocab} vs cand {cand_vocab}", file=sys.stderr)
        sys.exit(6)

    # --- tokenization equality gate ------------------------------------------
    for mode in MODES:
        for i, (a, b) in enumerate(zip(ref[mode]["ids"], cand[mode]["ids"])):
            if a != b and not args.skip_gates:
                print(f"\n!! prompt {i} ({mode}) tokenized differently between the two "
                      f"stages ({len(a)} vs {len(b)} tokens) — the KL would compare "
                      f"distributions over different prefixes.", file=sys.stderr)
                sys.exit(7)

    # --- score ----------------------------------------------------------------
    results = {}
    for mode in MODES:
        kl_f, kl_r, tv, top1 = divergences(ref[mode]["logprobs"], cand[mode]["logprobs"])
        results[mode] = {
            "harmless": summarize(kl_f[:split], kl_r[:split], tv[:split], top1[:split]),
            "harmful": summarize(kl_f[split:], kl_r[split:], tv[split:], top1[split:]),
        }
        hl = results[mode]["harmless"]["kl_median"]
        hf = results[mode]["harmful"]["kl_median"]
        results[mode]["selectivity_median_ratio"] = (hf / hl) if hl > 0 else None

    payload = {
        "ref": str(ref_dir), "cand": str(cand_dir),
        "metric": "first-token KL(ref || cand), full vocabulary, float64",
        "measured_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "corpus": prov,
        "tokenizer_digests": {f: {"ref": v[0], "cand": v[1]} for f, v in digests.items()},
        "vocab": ref_vocab,
        "noise_floor_self_kl": noise_floor,
        "modes": results,
        "env": {"torch": torch.__version__,
                "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
                "alloc_conf": alloc or None},
    }
    Path(args.out).write_text(json.dumps(payload, indent=2))

    print("\n" + "=" * 72)
    print(f"first-token KL(ref || cand)   ref={ref_dir.name}")
    print(f"                              cand={cand_dir.name}")
    for mode in MODES:
        r = results[mode]
        print(f"\n  [{mode} mode]")
        for cls in ("harmless", "harmful"):
            s = r[cls]
            print(f"    {cls:<9} n={s['n']:<4} median {s['kl_median']:.4f}  "
                  f"mean {s['kl_mean']:.4f}  p95 {s['kl_p95']:.4f}  max {s['kl_max']:.4f}  "
                  f"top1-agree {s['top1_agreement']:.1%}")
        ratio = r["selectivity_median_ratio"]
        shown = f"{ratio:.1f}x" if ratio else "n/a"
        print(f"    selectivity (harmful/harmless median KL): {shown}")
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()