"""Losslessly convert the upstream FP32 checkpoint to native MLX layout."""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import mlx.core as mx


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)

    weights = mx.load(str(args.source / "model.safetensors"))
    converted = {}
    for name, tensor in weights.items():
        if name.endswith("conv.conv.weight"):
            if tensor.ndim != 3 or tensor.shape[1:] != (1, 3):
                raise ValueError(f"Unexpected convolution shape for {name}: {tensor.shape}")
            tensor = tensor.transpose(0, 2, 1)
        if tensor.dtype != mx.float32:
            raise TypeError(f"Expected FP32 source tensor {name}, got {tensor.dtype}")
        converted[name] = tensor

    mx.save_safetensors(
        str(args.output / "model.safetensors"),
        converted,
        metadata={
            "format": "mlx",
            "source": "LiquidAI/LFM2.5-Encoder-350M",
            "source_revision": "b886781f7c6f10ca9b7096e21b83e30a073c2f39",
            "dtype": "float32",
        },
    )
    for filename in (
        "config.json",
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
        "LICENSE",
    ):
        shutil.copy2(args.source / filename, args.output / filename)

    manifest = {
        "source": "LiquidAI/LFM2.5-Encoder-350M",
        "source_revision": "b886781f7c6f10ca9b7096e21b83e30a073c2f39",
        "parameter_dtype": "float32",
        "parameter_count": int(sum(tensor.size for tensor in converted.values())),
        "tensor_count": len(converted),
        "transforms": ["conv weights [out,1,kernel] -> [out,kernel,1]"],
    }
    (args.output / "conversion_manifest.json").write_text(
        json.dumps(manifest, indent=2) + "\n"
    )


if __name__ == "__main__":
    main()