#!/usr/bin/env python3
"""Compare a candidate checkpoint's MTP head against a known-good reference head.

WHY THIS EXISTS
---------------
The expensive gate before quantizing a new Qwen3.8-27B candidate is "does its MTP
head actually work?" — presence is not acceptance, and a dead head silently costs
the entire +18% decode case. Measuring acceptance for real needs the model resident
in VRAM; at bf16 that is ~56 GB, which on a full 97.9 GB card means downing a
SECOND seat, not just `gen`.

This sidesteps that for the common case. Most community abliterations never touch
`mtp.*` at all: the `Qwen3_5ForConditionalGeneration` wrapper class does not load
the MTP head, so PEFT merges, Heretic runs, and llm-compressor passes all leave it
exactly as it came from the base. If a candidate's 15 `mtp.*` tensors are
numerically identical to a head we have already measured in production, its MTP is
that head — and the acceptance test is redundant.

Reference to compare against: `/tank/aimodels/qwen38-27b-uncensored-bf16`
(`model-mtp.safetensors`). Per its PROVENANCE that head was grafted verbatim from
the base, and it measures **47.7% acceptance** on the live gen seat through our
exact mixed-quant pipeline. That makes it a known-good baseline, not a guess.

READ THE RESULT HONESTLY
------------------------
- **IDENTICAL** -> the candidate carries the pristine base head. The bf16 acceptance
  gate buys nothing; go to quant and verify acceptance on the quantized build inside
  the freed `gen` budget (~22 GB, no second seat).
- **DIFFERENT** -> something edited the head. That is NOT automatically bad (an
  abliteration that deliberately includes `mtp.*` is a legitimate design — see
  hotdogs/Qwen3.8-27B-abliterated, which edits 2 mtp tensors on purpose), but it
  means the head is no longer one we have measured. Run the real acceptance gate.
- **MISSING** -> the head was dropped. Known failure mode; needs a graft.

CPU only. Reads just the shard(s) holding `mtp.*` — no full model load.
"""

import argparse
import hashlib
import json
import sys
from pathlib import Path

from safetensors import safe_open


def load_mtp(model_dir: Path) -> dict:
    """Return {tensor_name: tensor} for every mtp.* tensor, reading only the
    shards that actually hold them."""
    index = model_dir / "model.safetensors.index.json"
    if index.exists():
        weight_map = json.loads(index.read_text())["weight_map"]
        names = [k for k in weight_map if k.startswith("mtp")]
        shards = sorted({weight_map[n] for n in names})
    else:
        # unsharded checkpoint
        shards = ["model.safetensors"]
        names = None

    out = {}
    for shard in shards:
        path = model_dir / shard
        if not path.exists():
            raise SystemExit(f"missing shard: {path}")
        with safe_open(path, framework="pt") as f:
            for key in f.keys():
                if key.startswith("mtp"):
                    out[key] = f.get_tensor(key)
    if names is not None and len(out) != len(names):
        print(f"  warning: index listed {len(names)} mtp tensors, read {len(out)}",
              file=sys.stderr)
    return out


def digest(tensor) -> str:
    """SHA-256 over the raw tensor bytes. Dtype-sensitive by design — a head
    stored at a different precision is not the same head for our purposes.

    Goes through a uint8 reinterpret rather than `.numpy()`: numpy has no
    bfloat16, and these checkpoints are bf16, so the direct route raises
    `TypeError: Got unsupported ScalarType BFloat16`.
    """
    import torch

    flat = tensor.contiguous().flatten()
    return hashlib.sha256(flat.view(torch.uint8).numpy().tobytes()).hexdigest()


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("candidate", type=Path, help="candidate model dir")
    ap.add_argument("reference", type=Path,
                    help="known-good reference model dir (e.g. qwen38-27b-uncensored-bf16)")
    args = ap.parse_args()

    cand = load_mtp(args.candidate)
    ref = load_mtp(args.reference)

    print(f"candidate : {args.candidate}  ({len(cand)} mtp tensors)")
    print(f"reference : {args.reference}  ({len(ref)} mtp tensors)")
    print()

    if not cand:
        print("VERDICT: MISSING — candidate has no mtp.* tensors. Needs a graft.")
        return 2

    only_cand = sorted(set(cand) - set(ref))
    only_ref = sorted(set(ref) - set(cand))
    if only_cand or only_ref:
        print("  tensor-name mismatch:")
        for n in only_cand:
            print(f"    only in candidate: {n}")
        for n in only_ref:
            print(f"    only in reference: {n}")
        print()

    differing = []
    for name in sorted(set(cand) & set(ref)):
        c, r = cand[name], ref[name]
        if c.dtype != r.dtype or c.shape != r.shape:
            differing.append((name, f"dtype/shape {c.dtype}{tuple(c.shape)} vs "
                                    f"{r.dtype}{tuple(r.shape)}"))
            continue
        if digest(c) != digest(r):
            # quantify it — a tiny delta is a different story from a rewritten head
            delta = (c.float() - r.float()).abs().max().item()
            denom = r.float().abs().max().item() or 1.0
            differing.append((name, f"max|Δ| = {delta:.6g}  (rel {delta / denom:.3%})"))

    for name, why in differing:
        print(f"  DIFFERS  {name}: {why}")

    print()
    if not differing and not only_cand and not only_ref:
        print("VERDICT: IDENTICAL — candidate carries the reference MTP head verbatim.")
        print("  The bf16 acceptance gate is redundant: this head is already measured")
        print("  at 47.7% acceptance in production through our mixed-quant pipeline.")
        print("  Proceed to quant; verify acceptance on the quantized build.")
        return 0

    print(f"VERDICT: DIFFERENT — {len(differing)} of {len(cand)} mtp tensors diverge.")
    print("  The head is not one we have measured. Run the real bf16 acceptance gate")
    print("  (~56 GB resident) before spending quant GPU time.")
    return 1


if __name__ == "__main__":
    raise SystemExit(main())