#!/usr/bin/env python3
"""Convert Mistral Small 4 HF weights to Mistral native consolidated format.

The converter targets the BF16 path:

    HF Mistral3ForConditionalGeneration-style checkpoint
      -> Mistral native `params.json` + `consolidated*.safetensors`

It copies native runtime assets from an official native reference directory,
strips NVFP4 quantization config from `params.json`, rewrites tensor names, and
splits fused HF MoE expert tensors into native per-expert weights.
"""
from __future__ import annotations

import argparse
import json
import math
import re
import shutil
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

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


TEXT_LAYER_BF16_RE = re.compile(
    r"^model\.language_model\.model\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
)
TEXT_LAYER_NVFP4_RE = re.compile(
    r"^model\.language_model\.model\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
)
VISION_LAYER_BF16_RE = re.compile(
    r"^model\.vision_tower\.transformer\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
)
VISION_LAYER_NVFP4_RE = re.compile(
    r"^model\.vision_tower\.transformer\.layers\.(?P<layer>\d+)\.(?P<rest>.+)$"
)
NVFP4_EXPERT_RE = re.compile(
    r"^mlp\.experts\.(?P<expert>\d+)\."
    r"(?P<proj>gate_proj|up_proj|down_proj)\."
    r"(?P<artifact>weight_packed|weight_scale|weight_global_scale|input_global_scale)$"
)

TOP_LEVEL_MAP = {
    "model.language_model.model.embed_tokens.weight": "tok_embeddings.weight",
    "language_model.model.embed_tokens.weight": "tok_embeddings.weight",
    "model.language_model.model.norm.weight": "norm.weight",
    "language_model.model.norm.weight": "norm.weight",
    "language_model.lm_head.weight": "output.weight",
    "model.vision_tower.ln_pre.weight": "vision_encoder.ln_pre.weight",
    "vision_tower.ln_pre.weight": "vision_encoder.ln_pre.weight",
    "model.vision_tower.patch_conv.weight": "vision_encoder.patch_conv.weight",
    "vision_tower.patch_conv.weight": "vision_encoder.patch_conv.weight",
    "model.multi_modal_projector.linear_1.weight": "vision_language_adapter.w_in.weight",
    "multi_modal_projector.linear_1.weight": "vision_language_adapter.w_in.weight",
    "model.multi_modal_projector.linear_2.weight": "vision_language_adapter.w_out.weight",
    "multi_modal_projector.linear_2.weight": "vision_language_adapter.w_out.weight",
    "model.multi_modal_projector.norm.weight": "pre_mm_projector_norm.weight",
    "multi_modal_projector.norm.weight": "pre_mm_projector_norm.weight",
    "model.multi_modal_projector.patch_merger.merging_layer.weight": (
        "patch_merger.merging_layer.weight"
    ),
    "multi_modal_projector.patch_merger.merging_layer.weight": (
        "patch_merger.merging_layer.weight"
    ),
}

TEXT_LAYER_MAP = {
    "input_layernorm.weight": "attention_norm.weight",
    "post_attention_layernorm.weight": "ffn_norm.weight",
    "self_attn.q_a_proj.weight": "attention.wq_a.weight",
    "self_attn.q_b_proj.weight": "attention.wq_b.weight",
    "self_attn.q_a_layernorm.weight": "attention.q_a_norm.weight",
    "self_attn.kv_a_proj_with_mqa.weight": "attention.wkv_a_with_mqa.weight",
    "self_attn.kv_a_layernorm.weight": "attention.kv_a_norm.weight",
    "self_attn.kv_b_proj.weight": "attention.wkv_b.weight",
    "self_attn.o_proj.weight": "attention.wo.weight",
    "mlp.gate.weight": "gate.weight",
    "mlp.shared_experts.gate_proj.weight": "shared_experts.w1.weight",
    "mlp.shared_experts.up_proj.weight": "shared_experts.w3.weight",
    "mlp.shared_experts.down_proj.weight": "shared_experts.w2.weight",
}

VISION_LAYER_MAP = {
    "attention.q_proj.weight": "attention.wq.weight",
    "attention.k_proj.weight": "attention.wk.weight",
    "attention.v_proj.weight": "attention.wv.weight",
    "attention.o_proj.weight": "attention.wo.weight",
    "attention_norm.weight": "attention_norm.weight",
    "ffn_norm.weight": "ffn_norm.weight",
    "feed_forward.gate_proj.weight": "feed_forward.w1.weight",
    "feed_forward.up_proj.weight": "feed_forward.w3.weight",
    "feed_forward.down_proj.weight": "feed_forward.w2.weight",
}

NVFP4_EXPERT_PROJ_MAP = {
    "gate_proj": "w1",
    "down_proj": "w2",
    "up_proj": "w3",
}

ASSET_FILES = (
    "params.json",
    "tekken.json",
    "tokenizer_config.json",
    "processor_config.json",
    "chat_template.jinja",
)


@dataclass(frozen=True)
class TensorRef:
    name: str
    shard: str


@dataclass
class ConvertStats:
    copied: int = 0
    split: int = 0
    bytes_written: int = 0
    shards_written: int = 0


class ShardWriter:
    def __init__(self, out_dir: Path, max_shard_bytes: int) -> None:
        self.out_dir = out_dir
        self.max_shard_bytes = max_shard_bytes
        self.pending: dict[str, torch.Tensor] = {}
        self.pending_bytes = 0
        self.weight_map: dict[str, str] = {}
        self.shard_paths: list[Path] = []
        self.stats = ConvertStats()

    def add(self, name: str, tensor: torch.Tensor) -> None:
        owned = tensor.detach().cpu().contiguous()
        size = tensor_nbytes(owned)
        if self.pending and self.pending_bytes + size > self.max_shard_bytes:
            self.flush()
        self.pending[name] = owned
        self.pending_bytes += size

    def flush(self) -> None:
        if not self.pending:
            return
        shard_idx = len(self.shard_paths) + 1
        shard_name = f"consolidated-{shard_idx:05d}.safetensors"
        shard_path = self.out_dir / shard_name
        save_file(self.pending, shard_path)
        for name in self.pending:
            self.weight_map[name] = shard_name
        self.stats.bytes_written += self.pending_bytes
        self.stats.shards_written += 1
        self.shard_paths.append(shard_path)
        self.pending = {}
        self.pending_bytes = 0

    def write_index(self) -> None:
        self.flush()
        total_size = sum(path.stat().st_size for path in self.shard_paths)
        index = {
            "metadata": {"total_size": total_size},
            "weight_map": dict(sorted(self.weight_map.items())),
        }
        (self.out_dir / "consolidated.safetensors.index.json").write_text(
            json.dumps(index, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )


def tensor_nbytes(tensor: torch.Tensor) -> int:
    return tensor.numel() * tensor.element_size()


def parse_size_gib(value: str) -> int:
    size = float(value)
    if not math.isfinite(size) or size <= 0:
        raise argparse.ArgumentTypeError("--max-shard-size-gb must be positive")
    return int(size * 1024**3)


def load_weight_map(hf_dir: Path) -> list[TensorRef]:
    index_path = hf_dir / "model.safetensors.index.json"
    data = json.loads(index_path.read_text(encoding="utf-8"))
    weight_map = data.get("weight_map")
    if not isinstance(weight_map, dict):
        raise ValueError(f"{index_path} does not contain a weight_map object")
    return [TensorRef(name=name, shard=shard) for name, shard in sorted(weight_map.items())]


_SHARD_CACHE: dict = {}


def read_tensor(hf_dir: Path, ref: TensorRef) -> torch.Tensor:
    # Non-mmap shard read: safetensors safe_open() mmaps the whole shard, which
    # ENOMEMs on /tank (ZFS) for large shards (the 50 GB NVFP4 shard) regardless of
    # free RAM or overcommit (a MAP_SHARED file mmap never consults the commit limit).
    # Read the shard with a plain read() and deserialize from the in-memory buffer
    # instead. Caches one shard at a time, so the copy loop must iterate refs grouped
    # by shard (see the sorted(...) in convert()).
    global _SHARD_CACHE
    if _SHARD_CACHE.get("shard") != ref.shard:
        from safetensors.torch import load as _st_load
        with open(hf_dir / ref.shard, "rb") as _fh:
            _SHARD_CACHE = {"shard": ref.shard, "tensors": _st_load(_fh.read())}
    return _SHARD_CACHE["tensors"][ref.name]


def write_assets(native_ref_dir: Path, out_dir: Path, *, keep_quantization_config: bool) -> None:
    for filename in ASSET_FILES:
        src = native_ref_dir / filename
        if not src.exists():
            continue
        if filename == "params.json":
            params = json.loads(src.read_text(encoding="utf-8"))
            if not keep_quantization_config:
                params.pop("quantization_config", None)
            (out_dir / filename).write_text(
                json.dumps(params, indent=2, sort_keys=True) + "\n",
                encoding="utf-8",
            )
        else:
            shutil.copy2(src, out_dir / filename)


def mapped_name(name: str, *, output_format: str) -> str | None:
    if name in TOP_LEVEL_MAP:
        return TOP_LEVEL_MAP[name]

    text_match = (
        TEXT_LAYER_NVFP4_RE.match(name)
        if output_format == "nvfp4"
        else TEXT_LAYER_BF16_RE.match(name)
    )
    if text_match:
        layer = text_match.group("layer")
        rest = text_match.group("rest")
        suffix = TEXT_LAYER_MAP.get(rest)
        if suffix is not None:
            return f"layers.{layer}.{suffix}"

        if output_format == "nvfp4":
            expert_match = NVFP4_EXPERT_RE.match(rest)
            if expert_match:
                expert = expert_match.group("expert")
                native_proj = NVFP4_EXPERT_PROJ_MAP[expert_match.group("proj")]
                artifact = expert_match.group("artifact")
                return f"layers.{layer}.experts.{expert}.{native_proj}.{artifact}"

            for hf_proj, native_proj in NVFP4_EXPERT_PROJ_MAP.items():
                prefix = f"mlp.shared_experts.{hf_proj}."
                if rest.startswith(prefix):
                    artifact = rest.removeprefix(prefix)
                    if artifact in {
                        "weight_packed",
                        "weight_scale",
                        "weight_global_scale",
                        "input_global_scale",
                    }:
                        return f"layers.{layer}.shared_experts.{native_proj}.{artifact}"

        return None

    vision_match = (
        VISION_LAYER_NVFP4_RE.match(name)
        if output_format == "nvfp4"
        else VISION_LAYER_BF16_RE.match(name)
    )
    if vision_match:
        layer = vision_match.group("layer")
        rest = vision_match.group("rest")
        suffix = VISION_LAYER_MAP.get(rest)
        if suffix is not None:
            return f"vision_encoder.transformer.layers.{layer}.{suffix}"
        return None

    return None


def split_experts(
    *,
    writer: ShardWriter,
    layer: int,
    gate_up: torch.Tensor,
    down: torch.Tensor,
    expert_hidden_dim: int,
) -> None:
    if gate_up.ndim != 3:
        raise ValueError(f"layer {layer}: gate_up_proj must be rank 3, got {gate_up.shape}")
    if down.ndim != 3:
        raise ValueError(f"layer {layer}: down_proj must be rank 3, got {down.shape}")
    if gate_up.shape[1] != expert_hidden_dim * 2:
        raise ValueError(
            f"layer {layer}: gate_up second dim {gate_up.shape[1]} != "
            f"2 * expert_hidden_dim {expert_hidden_dim}"
        )
    if gate_up.shape[0] != down.shape[0]:
        raise ValueError(
            f"layer {layer}: gate_up experts {gate_up.shape[0]} != down experts {down.shape[0]}"
        )

    for expert in range(gate_up.shape[0]):
        writer.add(
            f"layers.{layer}.experts.{expert}.w1.weight",
            gate_up[expert, :expert_hidden_dim, :].clone(),
        )
        writer.add(
            f"layers.{layer}.experts.{expert}.w3.weight",
            gate_up[expert, expert_hidden_dim:, :].clone(),
        )
        writer.add(f"layers.{layer}.experts.{expert}.w2.weight", down[expert].clone())


def text_layer_id(name: str, suffix: str) -> int | None:
    match = TEXT_LAYER_BF16_RE.match(name)
    if match and match.group("rest") == suffix:
        return int(match.group("layer"))
    return None


def convert(
    *,
    hf_dir: Path,
    native_ref_dir: Path,
    out_dir: Path,
    max_shard_bytes: int,
    expert_hidden_dim: int,
    output_format: str,
    dry_run: bool,
) -> dict[str, Any]:
    refs = load_weight_map(hf_dir)
    if output_format == "bf16":
        gate_up_by_layer = {
            layer: ref
            for ref in refs
            if (layer := text_layer_id(ref.name, "mlp.experts.gate_up_proj")) is not None
        }
        down_by_layer = {
            layer: ref
            for ref in refs
            if (layer := text_layer_id(ref.name, "mlp.experts.down_proj")) is not None
        }
    else:
        gate_up_by_layer = {}
        down_by_layer = {}

    unmapped: list[str] = []
    mapped: dict[str, str] = {}
    split_layers = sorted(set(gate_up_by_layer) | set(down_by_layer))

    for ref in refs:
        if output_format == "bf16" and (
            text_layer_id(ref.name, "mlp.experts.gate_up_proj") is not None
            or text_layer_id(ref.name, "mlp.experts.down_proj") is not None
        ):
            continue
        native_name = mapped_name(ref.name, output_format=output_format)
        if native_name is None:
            unmapped.append(ref.name)
        else:
            mapped[ref.name] = native_name

    if dry_run:
        return {
            "mapped_tensors": len(mapped),
            "split_layers": split_layers,
            "unmapped": unmapped,
            "native_reference_diff": reference_name_diff(
                native_ref_dir=native_ref_dir,
                candidate_names=predicted_native_names(mapped.values(), split_layers),
                output_format=output_format,
            ),
        }

    out_dir.mkdir(parents=True, exist_ok=True)
    if any(out_dir.iterdir()):
        raise FileExistsError(f"output directory is not empty: {out_dir}")

    write_assets(
        native_ref_dir,
        out_dir,
        keep_quantization_config=output_format == "nvfp4",
    )
    writer = ShardWriter(out_dir=out_dir, max_shard_bytes=max_shard_bytes)

    for ref in sorted(refs, key=lambda r: (r.shard, r.name)):
        if ref.name not in mapped:
            continue
        writer.add(mapped[ref.name], read_tensor(hf_dir, ref))
        writer.stats.copied += 1

    for layer in split_layers:
        gate_ref = gate_up_by_layer.get(layer)
        down_ref = down_by_layer.get(layer)
        if gate_ref is None or down_ref is None:
            raise ValueError(f"layer {layer}: missing gate_up or down fused expert tensor")
        split_experts(
            writer=writer,
            layer=layer,
            gate_up=read_tensor(hf_dir, gate_ref),
            down=read_tensor(hf_dir, down_ref),
            expert_hidden_dim=expert_hidden_dim,
        )
        writer.stats.split += 1

    writer.write_index()

    report = {
        "copied_tensors": writer.stats.copied,
        "split_layers": writer.stats.split,
        "shards_written": writer.stats.shards_written,
        "bytes_written": writer.stats.bytes_written,
        "unmapped": unmapped,
        "native_reference_diff": reference_name_diff(
            native_ref_dir=native_ref_dir,
            candidate_names=writer.weight_map.keys(),
            output_format=output_format,
        ),
    }
    (out_dir / "conversion_report.json").write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return report


def predicted_native_names(mapped_names: Iterable[str], split_layers: Iterable[int]) -> set[str]:
    names = set(mapped_names)
    for layer in split_layers:
        for expert in range(128):
            for weight in ("w1", "w2", "w3"):
                names.add(f"layers.{layer}.experts.{expert}.{weight}.weight")
    return names


def iter_native_reference_names(native_ref_dir: Path, *, output_format: str) -> Iterable[str]:
    index_path = native_ref_dir / "consolidated.safetensors.index.json"
    if not index_path.exists():
        return ()
    data = json.loads(index_path.read_text(encoding="utf-8"))
    names = data.get("weight_map", {}).keys()
    if output_format == "nvfp4":
        return set(names)

    normalized = set()
    for name in names:
        if name.endswith(".weight_packed"):
            normalized.add(name.removesuffix("weight_packed") + "weight")
        elif not (
            name.endswith(".weight_scale")
            or name.endswith(".weight_global_scale")
            or name.endswith(".input_global_scale")
        ):
            normalized.add(name)
    return normalized


def reference_name_diff(
    *, native_ref_dir: Path, candidate_names: Iterable[str], output_format: str
) -> dict[str, list[str]]:
    reference_names = set(
        iter_native_reference_names(native_ref_dir, output_format=output_format)
    )
    if not reference_names:
        return {"missing_from_output": [], "extra_in_output": []}

    candidates = set(candidate_names)
    return {
        "missing_from_output": sorted(reference_names - candidates),
        "extra_in_output": sorted(candidates - reference_names),
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--hf-dir", type=Path, required=True)
    parser.add_argument("--native-ref-dir", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument(
        "--max-shard-size-gb",
        type=parse_size_gib,
        default=parse_size_gib("20"),
        help="Approximate max safetensors shard size in GiB before flushing.",
    )
    parser.add_argument(
        "--expert-hidden-dim",
        type=int,
        default=2048,
        help="Mistral Small 4 routed expert hidden dim used to split gate_up_proj.",
    )
    parser.add_argument(
        "--format",
        choices=("bf16", "nvfp4"),
        default="bf16",
        help="Output checkpoint format. BF16 splits fused HF experts; NVFP4 renames per-expert quant artifacts.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only report mapped/unmapped tensors; do not read or write weight shards.",
    )
    return parser


def main() -> int:
    args = build_parser().parse_args()
    report = convert(
        hf_dir=args.hf_dir,
        native_ref_dir=args.native_ref_dir,
        out_dir=args.out_dir,
        max_shard_bytes=args.max_shard_size_gb,
        expert_hidden_dim=args.expert_hidden_dim,
        output_format=args.format,
        dry_run=args.dry_run,
    )
    print(json.dumps(report, indent=2, sort_keys=True))
    return 1 if report.get("unmapped") else 0


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