#!/usr/bin/env python3
"""
Convert a Diffusers-format MiniMax-H3 transformer safetensors file back to the
native/ComfyUI MiniMax-H3 state-dict layout.

This is a streaming container conversion:
- direct-renamed tensors are copied byte-for-byte;
- Diffusers to_q/to_k/to_v tensors are concatenated into native qkv_proj;
- Diffusers SwiGLU [value; gate] is restored to native [gate; value];\n- pruned `folded_bias` F32 constants are restored as native AdaLN biases;
- rope.inv_freq, intentionally dropped by the Diffusers converter, is
  reconstructed from the MiniMax-H3 rope configuration.

The full ~40 GB model is never materialized in RAM.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import struct
from collections import Counter
from pathlib import Path

import torch
from safetensors import safe_open

CHUNK_SIZE = 64 * 1024 * 1024

Q_SUFFIX = ".attn.to_q.weight"
K_SUFFIX = ".attn.to_k.weight"
V_SUFFIX = ".attn.to_v.weight"
SWIGLU_SUFFIX = ".ff.net.0.proj.weight"

DTYPE_SIZE = {
    "BOOL": 1,
    "U8": 1,
    "I8": 1,
    "F8_E4M3": 1,
    "F8_E5M2": 1,
    "I16": 2,
    "U16": 2,
    "F16": 2,
    "BF16": 2,
    "I32": 4,
    "U32": 4,
    "F32": 4,
    "I64": 8,
    "U64": 8,
    "F64": 8,
}


def product(values):
    out = 1
    for value in values:
        out *= value
    return out


def read_header(path: Path):
    with path.open("rb") as f:
        raw = f.read(8)
        if len(raw) != 8:
            raise RuntimeError(f"{path}: truncated safetensors file")
        header_len = struct.unpack("<Q", raw)[0]
        if header_len <= 0 or header_len > 1024 * 1024 * 1024:
            raise RuntimeError(
                f"{path}: implausible safetensors header length: {header_len}"
            )
        raw_header = f.read(header_len)
        if len(raw_header) != header_len:
            raise RuntimeError(f"{path}: truncated safetensors header")

    try:
        header = json.loads(raw_header.decode("utf-8"))
    except Exception as exc:
        raise RuntimeError(f"{path}: invalid safetensors header: {exc}") from exc

    tensors = {k: v for k, v in header.items() if k != "__metadata__"}
    metadata = header.get("__metadata__", {})
    if metadata is None:
        metadata = {}
    if not isinstance(metadata, dict):
        raise RuntimeError(f"{path}: invalid __metadata__")

    return {
        "header_len": header_len,
        "data_start": 8 + header_len,
        "tensors": tensors,
        "metadata": metadata,
    }


def validate_tensor_info(path: Path, header, key: str):
    info = header["tensors"][key]
    dtype = info.get("dtype")
    shape = info.get("shape")
    offsets = info.get("data_offsets")

    if dtype not in DTYPE_SIZE:
        raise RuntimeError(f"{key}: unsupported dtype {dtype!r}")
    if not isinstance(shape, list) or not all(
        isinstance(x, int) and x >= 0 for x in shape
    ):
        raise RuntimeError(f"{key}: invalid shape {shape!r}")
    if (
        not isinstance(offsets, list)
        or len(offsets) != 2
        or not all(isinstance(x, int) for x in offsets)
    ):
        raise RuntimeError(f"{key}: invalid data_offsets {offsets!r}")

    start, end = offsets
    if start < 0 or end < start:
        raise RuntimeError(f"{key}: invalid data_offsets {offsets!r}")

    file_payload_size = path.stat().st_size - header["data_start"]
    if end > file_payload_size:
        raise RuntimeError(
            f"{key}: data_offsets extend beyond source payload "
            f"({end} > {file_payload_size})"
        )

    expected = product(shape) * DTYPE_SIZE[dtype]
    actual = end - start
    if expected != actual:
        raise RuntimeError(
            f"{key}: payload byte count does not match shape/dtype "
            f"({actual} != {expected})"
        )


def tensor_size(header, key: str):
    start, end = header["tensors"][key]["data_offsets"]
    return end - start


def copy_region(source_f, output_f, absolute_offset: int, size: int, digest):
    source_f.seek(absolute_offset)
    remaining = size
    while remaining:
        block = source_f.read(min(CHUNK_SIZE, remaining))
        if not block:
            raise RuntimeError(
                f"Unexpected EOF while copying; {remaining:,} bytes remained"
            )
        output_f.write(block)
        digest.update(block)
        remaining -= len(block)


def copy_tensor(source_f, output_f, source_header, key: str, digest):
    info = source_header["tensors"][key]
    begin, end = info["data_offsets"]
    copy_region(
        source_f,
        output_f,
        source_header["data_start"] + begin,
        end - begin,
        digest,
    )


def direct_native_key(key: str):
    if key.endswith(Q_SUFFIX) or key.endswith(K_SUFFIX) or key.endswith(V_SUFFIX):
        return None
    if key.endswith(SWIGLU_SUFFIX):
        return None

    # The Diffusers pruned H3 representation keeps the ComfyUI/Comfy-Org
    # AdaLN coordinate table under time_embedder.table. Native ComfyUI detects
    # the pruned architecture from the top-level adaln_t_table key.
    if key == "time_embedder.table":
        return "adaln_t_table"

    # multimodalart/diffusers pruned H3 keeps the reduced-rank AdaLN constant
    # term outside the bias-free Linear as a float32 `folded_bias` buffer.
    # Native ComfyUI's pruned H3 Linear is float32 and expects this exact
    # constant as its bias.
    if key == "norm_out.folded_bias":
        return "final_layer.adaln_proj.linear.bias"

    # Top-level module renames must be prefix-scoped. Using unrestricted
    # str.replace() here would corrupt names such as
    # time_embedder.linear_1 -> time_embedder.proj_in -> time_embedder.video_patch_proj.
    top_level_prefixes = (
        ("time_embedder.linear_1.", "time_embedder.proj_in."),
        ("time_embedder.linear_2.", "time_embedder.proj_out."),
        ("audio_proj_out.", "final_layer.audio_out."),
        ("audio_proj_in.", "audio_patch_proj."),
        ("proj_out.", "final_layer.video_out."),
        ("proj_in.", "video_patch_proj."),
        ("context_embedder.", "condition_proj."),
        ("norm_out.norm.", "final_layer.norm."),
        ("norm_out.linear.", "final_layer.adaln_proj.linear."),
    )
    for old, new in top_level_prefixes:
        if key.startswith(old):
            return new + key[len(old):]

    out = key
    if out.startswith("token_refiner.refiner_blocks."):
        out = out.replace(
            "token_refiner.refiner_blocks.", "token_refiner.blocks.", 1
        )
    elif out.startswith("transformer_blocks."):
        out = out.replace("transformer_blocks.", "blocks.", 1)

    out = out.replace(".adaln_proj.folded_bias", ".adaln_proj.linear.bias")
    out = out.replace(".attn.norm_q.", ".attn.q_norm.")
    out = out.replace(".attn.norm_k.", ".attn.k_norm.")
    out = out.replace(".attn.to_out.0.", ".attn.out_proj.")
    out = out.replace(".ff.net.2.", ".mlp.fc2.")
    return out

def qkv_native_key(q_key: str):
    base = q_key[: -len(Q_SUFFIX)]
    if base.startswith("transformer_blocks."):
        base = base.replace("transformer_blocks.", "blocks.", 1)
    elif base.startswith("token_refiner.refiner_blocks."):
        base = base.replace(
            "token_refiner.refiner_blocks.", "token_refiner.blocks.", 1
        )
    else:
        raise RuntimeError(f"Unexpected MiniMax-H3 Diffusers Q key: {q_key}")
    return base + ".attn.qkv_proj.weight"


def swiglu_native_key(key: str):
    out = key
    if out.startswith("transformer_blocks."):
        out = out.replace("transformer_blocks.", "blocks.", 1)
    elif out.startswith("token_refiner.refiner_blocks."):
        out = out.replace(
            "token_refiner.refiner_blocks.", "token_refiner.blocks.", 1
        )
    else:
        raise RuntimeError(f"Unexpected MiniMax-H3 Diffusers SwiGLU key: {key}")
    return out.replace(".ff.net.0.proj.weight", ".mlp.fc1.weight")


def make_rope_inv_freq_bytes():
    # Matches the upstream MiniMax-H3/Diffusers reconstruction:
    # 1 / 10000 ** (arange(0, 2 * 16, 2) / (2 * 16))
    inv_freq = 1.0 / (
        10000.0
        ** (
            torch.arange(0, 32, 2, dtype=torch.float32)
            / 32.0
        )
    )
    inv_freq = inv_freq.contiguous().cpu()
    if tuple(inv_freq.shape) != (16,) or inv_freq.dtype != torch.float32:
        raise RuntimeError("Internal rope.inv_freq reconstruction invariant failed")
    return inv_freq.numpy().tobytes(order="C")


def build_plan(source_path: Path, source_header):
    tensors = source_header["tensors"]
    keys = list(tensors)

    required_source_keys = (
        "proj_in.weight",
        "audio_proj_in.weight",
        "context_embedder.weight",
        "transformer_blocks.0.attn.to_q.weight",
        "token_refiner.refiner_blocks.0.attn.to_q.weight",
        "proj_out.weight",
        "audio_proj_out.weight",
    )
    missing = [key for key in required_source_keys if key not in tensors]
    if missing:
        raise RuntimeError(
            "Input does not look like the expected Diffusers MiniMax-H3 "
            "transformer. Missing:\n  " + "\n  ".join(missing)
        )

    native_markers = (
        "video_patch_proj.weight",
        "audio_patch_proj.weight",
        "blocks.0.attn.qkv_proj.weight",
        "rope.inv_freq",
    )
    present_native = [key for key in native_markers if key in tensors]
    if present_native:
        raise RuntimeError(
            "Input already contains native/ComfyUI MiniMax-H3 keys:\n  "
            + "\n  ".join(present_native)
        )

    is_pruned = "time_embedder.table" in tensors

    full_time_keys = {
        "time_embedder.linear_1.weight",
        "time_embedder.linear_1.bias",
        "time_embedder.linear_2.weight",
        "time_embedder.linear_2.bias",
    }

    if is_pruned:
        unexpected_full = sorted(full_time_keys & set(keys))
        if unexpected_full:
            raise RuntimeError(
                "Checkpoint mixes the pruned AdaLN table with full timestep-MLP "
                "weights:\n  " + "\n  ".join(unexpected_full)
            )

        table = tensors["time_embedder.table"]
        if table.get("shape") != [1025, 8]:
            raise RuntimeError(
                "Expected pruned MiniMax-H3 time_embedder.table shape [1025, 8], "
                f"got {table.get('shape')}"
            )
        if table.get("dtype") != "F32":
            raise RuntimeError(
                "Expected pruned MiniMax-H3 time_embedder.table dtype F32, "
                f"got {table.get('dtype')}"
            )

        folded_bias_keys = [
            f"transformer_blocks.{i}.adaln_proj.folded_bias" for i in range(50)
        ] + ["norm_out.folded_bias"]
        missing_folded = [k for k in folded_bias_keys if k not in tensors]
        if missing_folded:
            raise RuntimeError(
                "Pruned Diffusers MiniMax-H3 checkpoint is missing required "
                "float32 folded AdaLN constant terms:\n  "
                + "\n  ".join(missing_folded)
            )
        for i in range(50):
            info = tensors[f"transformer_blocks.{i}.adaln_proj.folded_bias"]
            if info.get("shape") != [96768] or info.get("dtype") != "F32":
                raise RuntimeError(
                    f"transformer_blocks.{i}.adaln_proj.folded_bias: expected "
                    f"[96768] F32, got {info.get('shape')} {info.get('dtype')}"
                )
        final_bias = tensors["norm_out.folded_bias"]
        if final_bias.get("shape") != [10752] or final_bias.get("dtype") != "F32":
            raise RuntimeError(
                "norm_out.folded_bias: expected [10752] F32, got "
                f"{final_bias.get('shape')} {final_bias.get('dtype')}"
            )

        for adaln_key in (
            "transformer_blocks.0.adaln_proj.linear.weight",
            "norm_out.linear.weight",
        ):
            if adaln_key not in tensors:
                raise RuntimeError(
                    f"Pruned MiniMax-H3 checkpoint is missing {adaln_key}"
                )
            shape = tensors[adaln_key].get("shape")
            if not isinstance(shape, list) or len(shape) != 2 or shape[1] != 8:
                raise RuntimeError(
                    f"{adaln_key}: expected pruned input width 8, got {shape}"
                )
    else:
        missing_time = sorted(full_time_keys - set(keys))
        if missing_time:
            raise RuntimeError(
                "Checkpoint has neither a complete full timestep MLP nor the "
                "pruned time_embedder.table. Missing:\n  "
                + "\n  ".join(missing_time)
            )

    for key in keys:
        validate_tensor_info(source_path, source_header, key)

    q_keys = [key for key in keys if key.endswith(Q_SUFFIX)]
    transformer_q = [
        key for key in q_keys if key.startswith("transformer_blocks.")
    ]
    refiner_q = [
        key
        for key in q_keys
        if key.startswith("token_refiner.refiner_blocks.")
    ]

    if len(transformer_q) != 50:
        raise RuntimeError(
            f"Expected 50 transformer attention blocks, found {len(transformer_q)}"
        )
    if len(refiner_q) != 2:
        raise RuntimeError(
            f"Expected 2 token-refiner attention blocks, found {len(refiner_q)}"
        )

    used = set()
    plan = []

    for key in keys:
        if key in used:
            continue

        if is_pruned and key in {"adaln_basis", "adaln_mean"}:
            used.add(key)
            continue

        if key.endswith(Q_SUFFIX):
            prefix = key[: -len(Q_SUFFIX)]
            k_key = prefix + K_SUFFIX
            v_key = prefix + V_SUFFIX
            if k_key not in tensors or v_key not in tensors:
                raise RuntimeError(f"Incomplete Q/K/V group for {key}")

            q_info = tensors[key]
            k_info = tensors[k_key]
            v_info = tensors[v_key]

            if not (
                q_info["dtype"] == k_info["dtype"] == v_info["dtype"]
            ):
                raise RuntimeError(f"Q/K/V dtype mismatch for {prefix}")

            if not (
                q_info["shape"] == k_info["shape"] == v_info["shape"]
            ):
                raise RuntimeError(f"Q/K/V shape mismatch for {prefix}")

            q_shape = q_info["shape"]
            if len(q_shape) != 2:
                raise RuntimeError(
                    f"Expected 2D Q/K/V weights for {prefix}, got {q_shape}"
                )

            plan.append(
                {
                    "kind": "qkv",
                    "out_key": qkv_native_key(key),
                    "dtype": q_info["dtype"],
                    "shape": [q_shape[0] * 3, q_shape[1]],
                    "nbytes": (
                        tensor_size(source_header, key)
                        + tensor_size(source_header, k_key)
                        + tensor_size(source_header, v_key)
                    ),
                    "src_keys": [key, k_key, v_key],
                }
            )
            used.update((key, k_key, v_key))
            continue

        if key.endswith(K_SUFFIX) or key.endswith(V_SUFFIX):
            # Its Q key will consume it. If not, the unhandled check below fails.
            continue

        if key.endswith(SWIGLU_SUFFIX):
            info = tensors[key]
            shape = info["shape"]
            nbytes = tensor_size(source_header, key)

            if len(shape) != 2 or shape[0] % 2:
                raise RuntimeError(
                    f"Invalid Diffusers SwiGLU tensor shape for {key}: {shape}"
                )
            if nbytes % 2:
                raise RuntimeError(
                    f"Invalid Diffusers SwiGLU byte count for {key}: {nbytes}"
                )

            plan.append(
                {
                    "kind": "swiglu",
                    "out_key": swiglu_native_key(key),
                    "dtype": info["dtype"],
                    "shape": shape,
                    "nbytes": nbytes,
                    "src_key": key,
                }
            )
            used.add(key)
            continue

        out_key = direct_native_key(key)
        if out_key is None:
            raise RuntimeError(f"Unhandled source key: {key}")

        info = tensors[key]
        plan.append(
            {
                "kind": "copy",
                "out_key": out_key,
                "dtype": info["dtype"],
                "shape": info["shape"],
                "nbytes": tensor_size(source_header, key),
                "src_key": key,
            }
        )
        used.add(key)

    unhandled = sorted(set(keys) - used)
    if unhandled:
        raise RuntimeError(
            "Source contains keys that were not converted:\n  "
            + "\n  ".join(unhandled)
        )

    # Diffusers intentionally drops native rope.inv_freq; restore it.
    rope_bytes = make_rope_inv_freq_bytes()
    plan.append(
        {
            "kind": "generated",
            "out_key": "rope.inv_freq",
            "dtype": "F32",
            "shape": [16],
            "nbytes": len(rope_bytes),
            "data": rope_bytes,
        }
    )

    counts = Counter(entry["out_key"] for entry in plan)
    duplicates = sorted(key for key, count in counts.items() if count != 1)
    if duplicates:
        raise RuntimeError(
            "Conversion would create duplicate output keys:\n  "
            + "\n  ".join(duplicates)
        )

    dropped_aux_count = sum(
        1 for k in ("adaln_basis", "adaln_mean") if is_pruned and k in tensors
    )
    expected_out_count = len(keys) - 2 * len(q_keys) + 1 - dropped_aux_count
    if len(plan) != expected_out_count:
        raise RuntimeError(
            f"Output tensor-count invariant failed: plan={len(plan)}, "
            f"expected={expected_out_count}"
        )

    return plan, is_pruned


def build_output_header(plan, source_metadata, is_pruned):
    offset = 0
    tensors = {}

    for entry in plan:
        begin = offset
        end = begin + entry["nbytes"]
        tensors[entry["out_key"]] = {
            "dtype": entry["dtype"],
            "shape": entry["shape"],
            "data_offsets": [begin, end],
        }
        offset = end

    metadata = {}
    for key, value in source_metadata.items():
        if not isinstance(key, str) or not isinstance(value, str):
            raise RuntimeError(
                f"Safetensors metadata must be string:string; got "
                f"{key!r}:{value!r}"
            )
        metadata[key] = value

    # A Diffusers/full-model config can carry time_embed_dim=2688. Keeping that
    # metadata on a pruned ComfyUI checkpoint would override the tensor-derived
    # width from adaln_t_table (8) in ComfyUI model detection.
    if is_pruned:
        metadata.pop("config", None)

    metadata["format"] = metadata.get("format", "pt")
    metadata["converted_from"] = "diffusers MiniMax-H3 transformer"
    metadata["converted_to"] = "native ComfyUI MiniMax-H3 diffusion model"

    full_header = {"__metadata__": metadata, **tensors}
    raw = json.dumps(
        full_header, separators=(",", ":"), ensure_ascii=False
    ).encode("utf-8")
    raw += b" " * ((-len(raw)) % 8)
    return raw, offset


def write_output(source_path, output_path, source_header, plan, is_pruned):
    raw_header, payload_size = build_output_header(
        plan, source_header["metadata"], is_pruned
    )
    output_path.parent.mkdir(parents=True, exist_ok=True)
    partial_path = output_path.with_name(output_path.name + ".partial")

    if partial_path.exists():
        partial_path.unlink()

    digest = hashlib.sha256()

    try:
        with source_path.open("rb") as source_f, partial_path.open("xb") as out_f:
            out_f.write(struct.pack("<Q", len(raw_header)))
            out_f.write(raw_header)

            total = len(plan)
            for index, entry in enumerate(plan, 1):
                kind = entry["kind"]

                if kind == "copy":
                    copy_tensor(
                        source_f,
                        out_f,
                        source_header,
                        entry["src_key"],
                        digest,
                    )

                elif kind == "qkv":
                    # Diffusers stores three separate, contiguous [q], [k], [v]
                    # matrices. ComfyUI's qkv_proj expects [q_all; k_all; v_all].
                    for key in entry["src_keys"]:
                        copy_tensor(
                            source_f, out_f, source_header, key, digest
                        )

                elif kind == "swiglu":
                    # Diffusers SwiGLU is [value; gate].
                    # Native MiniMax-H3 is [gate; value].
                    info = source_header["tensors"][entry["src_key"]]
                    begin, end = info["data_offsets"]
                    size = end - begin
                    half = size // 2
                    absolute = source_header["data_start"] + begin

                    copy_region(
                        source_f, out_f, absolute + half, half, digest
                    )
                    copy_region(source_f, out_f, absolute, half, digest)

                elif kind == "generated":
                    data = entry["data"]
                    out_f.write(data)
                    digest.update(data)

                else:
                    raise RuntimeError(f"Unknown conversion plan kind: {kind}")

                if index % 25 == 0 or index == total:
                    print(
                        f"\rConverted tensors: {index:,}/{total:,}",
                        end="",
                        flush=True,
                    )

            print()
            out_f.flush()
            os.fsync(out_f.fileno())

        expected_size = 8 + len(raw_header) + payload_size
        actual_size = partial_path.stat().st_size
        if actual_size != expected_size:
            raise RuntimeError(
                f"Output size mismatch: {actual_size:,} != {expected_size:,}"
            )

        os.replace(partial_path, output_path)

    except Exception:
        try:
            if partial_path.exists():
                partial_path.unlink()
        finally:
            raise

    return digest.hexdigest(), payload_size


def validate_output(output_path: Path, expected_count: int, is_pruned: bool):
    header = read_header(output_path)
    keys = set(header["tensors"])

    required = {
        "video_patch_proj.weight",
        "audio_patch_proj.weight",
        "condition_proj.weight",
        "blocks.0.attn.qkv_proj.weight",
        "blocks.49.attn.qkv_proj.weight",
        "token_refiner.blocks.0.attn.qkv_proj.weight",
        "token_refiner.blocks.1.attn.qkv_proj.weight",
        "final_layer.video_out.weight",
        "final_layer.audio_out.weight",
        "rope.inv_freq",
    }
    if is_pruned:
        required.add("adaln_t_table")
        required.add("final_layer.adaln_proj.linear.bias")
        required.update(
            f"blocks.{i}.adaln_proj.linear.bias" for i in range(50)
        )
    else:
        required.update({
            "time_embedder.proj_in.weight",
            "time_embedder.proj_out.weight",
        })

    forbidden = {
        "proj_in.weight",
        "audio_proj_in.weight",
        "context_embedder.weight",
        "transformer_blocks.0.attn.to_q.weight",
        "token_refiner.refiner_blocks.0.attn.to_q.weight",
        "time_embedder.table",
    }
    if is_pruned:
        forbidden.update({
            "time_embedder.proj_in.weight",
            "time_embedder.proj_in.bias",
            "time_embedder.proj_out.weight",
            "time_embedder.proj_out.bias",
            "norm_out.folded_bias",
            "adaln_basis",
            "adaln_mean",
        })
        forbidden.update(
            f"blocks.{i}.adaln_proj.folded_bias" for i in range(50)
        )
    else:
        forbidden.add("adaln_t_table")

    missing = sorted(required - keys)
    stale = sorted(forbidden & keys)

    if missing:
        raise RuntimeError(
            "Converted output is missing required native keys:\n  "
            + "\n  ".join(missing)
        )
    if stale:
        raise RuntimeError(
            "Converted output still contains Diffusers keys:\n  "
            + "\n  ".join(stale)
        )
    if len(keys) != expected_count:
        raise RuntimeError(
            f"Converted output tensor count is {len(keys)}, expected "
            f"{expected_count}"
        )

    if is_pruned:
        table_info = header["tensors"]["adaln_t_table"]
        if table_info["shape"] != [1025, 8] or table_info["dtype"] != "F32":
            raise RuntimeError(
                "Converted adaln_t_table does not preserve [1025, 8] F32"
            )
        for i in range(50):
            info = header["tensors"][f"blocks.{i}.adaln_proj.linear.bias"]
            if info["shape"] != [96768] or info["dtype"] != "F32":
                raise RuntimeError(
                    f"blocks.{i}.adaln_proj.linear.bias does not preserve "
                    "[96768] F32"
                )
        final_bias = header["tensors"]["final_layer.adaln_proj.linear.bias"]
        if final_bias["shape"] != [10752] or final_bias["dtype"] != "F32":
            raise RuntimeError(
                "final_layer.adaln_proj.linear.bias does not preserve [10752] F32"
            )
        if "config" in header["metadata"]:
            raise RuntimeError(
                "Pruned output still contains config metadata that can override "
                "shape-based ComfyUI detection"
            )

    for key in header["tensors"]:
        validate_tensor_info(output_path, header, key)

    with safe_open(str(output_path), framework="pt", device="cpu") as f:
        safe_keys = set(f.keys())
        if safe_keys != keys:
            raise RuntimeError(
                "safetensors.safe_open key set differs from parsed header"
            )

        rope = f.get_tensor("rope.inv_freq")
        expected_rope = 1.0 / (
            10000.0
            ** (
                torch.arange(0, 32, 2, dtype=torch.float32)
                / 32.0
            )
        )
        if not torch.equal(rope, expected_rope):
            raise RuntimeError("rope.inv_freq reconstruction mismatch")

    return len(keys)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()

    source_path = args.source.expanduser().resolve()
    output_path = args.output.expanduser().resolve()

    if not source_path.is_file():
        raise RuntimeError(f"Source file does not exist: {source_path}")
    if source_path == output_path:
        raise RuntimeError(
            "Source and output must be different files; validate before replacing"
        )
    if output_path.exists():
        raise RuntimeError(
            f"Output already exists; refusing to overwrite: {output_path}"
        )

    source_header = read_header(source_path)

    print(f"Source: {source_path}")
    print(f"Source tensors: {len(source_header['tensors']):,}")
    print("Validating source and building reverse-conversion plan...")

    plan, is_pruned = build_plan(source_path, source_header)
    qkv_count = sum(entry["kind"] == "qkv" for entry in plan)
    swiglu_count = sum(entry["kind"] == "swiglu" for entry in plan)

    print(f"Checkpoint variant: {'pruned AdaLN table' if is_pruned else 'full timestep MLP'}")
    print(f"Native output tensors: {len(plan):,}")
    print(f"QKV groups recombined: {qkv_count:,}")
    print(f"SwiGLU tensors unswapped: {swiglu_count:,}")
    print("Reconstructed tensors: 1 (rope.inv_freq)")
    print(f"Output: {output_path}")
    print()

    payload_hash, payload_size = write_output(
        source_path, output_path, source_header, plan, is_pruned
    )

    print(f"Converted payload: {payload_size / (1024**3):.3f} GiB")
    print(f"Converted payload SHA256: {payload_hash}")
    print("Validating complete output with safetensors...")

    count = validate_output(output_path, len(plan), is_pruned)

    print()
    print("safetensors.safe_open: OK")
    print(f"Native tensor count: {count:,}")
    print("Required native MiniMax-H3 keys: OK")
    if is_pruned:
        print("adaln_t_table mapping/shape: OK")
        print("Stale config metadata: removed")
    print("Stale Diffusers MiniMax-H3 keys: none")
    print("rope.inv_freq reconstruction: OK")
    print()
    print("CONVERSION SUCCESSFUL")
    print(output_path)


if __name__ == "__main__":
    main()