#!/home/olegk/venv/vllm/bin/python3
"""
🌸 splice_nvfp4_model.py - Splicer for Calibrated Experts-Only NVFP4 Drafter
Channels Nikola's Sparkly Village Charm for Haga's Seeker Quest! ✨🚀

Non-destructively links base model NVFP4 assets, updates config.json with
mtp_quantization_config, and patches model.safetensors.index.json for the
calibrated NVFP4 MTP drafter.
"""

import json
import os
import shutil
import torch
import safetensors.torch
from safetensors import safe_open


def splice_nvfp4(
    base_dir: str = "models/mazinb/Qwen3.8-Flash-Next-Uncensored-NVFP4",
    shadow_dir: str = "models/spliced/Qwen3.8-Flash-Next-MTP-NVFP4",
    drafter_file: str = "models/spliced/Qwen3.8-Flash-Next-MTP-NVFP4/model_mtp_tuned.safetensors",
):
    base_dir = os.path.abspath(base_dir)
    shadow_dir = os.path.abspath(shadow_dir)
    drafter_file = os.path.abspath(drafter_file)

    print("=" * 75)
    print("🌸 Splicing Calibrated NVFP4 MTP Model ✨🚀")
    print(f"Base Model:  {base_dir}")
    print(f"Shadow Dir:  {shadow_dir}")
    print(f"Drafter:     {drafter_file}")
    print("=" * 75)

    assert os.path.exists(drafter_file), f"Drafter file {drafter_file} does not exist!"

    # 1. Hardlink base model assets (excluding index, config, and drafter)
    skip_files = {"model.safetensors.index.json", "config.json", "model_mtp_tuned.safetensors"}
    linked = 0
    for fname in os.listdir(base_dir):
        if fname in skip_files:
            continue
        src = os.path.join(base_dir, fname)
        dst = os.path.join(shadow_dir, fname)
        if os.path.exists(dst) or os.path.islink(dst):
            os.remove(dst)
        if os.path.isdir(src):
            os.symlink(src, dst)
        else:
            os.link(src, dst)
        linked += 1
    print(f"🔗 Linked {linked} base model assets via hardlinks (0 bytes disk space used).")

    # 2. Configure config.json with mtp_quantization_config
    base_cfg_path = os.path.join(base_dir, "config.json")
    with open(base_cfg_path, "r") as f:
        cfg = json.load(f)

    mtp_quant_config = {
        "quant_method": "modelopt",
        "quant_algo": "NVFP4",
        "config_groups": {
            "group_0": {
                "input_activations": {
                    "dynamic": False,
                    "group_size": 16,
                    "num_bits": 4,
                    "type": "float",
                },
                "targets": ["Linear"],
                "weights": {
                    "dynamic": False,
                    "group_size": 16,
                    "num_bits": 4,
                    "type": "float",
                },
            }
        },
        "ignore": [
            "self_attn",
            "fc_embedding",
            "fc_hidden",
            "shared_expert",
            "mlp.gate",
            "shared_expert_gate",
            "hyper_connection",
            "norm",
            "embed_tokens",
            "lm_head",
        ],
        "producer": {"name": "modelopt", "version": "0.46.0-compatible-rtn"},
    }

    cfg["mtp_quantization_config"] = mtp_quant_config
    if "text_config" in cfg and isinstance(cfg["text_config"], dict):
        cfg["text_config"]["mtp_quantization_config"] = mtp_quant_config

    shadow_cfg_path = os.path.join(shadow_dir, "config.json")
    if os.path.exists(shadow_cfg_path):
        os.remove(shadow_cfg_path)
    with open(shadow_cfg_path, "w") as f:
        json.dump(cfg, f, indent=2)
    print("📝 Wrote updated config.json with mtp_quantization_config.")

    # 3. Pack Q output gate if q_proj is un-gated [6144, 2560]
    q_key = "mtp.layers.0.self_attn.q_proj.weight"
    drafter = safetensors.torch.load_file(drafter_file)
    if q_key in drafter and drafter[q_key].shape[0] == 6144:
        orig_sf11 = os.path.join(base_dir, "model-bf16-00011.safetensors")
        with safe_open(orig_sf11, framework="pt", device="cpu") as sf:
            orig_q_w = sf.get_tensor(q_key)
            orig_q_gate = orig_q_w[6144:, :]
        drafter[q_key] = torch.cat([drafter[q_key], orig_q_gate], dim=0).clone()
        safetensors.torch.save_file(drafter, drafter_file)
        print(f"  Packed Q projection with original output gate -> shape: {list(drafter[q_key].shape)}")

    # 4. Patch model.safetensors.index.json
    base_index_path = os.path.join(base_dir, "model.safetensors.index.json")
    with open(base_index_path, "r") as f:
        index_data = json.load(f)

    drafter_keys = list(drafter.keys())
    print(f"📦 Read {len(drafter_keys)} keys from {os.path.basename(drafter_file)}.")

    weight_map = index_data.get("weight_map", {})
    # Remove any old mtp keys pointing to model-bf16-*.safetensors
    old_mtp_keys = [k for k in weight_map.keys() if k.startswith("mtp.")]
    for k in old_mtp_keys:
        del weight_map[k]

    # Map all drafter keys to model_mtp_tuned.safetensors
    for k in drafter_keys:
        weight_map[k] = "model_mtp_tuned.safetensors"

    shadow_index_path = os.path.join(shadow_dir, "model.safetensors.index.json")
    if os.path.exists(shadow_index_path):
        os.remove(shadow_index_path)
    with open(shadow_index_path, "w") as f:
        json.dump(index_data, f, indent=2)
    print(f"🗺️ Patched model.safetensors.index.json: mapped {len(drafter_keys)} MTP keys.")

    print("\n✨ Calibrated NVFP4 Spliced Model Successfully Created! ✨🚀")
    print(f"Location: {shadow_dir}")
    print("=" * 75)


if __name__ == "__main__":
    splice_nvfp4()