#!/usr/bin/env python3
"""Graft the Qwen3.6-27B MTP (multi-token-prediction) head into Heretic2-Thinking.

DavidAU's Heretic2-Uncensored-Finetune-Thinking finetune dropped the MTP head
(config declares mtp_num_hidden_layers=1 but ships ZERO mtp.* weight tensors),
so the base Qwen/Qwen3.6-27B's 15 MTP tensors must be grafted back in BF16
before NVFP4 quantization — this is what lets the served seat use `qwen3_5_mtp`
speculative decoding (n=3), the pantheon-27b-mtp-nvfp4 pattern proven on our
Blackwell. (R36 fast-seat spike, brokkr calib spec 2026-07-14.)

CPU-only safetensors surgery — no GPU / no quant-window needed. Space-efficient:
symlinks Heretic2's large shards, adds one small model-mtp.safetensors, and
writes a merged index. Idempotent-ish: refuses to clobber a non-empty OUT_DIR.

Usage (on ana-ml2):
    python3 graft_mtp.py \
        --heretic2 /tank/aimodels/huggingface/hub/models--DavidAU--Qwen3.6-27B-Heretic2-Uncensored-Finetune-Thinking/snapshots/<hash> \
        --base     /tank/aimodels/huggingface/hub/models--Qwen--Qwen3.6-27B/snapshots/<hash> \
        --out      /tank/aimodels/heretic2-mtp-bf16
"""
import argparse
import json
import os
import shutil
import sys

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

MTP_SHARD = "model-mtp.safetensors"
# MTP-related config keys we ensure are present on the grafted model (copied from
# base if Heretic2's config is missing any).
MTP_CFG_KEYS = ("mtp_num_hidden_layers", "mtp_use_dedicated_embeddings")


def _is_mtp(name: str) -> bool:
    n = name.lower()
    return n.startswith("mtp.") or "mtp." in n or "nextn" in n


def _tensor_nbytes(t: torch.Tensor) -> int:
    return t.numel() * t.element_size()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--heretic2", required=True, help="Heretic2 BF16 snapshot dir (quant target)")
    ap.add_argument("--base", required=True, help="Qwen/Qwen3.6-27B snapshot dir (MTP source)")
    ap.add_argument("--out", required=True, help="output dir for the grafted BF16 model")
    args = ap.parse_args()

    her, base, out = args.heretic2, args.base, args.out
    for d in (her, base):
        if not os.path.isfile(os.path.join(d, "model.safetensors.index.json")):
            print(f"ERROR: no index.json in {d}", file=sys.stderr)
            return 2
    if os.path.isdir(out) and os.listdir(out):
        print(f"ERROR: {out} exists and is non-empty — refusing to clobber", file=sys.stderr)
        return 2
    os.makedirs(out, exist_ok=True)

    her_idx = json.load(open(os.path.join(her, "model.safetensors.index.json")))
    base_idx = json.load(open(os.path.join(base, "model.safetensors.index.json")))

    # sanity: Heretic2 must NOT already have MTP weights; base MUST.
    her_mtp = [k for k in her_idx["weight_map"] if _is_mtp(k)]
    base_mtp = [k for k in base_idx["weight_map"] if _is_mtp(k)]
    if her_mtp:
        print(f"ERROR: Heretic2 already has {len(her_mtp)} mtp tensors — graft not needed", file=sys.stderr)
        return 2
    if not base_mtp:
        print("ERROR: base has no mtp tensors — wrong source model", file=sys.stderr)
        return 2
    print(f"grafting {len(base_mtp)} MTP tensors from base into Heretic2 ({len(her_idx['weight_map'])} tensors)")

    # 1) stage Heretic2: symlink big shards, copy everything else.
    for fn in sorted(os.listdir(her)):
        src = os.path.join(her, fn)
        if not os.path.isfile(src):
            continue
        dst = os.path.join(out, fn)
        if fn.endswith(".safetensors"):
            os.symlink(os.path.realpath(src), dst)
        elif fn != "model.safetensors.index.json":  # index rewritten below
            shutil.copy2(src, dst)

    # 2) pull the MTP tensors out of base's shards into one new shard (BF16 preserved).
    base_mtp_shards = sorted({base_idx["weight_map"][k] for k in base_mtp})
    mtp_tensors = {}
    for shard in base_mtp_shards:
        with safe_open(os.path.join(base, shard), framework="pt") as f:
            for name in f.keys():
                if _is_mtp(name):
                    mtp_tensors[name] = f.get_tensor(name)
    assert len(mtp_tensors) == len(base_mtp), f"expected {len(base_mtp)} mtp tensors, got {len(mtp_tensors)}"
    dtypes = {str(t.dtype) for t in mtp_tensors.values()}
    print(f"  extracted {len(mtp_tensors)} mtp tensors, dtypes={dtypes}")
    save_file(mtp_tensors, os.path.join(out, MTP_SHARD), metadata={"format": "pt"})

    # 3) merged index: Heretic2 weight_map + the mtp tensors -> the new shard.
    new_map = dict(her_idx["weight_map"])
    added_bytes = 0
    for name, t in mtp_tensors.items():
        new_map[name] = MTP_SHARD
        added_bytes += _tensor_nbytes(t)
    meta = dict(her_idx.get("metadata", {}))
    if "total_size" in meta:
        meta["total_size"] = int(meta["total_size"]) + added_bytes
    out_idx = {"metadata": meta, "weight_map": new_map}
    json.dump(out_idx, open(os.path.join(out, "model.safetensors.index.json"), "w"), indent=2)

    # 4) ensure config has complete MTP config (copy from base if Heretic2 lacks a key).
    cfg_path = os.path.join(out, "config.json")
    cfg = json.load(open(cfg_path))
    base_cfg = json.load(open(os.path.join(base, "config.json")))

    def _get(d, k):
        return d.get(k, d.get("text_config", {}).get(k))

    changed = []
    for k in MTP_CFG_KEYS:
        if _get(cfg, k) is None and _get(base_cfg, k) is not None:
            cfg[k] = _get(base_cfg, k)
            changed.append(k)
    if changed:
        json.dump(cfg, open(cfg_path, "w"), indent=2)
    print(f"  config MTP keys filled from base: {changed or 'none (already complete)'}")

    total = len(new_map)
    print(f"DONE -> {out}")
    print(f"  total tensors: {total} (Heretic2 {len(her_idx['weight_map'])} + MTP {len(mtp_tensors)})")
    print(f"  new shard: {MTP_SHARD} (+{added_bytes/1e9:.2f} GB)")
    print("  next: NVFP4 quant (quant_nvfp4.py) — keeps mtp.*/GDN/vision/lm-head/norms in BF16")
    return 0


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