#!/usr/bin/env python3
"""Repair oMLX-quantized Qwen3.5-MoE checkpoints that store routed experts
in the legacy per-expert layout (mlp.experts.<E>.{gate,up,down}_proj.*).

mlx-lm's qwen3_5_moe sanitize() only stacks the fused `experts.gate_up_proj`
layout, so per-expert checkpoints fail to load with
"Received NNNNN parameters not in model". This script repairs the model
in place:

  1. Stacks every `<prefix>.mlp.experts.<E>.<proj>.<tensor>` group along a new
     leading axis into `<prefix>.mlp.switch_mlp.<proj>.<tensor>`.
  2. Rewrites per-path quantization overrides in config.json from raw HF key
     names (model.language_model.*) to post-sanitize module paths
     (language_model.model.*), which is how mlx-lm looks them up at load time.

New shards are written alongside the originals, verified bitwise against the
source tensors, and only then swapped in (originals deleted). A failure at any
point leaves the original model untouched. Needs free disk roughly equal to
the model size while running.

Usage:
    python3 repair_moe_experts.py <model_dir>

Requires mlx; no other dependencies.
"""

import json
import re
import struct
import sys
from pathlib import Path

import mlx.core as mx

EXPERT_RE = re.compile(
    r"^(?P<prefix>.+\.mlp)\.experts\.(?P<e>\d+)\.(?P<proj>\w+_proj)\.(?P<t>weight|scales|biases)$"
)
SHARD_BYTES = 5 * 1024**3
TMP_PREFIX = "tmp-shard-"


def quant_key_to_module_path(key):
    if key.startswith("model.language_model"):
        return key.replace("model.language_model", "language_model.model", 1)
    if key.startswith("language_model."):
        return key
    return "language_model." + key


def read_header(path):
    with open(path, "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        return json.loads(f.read(n))


def nbytes(a):
    return a.size * a.dtype.size


class ShardWriter:
    def __init__(self, out_dir, metadata):
        self.out_dir = out_dir
        self.metadata = metadata
        self.buffer = {}
        self.buffer_bytes = 0
        self.files = []  # [(tmp_path, [keys])]

    def add(self, key, array):
        self.buffer[key] = array
        self.buffer_bytes += nbytes(array)
        if self.buffer_bytes >= SHARD_BYTES:
            self.flush()

    def flush(self):
        if not self.buffer:
            return
        tmp = self.out_dir / f"{TMP_PREFIX}{len(self.files):05d}.safetensors"
        mx.save_safetensors(str(tmp), self.buffer, metadata=self.metadata)
        mx.clear_cache()
        self.files.append((tmp, list(self.buffer)))
        self.buffer = {}
        self.buffer_bytes = 0

    def tmp_weight_map(self):
        return {k: tmp for tmp, keys in self.files for k in keys}

    def commit(self, old_shards):
        """Delete the original shards and move tmp shards to final names."""
        for shard in old_shards:
            (self.out_dir / shard).unlink()
        n = len(self.files)
        weight_map, total = {}, 0
        for i, (tmp, keys) in enumerate(self.files, 1):
            name = f"model-{i:05d}-of-{n:05d}.safetensors"
            tmp.rename(self.out_dir / name)
            hdr = read_header(self.out_dir / name)
            for k, v in hdr.items():
                if k != "__metadata__":
                    total += v["data_offsets"][1] - v["data_offsets"][0]
            for k in keys:
                weight_map[k] = name
        index = {"metadata": {"total_size": total}, "weight_map": weight_map}
        with open(self.out_dir / "model.safetensors.index.json", "w") as f:
            json.dump(index, f, indent=2)
        return weight_map


def cleanup_tmp(model_dir):
    for p in model_dir.glob(f"{TMP_PREFIX}*.safetensors"):
        p.unlink()


def main():
    if len(sys.argv) != 2:
        sys.exit(__doc__)
    model_dir = Path(sys.argv[1]).resolve()
    index_file = model_dir / "model.safetensors.index.json"
    if not index_file.exists():
        sys.exit(f"error: {index_file} not found")
    cleanup_tmp(model_dir)  # leftovers from an interrupted run

    weight_map = json.load(open(index_file))["weight_map"]
    expert_keys = [k for k in weight_map if EXPERT_RE.match(k)]
    if not expert_keys:
        print("no per-expert tensors found — model is already repaired")
        return

    # group per-expert keys by their stacked target
    groups = {}  # stacked_key -> {expert_idx: source_key}
    for k in expert_keys:
        m = EXPERT_RE.match(k)
        stacked = f"{m['prefix']}.switch_mlp.{m['proj']}.{m['t']}"
        groups.setdefault(stacked, {})[int(m["e"])] = k
    n_experts = {len(v) for v in groups.values()}
    if len(n_experts) != 1:
        sys.exit(f"error: inconsistent expert counts per group: {sorted(n_experts)}")
    n_experts = n_experts.pop()
    key_to_stacked = {sk: stacked for stacked, exps in groups.items() for sk in exps.values()}
    print(f"{len(expert_keys)} per-expert tensors -> {len(groups)} stacked tensors ({n_experts} experts)")

    mx.set_default_device(mx.cpu)
    old_shards = sorted({v for v in weight_map.values()})
    src_metadata = read_header(model_dir / old_shards[0]).get("__metadata__") or {"format": "mlx"}
    writer = ShardWriter(model_dir, src_metadata)
    pending = {}  # stacked_key -> {expert_idx: array}

    try:
        for i, shard in enumerate(old_shards, 1):
            print(f"[{i}/{len(old_shards)}] {shard}")
            tensors = mx.load(str(model_dir / shard))
            for key, array in tensors.items():
                stacked = key_to_stacked.get(key)
                if stacked is None:
                    writer.add(key, array)
                    continue
                e = int(EXPERT_RE.match(key)["e"])
                pending.setdefault(stacked, {})[e] = array
                if len(pending[stacked]) == n_experts:
                    parts = pending.pop(stacked)
                    assert sorted(parts) == list(range(n_experts)), f"non-contiguous experts for {stacked}"
                    writer.add(stacked, mx.stack([parts[j] for j in range(n_experts)]))
            del tensors
        if pending:
            raise RuntimeError(f"incomplete expert groups: {list(pending)[:3]}")
        writer.flush()

        # verify before touching the originals
        tmp_map = writer.tmp_weight_map()
        expected = len(weight_map) - len(expert_keys) + len(groups)
        if len(tmp_map) != expected:
            raise RuntimeError(f"tensor count mismatch: {len(tmp_map)} != {expected}")
        check = [(sk, e) for sk in list(groups)[::max(1, len(groups) // 4)] for e in (0, n_experts - 1)]
        by_tmp = {}
        for sk, e in check:
            by_tmp.setdefault(tmp_map[sk], []).append((sk, e))
        for tmp, items in by_tmp.items():
            out_tensors = mx.load(str(tmp))
            for sk, e in items:
                src_key = groups[sk][e]
                orig = mx.load(str(model_dir / weight_map[src_key]))[src_key]
                if not mx.array_equal(out_tensors[sk][e], orig).item():
                    raise RuntimeError(f"bitwise mismatch at {sk}[{e}]")
            del out_tensors
        print(f"spot-check passed ({len(check)} slices)")
    except Exception as e:
        cleanup_tmp(model_dir)
        sys.exit(f"error: {e} — original model left untouched")

    # point of no return: swap repaired shards in, rewrite index and config
    writer.commit(old_shards)
    config_file = model_dir / "config.json"
    config = json.load(open(config_file))
    for section in ("quantization", "quantization_config"):
        if isinstance(config.get(section), dict):
            config[section] = {
                (quant_key_to_module_path(k) if isinstance(v, dict) else k): v
                for k, v in config[section].items()
            }
    with open(config_file, "w") as f:
        json.dump(config, f, indent=2)
    print(f"done: {model_dir} repaired in place")


if __name__ == "__main__":
    main()