#!/usr/bin/env python3
"""Graft an IN-BAND abliterated MTP head onto a Heretic-exported trunk.

THE PROBLEM
-----------
Heretic silently drops the MTP head on save (p-e-w/heretic PR #317, declined).
Measured on our own export: source 1199 tensors -> export 1184, all 15 `mtp.*`
gone, vision 333/333 intact, exit 0, no warning. That is also why
`MuXodious/Qwen3.8-27B-absolute-heresy` ships an MTP head byte-identical to its
base: not a design choice, the same bug.

WHY NOT JUST COPY THE ORIGINAL HEAD BACK
----------------------------------------
Because the head is conditioned on the trunk's hidden states, and Heretic just
removed a direction from that trunk. A pristine head grafted onto an abliterated
trunk is reading a distribution it was not trained on — which is exactly the
"byte-identical != behaviour-preserving" point. Our own in-band build measured
59.1% MTP acceptance against the grafted incumbent's 47.2%, so the in-band edit
is worth reproducing rather than shortcutting.

HOW THE DIRECTION IS RECOVERED
------------------------------
Heretic ran with `direction_scope = 0`, i.e. ONE shared direction across layers,
scaled per layer by its weight kernel. Directional ablation is

    W_out = W_in - w * d (d^T W_in)

so `delta = W_out - W_in` is rank-1 with left singular vector d. We therefore
recover d by SVD of a trunk delta, and the per-layer weight w from the ratio of
norms. This is verified, not assumed: the script asserts the delta really is
rank-1 (sigma2/sigma1 below a threshold) before trusting anything, and
cross-checks d against several layers for agreement.

SAFETY SCREENS CARRIED OVER FROM THE ROBINSON RECIPE
----------------------------------------------------
Dim 3994 is this architecture's massive-activation / attention-sink dimension.
Orthogonalising it out bricks the model, so its share of the direction's energy
is reported and gated. Vision tensors are never touched.
"""
from __future__ import annotations

import argparse
import json
import shutil
import struct
import sys
from glob import glob
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file

SINK_DIM = 3994
SINK_MAX = 0.01           # >1% of direction energy in the sink dim -> refuse
RANK1_MAX = 0.02          # sigma2/sigma1 above this -> the delta is not rank-1
MTP_WRITERS = ("mtp.layers.0.self_attn.o_proj.weight",
               "mtp.layers.0.mlp.down_proj.weight")


def shard_map(d: Path):
    """tensor key -> shard path, read from headers (no mmap; safe_open ENOMEMs on ZFS)."""
    m = {}
    for s in sorted(glob(str(d / "*.safetensors"))):
        with open(s, "rb") as f:
            n = struct.unpack("<Q", f.read(8))[0]
            hdr = json.loads(f.read(n))
        for k in hdr:
            if k != "__metadata__":
                m[k] = s
    return m


def load(path, key):
    with safe_open(path, framework="pt") as f:
        return f.get_tensor(key)


def recover_direction(src: Path, exp: Path, probe_layers, verbose=True):
    """Recover (unit direction, {layer: weight}) from the trunk deltas."""
    smap, emap = shard_map(src), shard_map(exp)
    dirs, weights = [], {}
    for li in probe_layers:
        key = f"model.language_model.layers.{li}.mlp.down_proj.weight"
        if key not in smap:
            key = f"model.layers.{li}.mlp.down_proj.weight"
        if key not in smap or key not in emap:
            continue
        Ws = load(smap[key], key).float()
        We = load(emap[key], key).float()
        delta = We - Ws
        if delta.abs().max() == 0:
            continue
        # rank-1 check on the delta
        U, S, _ = torch.linalg.svd(delta, full_matrices=False)
        ratio = (S[1] / S[0]).item()
        if ratio > RANK1_MAX:
            print(f"!! layer {li}: delta is NOT rank-1 (s2/s1={ratio:.4f}). Heretic's edit "
                  f"is not a simple directional ablation here — refusing to extrapolate.",
                  file=sys.stderr)
            sys.exit(3)
        d = U[:, 0]
        d = d / d.norm()
        # w from || delta || / || d (d^T Ws) ||
        proj = torch.outer(d, d @ Ws)
        w = (delta.norm() / proj.norm()).item()
        # sign: delta should oppose the projection
        if (delta * proj).sum() > 0:
            w = -w
        dirs.append(d)
        weights[li] = abs(w)
        if verbose:
            print(f"   layer {li:>2}: rank-1 ok (s2/s1 {ratio:.5f})   weight {abs(w):.4f}")
    if not dirs:
        print("!! no modified trunk layers found — is this really a Heretic export?",
              file=sys.stderr)
        sys.exit(4)
    # agreement between layers (scope=0 means they should be the same direction)
    ref = dirs[0]
    agree = [min(abs(float(ref @ x)), 1.0) for x in dirs]
    worst = min(agree)
    if verbose:
        print(f"   cross-layer |cos| agreement: min {worst:.4f} over {len(dirs)} layers")
    if worst < 0.95:
        print(f"!! layers disagree on the direction (min |cos| {worst:.3f}). That implies "
              f"per-layer directions, not the single shared direction this run used. "
              f"Refusing to pick one.", file=sys.stderr)
        sys.exit(5)
    # average, re-sign to ref, normalise
    acc = torch.zeros_like(ref)
    for x in dirs:
        acc += x if (ref @ x) >= 0 else -x
    d = acc / acc.norm()
    return d, weights


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", required=True, help="original bf16 (has the MTP head)")
    ap.add_argument("--trunk", required=True, help="Heretic export (MTP dropped)")
    ap.add_argument("--out", required=True, help="new dir; trunk is hardlinked, MTP added")
    ap.add_argument("--mtp-weight-scale", type=float, default=1.0,
                    help="multiplier on the trunk's deepest-layer ablation weight. 1.0 = "
                         "match the trunk (what our 59.1%%-acceptance build did). The "
                         "2026-08-21 panel suggests sweeping 0.5-0.7 as well.")
    ap.add_argument("--probe-layers", default="20,30,40,50,60")
    ap.add_argument("--pristine", action="store_true",
                    help="graft the head VERBATIM, no in-band ablation. This is what "
                         "heresy ships (~47%% acceptance vs our in-band 59.1%%). Use it "
                         "when the in-band edit is unsafe or not on the critical path: "
                         "MTP is a decode-throughput feature, not an output-quality one.")
    args = ap.parse_args()

    src, trunk, out = Path(args.source), Path(args.trunk), Path(args.out)
    print(f"source: {src}\ntrunk : {trunk}\nout   : {out}\n")

    if args.pristine:
        # No direction needed, so no sink screen applies: nothing is being
        # orthogonalised. The head arrives byte-identical to the base.
        print("PRISTINE graft — head copied verbatim, no ablation applied.")
        print("  Trade-off on the record: the head is conditioned on trunk hidden")
        print("  states that Heretic has since altered, so acceptance will sit nearer")
        print("  the grafted ~47% than our in-band 59.1%. Byte-identical is NOT")
        print("  behaviour-preserving; it is simply the safe option here.\n")
        d, w_mtp = None, 0.0
    else:
        print("recovering Heretic's refusal direction from the trunk deltas:")
        probes = [int(x) for x in args.probe_layers.split(",")]
        d, weights = recover_direction(src, trunk, probes)

        sink = (d[SINK_DIM] ** 2 / (d @ d)).item()
        print(f"\n   sink dim {SINK_DIM} energy: {sink:.4%}  (threshold {SINK_MAX:.1%})")
        if sink > SINK_MAX:
            print("!! direction is sink-dominated; orthogonalising it out bricks this "
                  "architecture. Refusing.\n"
                  "   NOTE: Heretic itself survives this direction because it uses "
                  "magnitude-preserving ablation (row_normalization=FULL); plain "
                  "projection here does not, so the screen is correct for THIS "
                  "operation even though the trunk is fine. Use --pristine, or "
                  "implement MPOA.", file=sys.stderr)
            sys.exit(6)

        deepest = max(weights)
        w_mtp = weights[deepest] * args.mtp_weight_scale
        print(f"   deepest probed layer {deepest} weight {weights[deepest]:.4f}"
              f"  x scale {args.mtp_weight_scale}  ->  MTP weight {w_mtp:.4f}")

    # --- build the MTP shard -------------------------------------------------
    smap = shard_map(src)
    mtp_keys = sorted(k for k in smap if k.startswith("mtp."))
    print(f"\ngrafting {len(mtp_keys)} mtp.* tensors ({len(MTP_WRITERS)} abliterated in-band):")
    tensors, edited = {}, 0
    for k in mtp_keys:
        t = load(smap[k], k)
        if k in MTP_WRITERS and d is not None:
            W = t.float()
            dd = d.to(W.dtype)
            W = W - w_mtp * torch.outer(dd, dd @ W)
            before = t.float()
            t = W.to(t.dtype)
            delta = (t.float() - before).norm().item()
            print(f"   [ABLITERATED] {k}  ||delta|| {delta:.4f}")
            edited += 1
        else:
            print(f"   [verbatim]    {k}")
        tensors[k] = t
    expect = 0 if d is None else len(MTP_WRITERS)
    if edited != expect:
        print(f"!! expected {expect} residual writers edited, got {edited}", file=sys.stderr)
        sys.exit(7)

    # --- assemble output: hardlink trunk shards, add the MTP shard -----------
    out.mkdir(parents=True, exist_ok=True)
    for f in trunk.iterdir():
        if f.is_file():
            dst = out / f.name
            if not dst.exists():
                try:
                    dst.hardlink_to(f)
                except Exception:
                    shutil.copy2(f, dst)
    mtp_file = "model-mtp.safetensors"
    save_file(tensors, str(out / mtp_file), metadata={"format": "pt"})
    print(f"\nwrote {mtp_file}")

    # merge the index (mirrors unsloth's layout)
    idx_p = out / "model.safetensors.index.json"
    idx = json.loads(idx_p.read_text())
    for k in mtp_keys:
        idx["weight_map"][k] = mtp_file
    idx["metadata"]["total_size"] = idx["metadata"].get("total_size", 0) + \
        sum(t.numel() * t.element_size() for t in tensors.values())
    idx_p.write_text(json.dumps(idx, indent=2))
    print("merged index")

    # restore the config files Heretic omits (the wrapper-class omission that
    # crash-loops a seat with "Can't load image processor")
    for name in ("preprocessor_config.json", "video_preprocessor_config.json", "vocab.json"):
        s = src / name
        if s.exists() and not (out / name).exists():
            shutil.copy2(s, out / name)
            print(f"restored {name}")

    print("\ndone.")


if __name__ == "__main__":
    main()