"""Convert Pixal3D checkpoints to MLX-ready safetensors.

Much lighter than LATO.2's converter: Pixal3D ships safetensors already (not pickled
.pt), and each checkpoint has a sibling .json naming the model class and its exact
constructor args — so there is no architecture to infer from tensor shapes.

The only real work is the rank-5 layout question, handled by `trellis_sparse_mlx.convert`
(spconv KRSC vs dense torch Conv3d — indistinguishable by rank, silently wrong if
guessed). Measured against the released weights:

    ss_flow / slat_flow x3   ZERO 5-D tensors — pure transformers, pass straight through
    shape_dec / tex_dec      40 KRSC sparse kernels each
    ss_dec                   dense Conv3d

So the ~20GB of flow models need no conversion at all; only the decoders do.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import mlx.core as mx
import numpy as np

from trellis_sparse_mlx.convert import convert_tensor


def load_config(ckpt: Path) -> dict:
    """Each checkpoint has a sibling .json with {name, args} — the authoritative config."""
    cfg = ckpt.with_suffix(".json")
    if not cfg.exists():
        raise FileNotFoundError(f"no config beside {ckpt.name}")
    return json.loads(cfg.read_text())


def convert_checkpoint(src: Path, dst: Path, flip_kernel: bool = False) -> dict:
    w = mx.load(str(src))
    out, counts = {}, {"dense": 0, "krsc": 0, "conv3d": 0}
    for k, v in w.items():
        dt = v.dtype
        # Only rank-5 tensors are ever rearranged, and the rearrangement is a pure
        # transpose/reshape — so round-trip through float32 for numpy only where needed
        # and restore the original dtype. Upcasting everything to fp32 would double
        # 24GB of checkpoints on disk for no benefit.
        if v.ndim != 5:
            out[k] = (v, "dense")
            counts["dense"] += 1
            continue
        arr = np.asarray(v.astype(mx.float32))
        arr, kind = convert_tensor(k, arr, flip_kernel)
        counts[kind] += 1
        out[k] = (mx.array(arr).astype(dt), kind)
    dst.parent.mkdir(parents=True, exist_ok=True)
    mx.save_safetensors(str(dst), {k: v for k, (v, _) in out.items()})
    cfg = load_config(src)
    return {
        "model": cfg.get("name"),
        "tensors": len(out),
        **counts,
        "gb": dst.stat().st_size / 1e9,
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="Pixal3D -> MLX weight converter")
    ap.add_argument(
        "--ckpt",
        default="/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts",
    )
    ap.add_argument("--out", default="weights")
    ap.add_argument("--only", default=None, help="substring filter, e.g. 'dec'")
    ap.add_argument(
        "--flip-kernel",
        action="store_true",
        help="mirror sparse kernel offsets (orientation experiment; see shared core)",
    )
    ap.add_argument(
        "--skip-passthrough",
        action="store_true",
        help="skip checkpoints with no 5-D tensors — they need no conversion",
    )
    a = ap.parse_args()

    src_dir, out_dir = Path(a.ckpt), Path(a.out)
    files = sorted(src_dir.glob("*.safetensors"))
    if a.only:
        files = [f for f in files if a.only in f.name]
    if not files:
        raise SystemExit(f"no .safetensors in {src_dir}")

    for f in files:
        if a.skip_passthrough:
            w = mx.load(str(f))
            if not any(v.ndim == 5 for v in w.values()):
                print(f"  {f.stem[:44]:46s} skipped (no 5-D tensors)")
                continue
        info = convert_checkpoint(f, out_dir / f.name, a.flip_kernel)
        print(
            f"  {f.stem[:44]:46s} {info['model'][:26]:28s} "
            f"{info['tensors']:4d}t  krsc {info['krsc']:3d}  conv3d {info['conv3d']:3d}  "
            f"{info['gb']:5.2f} GB"
        )


if __name__ == "__main__":
    main()