"""Convert the official BEVFormer-small checkpoint to a safetensors state_dict for the HF weights mirror.

Reads the mmcv checkpoint (as data only: torch.load(weights_only=True)), keeps `state_dict` (drops `optimizer` and the
training `meta`), and writes `bevformer_small_epoch_24.safetensors` plus a provenance JSON. Verifies a bit-exact
round trip of every tensor.

Usage:
    python -I make_weights_mirror.py <checkpoint.pth> <out_dir>
"""

import hashlib
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file

SOURCE_URL = "https://github.com/zhiqi-li/storage/releases/download/v1.0/bevformer_small_epoch_24.pth"
EXPECTED_SHA256 = "4ebc5810201ca1e609c29452f07c08bc4d9c9ebfda5d7f8ee1034dd446274bc5"


def sha256_of(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def main() -> None:
    ckpt_path = Path(sys.argv[1])
    out_dir = Path(sys.argv[2])
    out_dir.mkdir(parents=True, exist_ok=True)

    src_sha = sha256_of(ckpt_path)
    if src_sha != EXPECTED_SHA256:
        raise SystemExit(f"sha256 mismatch: {src_sha} != {EXPECTED_SHA256}")

    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    state = ckpt["state_dict"]
    tensors = {k: v.detach().contiguous().clone() for k, v in state.items()}
    meta = {
        "source_url": SOURCE_URL,
        "source_sha256": src_sha,
        "source_format": "mmcv checkpoint (state_dict only; optimizer and meta dropped)",
        "model": "BEVFormer-small (bevformer_small_epoch_24), nuScenes, 24 epochs",
    }
    out_file = out_dir / "bevformer_small_epoch_24.safetensors"
    save_file(tensors, str(out_file), metadata=meta)

    back = load_file(str(out_file))
    assert back.keys() == tensors.keys()
    for k, v in tensors.items():
        assert back[k].dtype == v.dtype and back[k].shape == v.shape and torch.equal(back[k], v), k

    n_params = sum(v.numel() for v in tensors.values() if v.is_floating_point())
    dtypes = sorted({str(v.dtype) for v in tensors.values()})
    prov = dict(meta)
    prov.update(
        {
            "num_tensors": len(tensors),
            "num_float_elements": n_params,
            "dtypes": dtypes,
            "safetensors_file": out_file.name,
            "safetensors_sha256": sha256_of(out_file),
            "safetensors_bytes": out_file.stat().st_size,
            "epoch": ckpt.get("meta", {}).get("epoch"),
            "iter": ckpt.get("meta", {}).get("iter"),
        }
    )
    (out_dir / "provenance.json").write_text(json.dumps(prov, indent=2) + "\n")
    print(json.dumps(prov, indent=2))


if __name__ == "__main__":
    main()