#!/usr/bin/env python
"""Strip torch.compile's `_orig_mod.` prefix from a saved checkpoint.

The physical_ai_ft run was launched with `compile: true` and saved the COMPILED
module, so 167 of its 489 tensors are named `expert._orig_mod.*` while the model
class expects `expert.*`.

This matters more than a normal key mismatch. LeRobot's `PreTrainedPolicy.
from_pretrained` loads with `strict=False` and merely LOGS missing keys, and
`FlowMatchingExpert.__init__` zero-initialises `action_out` — so a model loaded
from the raw checkpoint predicts velocity 0, and `predict_action_chunk` returns
its initial Gaussian noise unchanged. That is a plausible-looking bad number, not
a crash: nothing anywhere in the stack tells you the action expert is random.

Usage:
    python scripts/rekey_checkpoint.py <src_ckpt_dir> <dst_ckpt_dir>
"""

from __future__ import annotations

import shutil
import sys
from pathlib import Path

from safetensors.torch import load_file, save_file

PREFIX = "expert._orig_mod."


def main():
    src = Path(sys.argv[1] if len(sys.argv) > 1
               else "/home/alexw/tinyvla_data/b200/outputs/physical_ai_ft/final")
    dst = Path(sys.argv[2] if len(sys.argv) > 2
               else "/home/alexw/tinyvla/outputs/physical_ai_ft_fixed")

    dst.mkdir(parents=True, exist_ok=True)
    for f in src.iterdir():
        if f.suffix == ".json" or f.name.startswith("README"):
            shutil.copy(f, dst / f.name)

    sd = load_file(str(src / "model.safetensors"))
    renamed = {k.replace(PREFIX, "expert."): v for k, v in sd.items()}
    n_hit = sum(1 for k in sd if k.startswith(PREFIX))
    assert len(renamed) == len(sd), "rename collided — two keys mapped to one name"
    print(f"renamed {n_hit}/{len(sd)} tensors ({PREFIX} -> expert.)")

    save_file(renamed, str(dst / "model.safetensors"), metadata={"format": "pt"})

    # Hard gate: build the model from the checkpoint's own config and demand a
    # STRICT load. from_pretrained would silently accept a missing expert.
    from tinyvla.modeling_tinyvla import TinyVLAPolicy

    pol = TinyVLAPolicy.from_pretrained(str(dst))
    fresh = load_file(str(dst / "model.safetensors"))
    missing, unexpected = pol.load_state_dict(fresh, strict=True), None
    n_model = len(pol.state_dict())
    aow = pol.expert.action_out.weight.abs().mean().item()

    print(f"{len(fresh)}/{n_model} tensors, strict=True OK")
    print(f"expert.action_out |w| = {aow:.4f}  (zero at init -> nonzero proves training)")
    assert aow > 1e-3, "action_out is at its zero init — the expert never trained"
    print(f"wrote {dst}")


if __name__ == "__main__":
    main()