#!/usr/bin/env python3
"""Stream a Qwen3.6 ModelOpt FP8/NVFP4 checkpoint into MLX format.

This is intentionally a format conversion, not a re-quantization:

* FP8 E4M3 weight bytes are packed four-at-a-time into MLX uint32 tensors.
  Unit E8M0 block scales make MLX's MXFP8 kernel decode the original FP8
  values exactly; the original ModelOpt tensor scale is retained separately.
* NVFP4 E2M1 weight nibbles and E4M3 block-scale bytes are repacked without
  numerical modification.  The original FP32 tensor scales are retained.
* Expert tensors are stacked into MLX's SwitchGLU layout one layer at a time.
* The original BF16 vision tower is preserved unchanged in a dedicated shard
  and loaded by the model-local MLX-VLM runtime.
* ModelOpt activation scales are recorded as dropped because this runtime uses
  weight-only quantized kernels and keeps activations in the model dtype.

Each transformer layer is written as its own safetensors shard, bounding peak
memory to roughly one layer rather than the whole model.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import resource
import shutil
import sys
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, Iterable, Mapping

import mlx.core as mx


RUNTIME_FILE = "modeling_mlx_qwen36_modelopt_hybrid.py"
VLM_RUNTIME_FILE = "modeling_mlx_vlm_qwen36_modelopt_hybrid.py"
EXPERT_RE = re.compile(
    r"^model\.language_model\.layers\.(\d+)\.mlp\.experts\.(\d+)\."
    r"(gate_proj|up_proj|down_proj)\."
    r"(weight|weight_scale|weight_scale_2|input_scale)$"
)
LAYER_RE = re.compile(r"^model\.language_model\.layers\.(\d+)\.")
NORM_SUFFIXES = (
    ".input_layernorm.weight",
    ".post_attention_layernorm.weight",
    "model.norm.weight",
    ".q_norm.weight",
    ".k_norm.weight",
)


def log(message: str) -> None:
    now = datetime.now().strftime("%H:%M:%S")
    rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024**3)
    print(f"[{now}] [peak RSS {rss:.2f} GiB] {message}", flush=True)


def pack_u8x4(x: mx.array) -> mx.array:
    """Pack four consecutive bytes into one little-endian uint32."""
    if x.dtype != mx.uint8:
        raise TypeError(f"Expected uint8 storage, got {x.dtype}")
    if x.shape[-1] % 4:
        raise ValueError(f"Last dimension {x.shape[-1]} is not divisible by four")
    y = x.reshape(*x.shape[:-1], x.shape[-1] // 4, 4).astype(mx.uint32)
    return y[..., 0] | (y[..., 1] << 8) | (y[..., 2] << 16) | (y[..., 3] << 24)


def sanitize_name(key: str) -> str:
    prefix = "model.language_model."
    if key.startswith(prefix):
        return "language_model.model." + key[len(prefix) :]
    if key.startswith("language_model."):
        return key
    return "language_model." + key


def should_drop(key: str) -> bool:
    return (
        key.startswith("model.visual")
        or key.startswith("vision_tower")
        or key.startswith("mtp.")
        or ".mtp." in key
    )


class SourceWeights:
    def __init__(self, root: Path, weight_map: Mapping[str, str]):
        self.root = root
        self.weight_map = weight_map
        self._shards: Dict[str, Dict[str, mx.array]] = {}

    def get(self, key: str) -> mx.array:
        shard_name = self.weight_map[key]
        if shard_name not in self._shards:
            log(f"Opening source shard {shard_name} lazily")
            self._shards[shard_name] = mx.load(str(self.root / shard_name))
        return self._shards[shard_name][key]


def group_id_for_key(key: str) -> int | None:
    match = LAYER_RE.match(key)
    return int(match.group(1)) if match else None


def group_id_for_prefix(prefix: str) -> int | None:
    return group_id_for_key(prefix + ".weight")


def add_output(
    output: Dict[str, mx.array],
    key: str,
    value: mx.array,
) -> None:
    if key in output:
        raise KeyError(f"Duplicate output tensor {key}")
    output[key] = value


def convert_dense_nvfp4(
    source: SourceWeights,
    prefix: str,
    output: Dict[str, mx.array],
    quantization: Dict[str, str],
) -> None:
    raw = source.get(prefix + ".weight")
    scales = source.get(prefix + ".weight_scale")
    tensor_scale = source.get(prefix + ".weight_scale_2")
    if raw.dtype != mx.uint8 or scales.dtype != mx.uint8:
        raise TypeError(
            f"Unexpected NVFP4 storage for {prefix}: {raw.dtype}, {scales.dtype}"
        )
    if raw.shape[-1] != scales.shape[-1] * 8:
        raise ValueError(
            f"NVFP4 shape mismatch for {prefix}: weight={raw.shape}, scales={scales.shape}"
        )
    out_prefix = sanitize_name(prefix)
    add_output(output, out_prefix + ".weight", pack_u8x4(raw))
    add_output(output, out_prefix + ".scales", scales)
    add_output(output, out_prefix + ".global_scale", tensor_scale)
    quantization[out_prefix] = "scaled_nvfp4"


def convert_dense_fp8(
    source: SourceWeights,
    prefix: str,
    output: Dict[str, mx.array],
    quantization: Dict[str, str],
) -> None:
    raw = source.get(prefix + ".weight")
    tensor_scale = source.get(prefix + ".weight_scale")
    if raw.dtype != mx.uint8:
        raise TypeError(f"Unexpected FP8 storage for {prefix}: {raw.dtype}")
    if raw.shape[-1] % 32:
        raise ValueError(
            f"FP8 input dimension for {prefix} is not divisible by 32: {raw.shape}"
        )
    out_prefix = sanitize_name(prefix)
    add_output(output, out_prefix + ".weight", pack_u8x4(raw))
    # E8M0 byte 0x7f represents exactly 1.0.  With unit scales, MLX's
    # MXFP8 decoder reproduces ModelOpt's E4M3 weight bytes losslessly.
    scale_shape = (*raw.shape[:-1], raw.shape[-1] // 32)
    add_output(
        output,
        out_prefix + ".scales",
        mx.full(scale_shape, 0x7F, dtype=mx.uint8),
    )
    add_output(output, out_prefix + ".global_scale", tensor_scale)
    quantization[out_prefix] = "scaled_mxfp8"


def convert_expert_projection(
    source: SourceWeights,
    layer: int,
    projection: str,
    expert_lookup: Mapping[tuple[int, str, str], Mapping[int, str]],
    num_experts: int,
    output: Dict[str, mx.array],
    quantization: Dict[str, str],
) -> None:
    expected = list(range(num_experts))
    suffix_maps = {
        suffix: expert_lookup[(layer, projection, suffix)]
        for suffix in ("weight", "weight_scale", "weight_scale_2")
    }
    for suffix, mapping in suffix_maps.items():
        present = sorted(mapping)
        if present != expected:
            missing = sorted(set(expected) - set(present))
            raise ValueError(
                f"Layer {layer} {projection} {suffix}: expected {num_experts} "
                f"experts, missing {missing[:20]}"
            )

    raw = mx.stack(
        [source.get(suffix_maps["weight"][expert]) for expert in expected], axis=0
    )
    scales = mx.stack(
        [source.get(suffix_maps["weight_scale"][expert]) for expert in expected],
        axis=0,
    )
    tensor_scales = mx.stack(
        [source.get(suffix_maps["weight_scale_2"][expert]) for expert in expected],
        axis=0,
    )
    if raw.dtype != mx.uint8 or scales.dtype != mx.uint8:
        raise TypeError(
            f"Unexpected expert NVFP4 storage at layer {layer} {projection}: "
            f"{raw.dtype}, {scales.dtype}"
        )
    if raw.shape[-1] != scales.shape[-1] * 8:
        raise ValueError(
            f"Expert NVFP4 shape mismatch at layer {layer} {projection}: "
            f"weight={raw.shape}, scales={scales.shape}"
        )

    out_prefix = (
        f"language_model.model.layers.{layer}.mlp.switch_mlp.{projection}"
    )
    add_output(output, out_prefix + ".weight", pack_u8x4(raw))
    add_output(output, out_prefix + ".scales", scales)
    add_output(output, out_prefix + ".global_scales", tensor_scales)
    quantization[out_prefix] = "scaled_nvfp4_switch"


def transform_standard_weight(
    key: str,
    value: mx.array,
    shift_norm_weights: bool,
) -> tuple[str, mx.array]:
    out_key = sanitize_name(key)
    if "conv1d.weight" in out_key and value.shape[-1] != 1:
        value = value.moveaxis(2, 1)
    if (
        shift_norm_weights
        and value.ndim == 1
        and any(out_key.endswith(suffix) for suffix in NORM_SUFFIXES)
    ):
        value = value + 1.0
    return out_key, value


def write_shard(
    partial: Path,
    filename: str,
    tensors: Dict[str, mx.array],
    output_weight_map: Dict[str, str],
) -> int:
    tensors = dict(sorted(tensors.items()))
    total_bytes = sum(array.nbytes for array in tensors.values())
    log(
        f"Writing {filename}: {len(tensors)} tensors, "
        f"{total_bytes / (1024**3):.2f} GiB"
    )
    mx.save_safetensors(
        str(partial / filename),
        tensors,
        metadata={"format": "mlx"},
    )
    for key in tensors:
        if key in output_weight_map:
            raise KeyError(f"Tensor {key} was already assigned to a shard")
        output_weight_map[key] = filename
    del tensors
    try:
        mx.clear_cache()
    except AttributeError:
        # Compatibility with older MLX releases.
        mx.metal.clear_cache()
    return total_bytes


def copy_metadata(source: Path, partial: Path) -> None:
    names = [
        "README.md",
        "LICENSE",
        "chat_template.jinja",
        "generation_config.json",
        "preprocessor_config.json",
        "video_preprocessor_config.json",
        "special_tokens_map.json",
        "tokenizer.json",
        "tokenizer.model",
        "tokenizer_config.json",
        "vocab.json",
        "merges.txt",
    ]
    for name in names:
        src = source / name
        if src.exists():
            shutil.copy2(src, partial / name, follow_symlinks=True)
    assets = source / "assets"
    if assets.exists():
        shutil.copytree(assets, partial / "assets")


def convert(
    source: Path,
    output: Path,
    runtime_source: Path,
    vlm_runtime_source: Path,
) -> None:
    source = source.resolve()
    output = output.resolve()
    partial = output.with_name(output.name + ".partial")
    if output.exists():
        raise FileExistsError(f"Output already exists: {output}")
    if partial.exists():
        raise FileExistsError(
            f"Partial output already exists: {partial}. Remove it or choose another path."
        )
    if not (source / "config.json").exists():
        raise FileNotFoundError(f"Missing config.json in {source}")
    if not (source / "model.safetensors.index.json").exists():
        raise FileNotFoundError(f"Missing model.safetensors.index.json in {source}")

    config = json.loads((source / "config.json").read_text())
    index = json.loads((source / "model.safetensors.index.json").read_text())
    weight_map: Dict[str, str] = index["weight_map"]
    all_keys = sorted(weight_map)
    vision_keys = [key for key in all_keys if key.startswith("model.visual.")]
    text_config = config.get("text_config", config)
    num_layers = int(text_config["num_hidden_layers"])
    num_experts = int(text_config["num_experts"])
    log(
        f"Source {source}: {len(all_keys)} tensors, {num_layers} layers, "
        f"{num_experts} experts"
    )

    nvfp4_prefixes = {
        key[: -len(".weight_scale_2")]
        for key in all_keys
        if key.endswith(".weight_scale_2") and not should_drop(key)
    }
    all_scale_prefixes = {
        key[: -len(".weight_scale")]
        for key in all_keys
        if key.endswith(".weight_scale") and not should_drop(key)
    }
    fp8_prefixes = all_scale_prefixes - nvfp4_prefixes

    expert_lookup: dict[tuple[int, str, str], dict[int, str]] = defaultdict(dict)
    expert_keys = set()
    for key in all_keys:
        match = EXPERT_RE.match(key)
        if not match:
            continue
        layer, expert, projection, suffix = match.groups()
        expert_lookup[(int(layer), projection, suffix)][int(expert)] = key
        expert_keys.add(key)

    expert_prefix_marker = ".mlp.experts."
    dense_nvfp4 = sorted(
        prefix for prefix in nvfp4_prefixes if expert_prefix_marker not in prefix
    )
    dense_fp8 = sorted(
        prefix for prefix in fp8_prefixes if expert_prefix_marker not in prefix
    )
    log(
        f"Detected {len(dense_fp8)} dense FP8 modules, "
        f"{len(dense_nvfp4)} dense NVFP4 modules, and "
        f"{len(expert_keys)} expert component tensors"
    )

    layer_standard_keys: dict[int, list[str]] = defaultdict(list)
    global_standard_keys: list[str] = []
    quantized_prefixes = nvfp4_prefixes | fp8_prefixes
    quant_metadata_suffixes = (
        ".input_scale",
        ".weight_scale",
        ".weight_scale_2",
    )
    for key in all_keys:
        if should_drop(key) or key in expert_keys:
            continue
        if key.endswith(quant_metadata_suffixes):
            continue
        if key.endswith(".weight") and key[: -len(".weight")] in quantized_prefixes:
            continue
        group = group_id_for_key(key)
        if group is None:
            global_standard_keys.append(key)
        else:
            layer_standard_keys[group].append(key)

    has_mtp = any(key.startswith("mtp.") or ".mtp." in key for key in all_keys)
    has_unsanitized_conv = False
    source_weights = SourceWeights(source, weight_map)
    for key in all_keys:
        if "conv1d.weight" in key and not should_drop(key):
            if source_weights.get(key).shape[-1] != 1:
                has_unsanitized_conv = True
                break
    shift_norm_weights = has_mtp or has_unsanitized_conv
    log(
        f"Qwen sanitizer flags: has_mtp={has_mtp}, "
        f"unsanitized_conv1d={has_unsanitized_conv}, "
        f"shift_norm_weights={shift_norm_weights}"
    )

    partial.mkdir(parents=True)
    copy_metadata(source, partial)
    shutil.copy2(runtime_source, partial / RUNTIME_FILE)
    shutil.copy2(vlm_runtime_source, partial / VLM_RUNTIME_FILE)

    quantization: Dict[str, str] = {}
    output_weight_map: Dict[str, str] = {}
    total_size = 0
    shard_count = num_layers + 1 + bool(vision_keys)

    # Global tensors: embeddings, final norm, and LM head.
    global_output: Dict[str, mx.array] = {}
    for prefix in dense_fp8:
        if group_id_for_prefix(prefix) is None:
            convert_dense_fp8(source_weights, prefix, global_output, quantization)
    for prefix in dense_nvfp4:
        if group_id_for_prefix(prefix) is None:
            convert_dense_nvfp4(source_weights, prefix, global_output, quantization)
    for key in global_standard_keys:
        out_key, value = transform_standard_weight(
            key, source_weights.get(key), shift_norm_weights
        )
        add_output(global_output, out_key, value)
    total_size += write_shard(
        partial,
        f"model-{1:05d}-of-{shard_count:05d}.safetensors",
        global_output,
        output_weight_map,
    )

    for layer in range(num_layers):
        layer_output: Dict[str, mx.array] = {}
        log(f"Converting transformer layer {layer + 1}/{num_layers}")
        for prefix in dense_fp8:
            if group_id_for_prefix(prefix) == layer:
                convert_dense_fp8(source_weights, prefix, layer_output, quantization)
        for prefix in dense_nvfp4:
            if group_id_for_prefix(prefix) == layer:
                convert_dense_nvfp4(source_weights, prefix, layer_output, quantization)
        for projection in ("gate_proj", "up_proj", "down_proj"):
            convert_expert_projection(
                source_weights,
                layer,
                projection,
                expert_lookup,
                num_experts,
                layer_output,
                quantization,
            )
        for key in layer_standard_keys[layer]:
            out_key, value = transform_standard_weight(
                key, source_weights.get(key), shift_norm_weights
            )
            add_output(layer_output, out_key, value)

        total_size += write_shard(
            partial,
            f"model-{layer + 2:05d}-of-{shard_count:05d}.safetensors",
            layer_output,
            output_weight_map,
        )
        log(f"Finished transformer layer {layer + 1}/{num_layers}")

    # Vision is unchanged by the AntiLoop adapter. Keep the original BF16
    # tensors under their source names; the MLX-VLM runtime maps and sanitizes
    # them. Arrays returned by mx.load remain lazy until the safetensors writer
    # consumes them, so this does not materialize the source shard twice.
    if vision_keys:
        vision_output = {key: source_weights.get(key) for key in vision_keys}
        total_size += write_shard(
            partial,
            f"model-{shard_count:05d}-of-{shard_count:05d}.safetensors",
            vision_output,
            output_weight_map,
        )
        log(f"Preserved {len(vision_keys)} vision tensors without requantization")

    output_index = {
        "metadata": {"total_size": total_size},
        "weight_map": dict(sorted(output_weight_map.items())),
    }
    (partial / "model.safetensors.index.json").write_text(
        json.dumps(output_index, indent=2, sort_keys=True) + "\n"
    )

    original_quantization = config.pop("quantization_config", None)
    config.pop("quantization", None)
    if isinstance(config.get("text_config"), dict):
        config["text_config"].pop("quantization_config", None)
        config["text_config"].pop("quantization", None)
    config["model_file"] = RUNTIME_FILE
    config["vlm_model_file"] = VLM_RUNTIME_FILE
    config["mlx_modelopt_quantization"] = dict(sorted(quantization.items()))
    config["mlx_hybrid_format"] = {
        "format": "modelopt_fp8_nvfp4_v1",
        "fp8_storage": "mxfp8_carrier_with_unit_e8m0_scales",
        "nvfp4_storage": "native_e2m1_e4m3_with_output_tensor_scale",
        "activations": "model_dtype_weight_only_quantized_matmul",
        "source_activation_scales_retained": False,
        "source_quantization_config": original_quantization,
    }
    (partial / "config.json").write_text(
        json.dumps(config, indent=2, sort_keys=True) + "\n"
    )

    input_scales_dropped = sum(key.endswith(".input_scale") for key in all_keys)
    manifest = {
        "source": str(source),
        "output": str(output),
        "created_at": datetime.now(timezone.utc).isoformat(),
        "converter": str(Path(__file__).resolve()),
        "runtime_file": RUNTIME_FILE,
        "vlm_runtime_file": VLM_RUNTIME_FILE,
        "num_layers": num_layers,
        "num_experts": num_experts,
        "source_tensor_count": len(all_keys),
        "output_tensor_count": len(output_weight_map),
        "output_shards": shard_count,
        "output_total_size_bytes": total_size,
        "quantized_module_counts": dict(Counter(quantization.values())),
        "input_scales_dropped": input_scales_dropped,
        "weight_bytes_requantized": False,
        "vision_tensor_count": len(vision_keys),
        "vision_weight_bytes_requantized": False,
        "norm_weights_shifted": shift_norm_weights,
    }
    (partial / "mlx_conversion_manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n"
    )

    partial.rename(output)
    log(
        f"Conversion complete: {output} ({total_size / (1024**3):.2f} GiB, "
        f"{len(output_weight_map)} tensors)"
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--runtime-source",
        type=Path,
        default=Path(__file__).with_name(RUNTIME_FILE),
    )
    parser.add_argument(
        "--vlm-runtime-source",
        type=Path,
        default=Path(__file__).with_name(VLM_RUNTIME_FILE),
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    convert(
        args.source,
        args.output,
        args.runtime_source.resolve(),
        args.vlm_runtime_source.resolve(),
    )


if __name__ == "__main__":
    main()