#!/usr/bin/env python3
"""
One-time conversion: Xiaomi MiMo-V2.5 FP8 shards -> single pre-stacked
safetensors file for MLX block_fp8 path.

Strategy:
  - Read all shard headers up front.
  - For each output key:
      * Non-expert keys: copy raw bytes from source shard, preserve dtype.
      * Expert keys (MoE): read each expert's raw bytes, concatenate, write
        as one larger tensor with prepended expert-axis.
  - Write a single output safetensors file constructed from raw bytes.
  - Peak memory: one stack's worth (~2 GB) + open file handles.

Output:
  /Volumes/TB5/llm/MiMo-V2.5/mimo_v2.5_block_fp8.safetensors
"""

import json
import struct
import glob
import os
import time
import sys
import argparse
from collections import defaultdict

# Set these to wherever Xiaomi's MiMo-V2.5 shards live, and where you want
# the converted single-file safetensors written. Override via CLI:
#   python3 convert_mimo.py --src /path/to/MiMo-V2.5 --out /path/to/output.safetensors
SRC_DIR = os.environ.get("MIMO_SRC", os.path.expanduser("~/llm/MiMo-V2.5"))
OUT_PATH = os.environ.get("MIMO_OUT", os.path.expanduser("~/llm/MiMo-V2.5/mimo_v2.5_block_fp8.safetensors"))

SKIP_PREFIXES = (
    "model.mtp.",
    "visual.",
    "audio_encoder.",
    "speech_embeddings.",
)

# safetensors dtype name -> bytes per element
DTYPE_SIZE = {
    "F8_E4M3": 1, "F8_E5M2": 1,
    "BF16": 2, "F16": 2,
    "F32": 4, "F64": 8,
    "U8": 1, "I8": 1, "U16": 2, "I16": 2,
    "U32": 4, "I32": 4, "U64": 8, "I64": 8,
    "BOOL": 1,
}


def read_shard_headers(shard_paths):
    print(f"[scan] reading {len(shard_paths)} shard headers...", flush=True)
    # key -> (shard_path, dtype_str, shape_list, data_offset_in_shard)
    key_info = {}
    for sp in shard_paths:
        with open(sp, "rb") as f:
            hlen = struct.unpack("<Q", f.read(8))[0]
            hdr = json.loads(f.read(hlen))
            base = 8 + hlen  # absolute byte offset where data section begins
        for k, meta in hdr.items():
            if k == "__metadata__":
                continue
            offs = meta["data_offsets"]
            key_info[k] = (sp, meta["dtype"], meta["shape"], base + offs[0], base + offs[1])
    print(f"[scan] {len(key_info)} unique keys", flush=True)
    return key_info


def filter_keys(key_info):
    before = len(key_info)
    key_info = {k: v for k, v in key_info.items()
                if not k.startswith(SKIP_PREFIXES)}
    print(f"[filter] dropped {before - len(key_info)} mtp/vision/audio keys",
          flush=True)
    return key_info


def detect_n_experts(key_info):
    max_eid = -1
    for k in key_info:
        if ".experts." in k:
            try:
                eid = int(k.split(".experts.")[1].split(".")[0])
                if eid > max_eid:
                    max_eid = eid
            except (ValueError, IndexError):
                pass
    return max_eid + 1


def detect_n_layers(key_info):
    max_lid = -1
    for k in key_info:
        if k.startswith("model.layers."):
            try:
                lid = int(k.split(".")[2])
                if lid > max_lid:
                    max_lid = lid
            except (ValueError, IndexError):
                pass
    return max_lid + 1


def read_raw_bytes(shard_path, byte_start, byte_end, fh_cache):
    """Read tensor bytes from shard at given byte range. Caches file handles."""
    fh = fh_cache.get(shard_path)
    if fh is None:
        fh = open(shard_path, "rb")
        fh_cache[shard_path] = fh
    fh.seek(byte_start)
    return fh.read(byte_end - byte_start)


def main():
    global SRC_DIR, OUT_PATH
    ap = argparse.ArgumentParser(description="Convert Xiaomi MiMo-V2.5 FP8 shards to single safetensors for MLX block_fp8.")
    ap.add_argument("--src", default=SRC_DIR, help="Directory containing Xiaomi's *.safetensors shards (default: %(default)s)")
    ap.add_argument("--out", default=OUT_PATH, help="Output single safetensors path (default: %(default)s)")
    args = ap.parse_args()
    SRC_DIR = args.src
    OUT_PATH = args.out
    print(f"[config] src: {SRC_DIR}", flush=True)
    print(f"[config] out: {OUT_PATH}", flush=True)

    src_shards = sorted(glob.glob(os.path.join(SRC_DIR, "*.safetensors")))
    src_shards = [s for s in src_shards if not s.endswith("mimo_v2.5_block_fp8.safetensors")]
    if not src_shards:
        print(f"ERROR: no source shards in {SRC_DIR}", file=sys.stderr)
        sys.exit(1)

    key_info = read_shard_headers(src_shards)
    key_info = filter_keys(key_info)

    n_experts = detect_n_experts(key_info)
    n_layers = detect_n_layers(key_info)
    print(f"[detect] n_layers={n_layers}, n_routed_experts={n_experts}", flush=True)

    # Group MoE keys by (layer, proj, suffix)
    moe_groups = defaultdict(dict)
    moe_keys = set()
    for k in key_info:
        parts = k.split(".")
        if (len(parts) >= 7 and parts[0] == "model"
                and parts[1] == "layers" and parts[3] == "mlp"
                and parts[4] == "experts"):
            try:
                L = int(parts[2])
                E = int(parts[5])
                proj = parts[6]
                suffix = ".".join(parts[7:])
                moe_groups[(L, proj, suffix)][E] = k
                moe_keys.add(k)
            except (ValueError, IndexError):
                pass

    passthrough_keys = [k for k in key_info if k not in moe_keys]
    print(f"[plan] {len(moe_groups)} MoE groups, {len(passthrough_keys)} passthrough",
          flush=True)

    # Plan the output. Build the output header in memory; data is streamed.
    # We need to know byte offsets up front for the header, then write data
    # in the same order.
    fh_cache = {}
    t0 = time.time()

    # Build (out_key, dtype, shape, source_descriptor) plan
    # source_descriptor:
    #   ("passthrough", src_key)
    #   ("moe_stack",   [src_key per expert in order 0..n_experts-1])
    plan = []  # list of dicts
    print("[plan] building output plan...", flush=True)
    for k in sorted(passthrough_keys):
        sp, dtype, shape, bs, be = key_info[k]
        nbytes = be - bs
        plan.append({"out_key": k, "dtype": dtype, "shape": shape,
                     "nbytes": nbytes, "kind": "passthrough", "src": k})

    for L in range(n_layers):
        for proj in ("gate_proj", "down_proj", "up_proj"):
            for suffix in ("weight", "weight_scale_inv"):
                key = (L, proj, suffix)
                if key not in moe_groups:
                    continue
                em = moe_groups[key]
                if len(em) != n_experts:
                    raise RuntimeError(f"incomplete group {key}: {len(em)}/{n_experts}")
                # All experts should have the same dtype + shape
                first_src = em[0]
                _, dtype, shape, _, _ = key_info[first_src]
                per_expert_bytes = 1
                for d in shape:
                    per_expert_bytes *= d
                per_expert_bytes *= DTYPE_SIZE[dtype]
                stacked_shape = [n_experts] + shape
                stacked_nbytes = per_expert_bytes * n_experts
                src_keys = [em[E] for E in range(n_experts)]
                out_key = f"model.layers.{L}.mlp.switch_mlp.{proj}.{suffix}"
                plan.append({
                    "out_key": out_key, "dtype": dtype, "shape": stacked_shape,
                    "nbytes": stacked_nbytes, "kind": "moe_stack", "src": src_keys,
                })

    total_payload = sum(p["nbytes"] for p in plan)
    print(f"[plan] {len(plan)} output tensors, payload {total_payload/1024**3:.1f} GB",
          flush=True)

    # Build the safetensors header. data_offsets are relative to the start of
    # the data section (post-header).
    print("[header] building output header...", flush=True)
    out_hdr = {}
    cursor = 0
    for p in plan:
        out_hdr[p["out_key"]] = {
            "dtype": p["dtype"],
            "shape": p["shape"],
            "data_offsets": [cursor, cursor + p["nbytes"]],
        }
        cursor += p["nbytes"]
    out_hdr["__metadata__"] = {"format": "block_fp8_v1"}

    hdr_bytes = json.dumps(out_hdr, separators=(",", ":")).encode("utf-8")
    # safetensors spec: header length is little-endian uint64, then header,
    # then data. Pad header to 8-byte alignment.
    pad = (-len(hdr_bytes)) % 8
    hdr_bytes_padded = hdr_bytes + b" " * pad
    hdr_len = len(hdr_bytes_padded)
    print(f"[header] header size: {hdr_len/1024:.1f} KB", flush=True)

    # Stream-write the output file
    print(f"[write] streaming to {OUT_PATH}...", flush=True)
    with open(OUT_PATH, "wb") as out_fh:
        out_fh.write(struct.pack("<Q", hdr_len))
        out_fh.write(hdr_bytes_padded)

        # Write data in the same order as plan
        for i, p in enumerate(plan):
            if p["kind"] == "passthrough":
                sp, dtype, shape, bs, be = key_info[p["src"]]
                buf = read_raw_bytes(sp, bs, be, fh_cache)
                out_fh.write(buf)
            else:  # moe_stack
                for src_key in p["src"]:
                    sp, dtype, shape, bs, be = key_info[src_key]
                    buf = read_raw_bytes(sp, bs, be, fh_cache)
                    out_fh.write(buf)
            # Progress
            if (i + 1) % 50 == 0 or i + 1 == len(plan):
                elapsed = time.time() - t0
                done_bytes = sum(plan[j]["nbytes"] for j in range(i + 1))
                pct = 100 * done_bytes / total_payload
                rate = done_bytes / elapsed / 1024**3 if elapsed > 0 else 0
                eta = (total_payload - done_bytes) / 1024**3 / max(rate, 0.01)
                print(f"  [{i+1}/{len(plan)}] {pct:.1f}% wrote "
                      f"{done_bytes/1024**3:.1f}/{total_payload/1024**3:.1f} GB "
                      f"@ {rate:.2f} GB/s ETA {eta:.0f}s", flush=True)

    # Close any cached handles
    for fh in fh_cache.values():
        fh.close()

    elapsed = time.time() - t0
    print(f"[done] elapsed={elapsed:.0f}s output={OUT_PATH}", flush=True)
    print(f"[done] avg throughput: {total_payload/elapsed/1024**3:.2f} GB/s",
          flush=True)


if __name__ == "__main__":
    main()