"""Tiny random Mistral4 with fp8 (attention + routed experts) and nvfp4 (shared experts) layers, saved pre-quantized
with a per-layer quantization plan in config.json. Usage: PYTHONPATH=<worktree>/src python make.py <out_dir>"""

import fnmatch
import json
import sys

import torch
from safetensors.torch import save_file

from transformers import FineGrainedFP8Config, Mistral4Config, Mistral4ForCausalLM, NVFP4Config
from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize
from transformers.integrations.mxfp4 import FP4_VALUES as FP4_E2M1_VALUES
from transformers.integrations.nvfp4 import SF_VEC_SIZE, NVFP4Dequantize


out = sys.argv[1]
torch.manual_seed(0)
config = Mistral4Config(
    vocab_size=128, hidden_size=256, intermediate_size=512, moe_intermediate_size=256, num_hidden_layers=2,
    num_attention_heads=2, num_key_value_heads=2, n_routed_experts=4, num_experts_per_tok=2, n_shared_experts=1,
    q_lora_rank=64, kv_lora_rank=64, qk_rope_head_dim=32, qk_nope_head_dim=32, v_head_dim=32, first_k_dense_replace=0,
)  # fmt: skip
model = Mistral4ForCausalLM(config).to(torch.bfloat16).eval()

fp8 = FineGrainedFP8Config(weight_block_size=(32, 32))
# The experts get the real-world 128x128 blocks: the grouped kernel needs them
plan = {"*.self_attn.*_proj*": fp8, "*.mlp.experts.*": FineGrainedFP8Config(weight_block_size=(128, 128)), "*.mlp.shared_experts.*": NVFP4Config()}


def quantize_nvfp4(weight):
    "Packed e2m1 (low nibble first), row-major e4m3 block scales of 16, and the per-tensor global scale"
    weight = weight.float()
    global_scale = 448.0 * 6.0 / weight.abs().amax()
    blocks = weight.view(weight.shape[0], -1, SF_VEC_SIZE)
    scales = (blocks.abs().amax(-1) / 6.0 * global_scale).clamp(max=448.0).to(torch.float8_e4m3fn)
    scaled = blocks * global_scale / scales.float().clamp(min=1e-12).unsqueeze(-1)
    codes = (scaled.unsqueeze(-1) - torch.tensor(FP4_E2M1_VALUES)).abs().argmin(-1).to(torch.uint8).flatten(1)
    packed = codes[:, 0::2] | (codes[:, 1::2] << 4)
    return {"weight": packed, "weight_sf_rowmajor": scales.view(torch.uint8), "weight_global_scale": global_scale.reshape(1)}


state_dict, report = {}, []
for name, tensor in model.state_dict().items():
    pattern = next((p for p in plan if fnmatch.fnmatch(name, p)), None)
    if pattern is None or tensor.ndim < 2:
        state_dict[name] = tensor.contiguous()
        continue
    prefix = name.removesuffix("weight")
    if isinstance(plan[pattern], NVFP4Config):
        quantized = quantize_nvfp4(tensor)
        dense = NVFP4Dequantize(None).convert({"weight": quantized["weight"], "weight_sf_rowmajor": quantized["weight_sf_rowmajor"], "weight_global_scale": quantized["weight_global_scale"]}, model=model, full_layer_name=name)["weight"][0]
        state_dict.update({prefix + key: value for key, value in quantized.items()})
    else:
        hf_quantizer = type("Q", (), {"quantization_config": plan[pattern]})()
        quantized = Fp8Quantize(hf_quantizer).convert({name: tensor})
        scale_key = next(key for key in quantized if key != name)
        dense = Fp8Dequantize(hf_quantizer).convert({"weight": quantized[name], "weight_scale_inv": quantized[scale_key]}, full_layer_name=name)[name]
        state_dict.update(quantized)
    error = ((dense.float() - tensor.float()).norm() / tensor.float().norm()).item()
    report.append(f"{type(plan[pattern]).__name__:22} {name:52} rel err {error:.3f}")
    assert error < 0.2, (name, error)

print("\n".join(report))
model.config.quantization_config = {pattern: cfg.to_dict() for pattern, cfg in plan.items()}
model.config.save_pretrained(out)
save_file({k: v.contiguous() for k, v in state_dict.items()}, f"{out}/model.safetensors", metadata={"format": "pt"})
torch.save({k: v for k, v in model.state_dict().items()}, f"{out}/dense_reference.pt")
print(json.dumps(model.config.quantization_config, indent=1)[:800])