#!/usr/bin/env python3
"""Export an NVFP4 draft (W4A16: quantized routed experts, BF16 the rest)
from the BF16 GLM-5.3-Flash mtp/ folder.

Uses ModelOpt's own NVFP4QTensor.quantize (same math as the target
checkpoint; packing: byte k = (even nibble | odd nibble<<4), one-level
e4m3 block scales, scale_2 = 1.0). Declares input_activations: null so
vLLM routes the experts through its W4A16 NVFP4 path (no input_scale
needed, mirroring the target's expert behavior).

Output: ~/github/GLM5.3-Flash-E224-DGX-Spark/mtp-nvfp4/
"""
import json
import shutil
import struct
import sys
import time
from pathlib import Path

import torch
from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
from safetensors.torch import save_file

torch.set_grad_enabled(False)

# Usage: python3 export_mtp_nvfp4.py /path/to/GLM5.3-Flash-E224-DGX-Spark
ROOT = (Path(sys.argv[1]).expanduser() if len(sys.argv) > 1
        else Path.home() / "github/GLM5.3-Flash-E224-DGX-Spark")
SRC = ROOT / "mtp"
DST = ROOT / "mtp-nvfp4"


def read_shard(p: Path):
    with open(p, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        hdr = json.loads(f.read(n))
        base = f.tell()
        blobs = {}
        for k, v in hdr.items():
            if k == "__metadata__":
                continue
            s, e = v["data_offsets"]
            f.seek(base + s)
            blobs[k] = f.read(e - s)
    return hdr, blobs


def main():
    DST.mkdir(exist_ok=True)
    for f in SRC.iterdir():
        if f.suffix in (".json", ".jinja") and f.name not in (
            "hf_quant_config.json",
            "model.safetensors.index.json",
        ):
            shutil.copy(f, DST / f.name)

    out_tensors = {}
    n_quant = 0
    n_bf16 = 0
    t0 = time.time()

    for si in range(1, 5):
        p = SRC / f"model-{si:05d}-of-00004.safetensors"
        hdr, blobs = read_shard(p)
        for name, blob in blobs.items():
            short = name.replace("model.language_model.layers.45.", "")
            quantize_this = (
                short.startswith("mlp.experts.")
                or short.startswith("mlp.shared_experts.")
            ) and short.endswith(("gate_proj.weight", "up_proj.weight", "down_proj.weight"))
            if quantize_this:
                shape = hdr[name]["shape"]
                w = torch.frombuffer(bytearray(blob), dtype=torch.bfloat16).clone().view(shape)
                qt, block_scale, scale_2 = NVFP4QTensor.quantize(
                    w, 16, weights_scaling_factor_2=torch.tensor(1.0), try_tensorrt=False
                )
                packed = qt.quantized_data if hasattr(qt, "quantized_data") else qt._quantized_data
                base = name[: -len(".weight")]
                out_tensors[base + ".weight"] = packed.contiguous()
                out_tensors[base + ".weight_scale"] = block_scale.contiguous()
                out_tensors[base + ".input_scale"] = torch.ones((), dtype=torch.float32)
                out_tensors[base + ".weight_scale_2"] = torch.ones((), dtype=torch.float32)
                n_quant += 1
                if short == "eh_proj.weight":
                    out_tensors["model.layers.45.eh_proj.input_scale"] = torch.ones((), dtype=torch.float32)
                    out_tensors["model.layers.45.eh_proj.weight_scale_2"] = torch.ones((), dtype=torch.float32)
                if n_quant % 100 == 0:
                    print(f"  {n_quant} expert linears quantized ({time.time()-t0:.0f}s)")
            else:
                shape = hdr[name]["shape"]
                dt = {"BF16": torch.bfloat16, "F32": torch.float32,
                      "F8_E4M3": torch.float8_e4m3fn, "U8": torch.uint8}[hdr[name]["dtype"]]
                w = torch.frombuffer(bytearray(blob), dtype=dt).clone()
                out_tensors[name] = w.view(shape) if w.numel() == torch.Size(shape).numel() else w[: torch.Size(shape).numel()].view(shape)
                n_bf16 += 1

    print(f"quantized: {n_quant} expert linears; kept BF16: {n_bf16} tensors")
    save_file(out_tensors, str(DST / "model.safetensors"))

    # quant config: W4A16 NVFP4 (input_activations: null)
    quant_cfg = {
        "producer": {"name": "modelopt-handoff", "version": "0.47.0-direct"},
        "quantization": {
            "quant_algo": "NVFP4",
            "kv_cache_quant_algo": None,
            "group_size": 16,
            "exclude_modules": [
                # keep attention + indexer + norms + embed + head + eh_proj
                # in BF16; experts and shared_experts ARE quantized (NVFP4)
                "model.layers.45.mtp_block.mlp.gate",
                "model.layers.45.mtp_block.self_attn*",
                "model.layers.45.mtp_block.input_layernorm*",
                "model.layers.45.mtp_block.post_attention_layernorm*",
                "model.layers.45.eh_proj",
                "model.layers.45.enorm",
                "model.layers.45.hnorm",
                "model.layers.45.shared_head*",
                "model.layers.45.embed_tokens*",
            ],
            "config_groups": {
                "group_0": {
                    "weights": {"num_bits": 4, "type": "float", "dynamic": False, "group_size": 16},
                    "input_activations": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": 16},
                    "targets": ["Linear"],
                }
            },
            "quant_method": "modelopt",
        },
    }
    (DST / "hf_quant_config.json").write_text(json.dumps(quant_cfg, indent=2))

    # get_quant_config() reads hf_config.quantization_config (config.json),
    # not hf_quant_config.json — keep both consistent.
    cfg = json.loads((DST / "config.json").read_text())
    cfg["quantization_config"] = quant_cfg["quantization"]
    (DST / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False))

    print("wrote", DST / "model.safetensors",
          f"({sum(t.numel() * t.element_size() for t in out_tensors.values()) / 1e9:.2f} GB)")


if __name__ == "__main__":
    main()