#!/usr/bin/env python
"""Grow a checkpoint's embodiment/state/action dimensions without losing what it learned.

Adding the unitree + navigation packs widens three things:
  max_state_dim   16 -> 32   (G1 Dex3 has 28 joint channels)
  max_action_dim   8 -> 32   (same, and nav actions are 3-dim)
  num_embodiments 16 -> 24   (packed ids reach 19)

Every affected weight is grown by ZERO padding, never re-initialised:

  state_proj.weight   (d, 16) -> (d, 32)   new input columns = 0
  action_in.weight    (d,  8) -> (d, 32)   new input columns = 0
  action_out.weight   (8,  d) -> (32,  d)  new output rows   = 0
  action_out.bias     (8,)    -> (32,)     new outputs       = 0
  embodiment emb      (16, d) -> (24, d)   new rows          = 0

Zeroed input columns receive zero-padded inputs anyway, and zeroed output rows
predict 0 for dims no old robot had — so on the OLD data the expanded model is
bit-for-bit the old model, and the new dims start from a clean slate. That makes
"did the new data help or hurt?" answerable, which a re-initialised head would not.

Usage:
    python scripts/expand_checkpoint.py outputs/tv2_C_scaled/final outputs/tv2_C_scaled_x32
"""

from __future__ import annotations

import json
import shutil
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file

# param name suffix -> axis that must grow (0 = rows/out, 1 = cols/in)
GROW = {
    "state_proj.weight": 1,
    "action_in.weight": 1,
    "action_out.weight": 0,
    "action_out.bias": 0,
}


def main():
    src = Path(sys.argv[1])
    dst = Path(sys.argv[2])
    new_state = int(sys.argv[3]) if len(sys.argv) > 3 else 32
    new_action = int(sys.argv[4]) if len(sys.argv) > 4 else 32
    new_emb = int(sys.argv[5]) if len(sys.argv) > 5 else 24

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

    cfg_path = dst / "config.json"
    cfg = json.loads(cfg_path.read_text())
    old_state = cfg.get("max_state_dim", 16)
    old_action = cfg.get("max_action_dim", 8)
    old_emb = cfg.get("num_embodiments", 16)
    cfg["max_state_dim"] = new_state
    cfg["max_action_dim"] = new_action
    cfg["num_embodiments"] = new_emb
    # input/output_features carry their own copy of the dims. Leaving them stale
    # makes the checkpoint load with a 32-wide expert but an 8-wide declared
    # action, so anything reading the feature spec silently truncates.
    if "action" in cfg.get("output_features", {}):
        cfg["output_features"]["action"]["shape"] = [new_action]
    if "observation.state" in cfg.get("input_features", {}):
        cfg["input_features"]["observation.state"]["shape"] = [new_state]
    cfg_path.write_text(json.dumps(cfg, indent=2))
    print(f"config: state {old_state}->{new_state}  action {old_action}->{new_action}  "
          f"embodiments {old_emb}->{new_emb}")

    weights = src / "model.safetensors"
    sd = load_file(str(weights))
    grown = []
    for k, v in list(sd.items()):
        target = None
        for suffix, axis in GROW.items():
            if k.endswith(suffix):
                want = new_state if "state_proj" in k else new_action
                if v.shape[axis] < want:
                    target = (axis, want)
                break
        # embodiment table: a 2-D embedding whose row count equals num_embodiments
        if target is None and v.dim() == 2 and v.shape[0] == old_emb and "emb" in k.lower():
            target = (0, new_emb)
        if target is None:
            continue
        axis, want = target
        shape = list(v.shape)
        pad = want - shape[axis]
        if pad <= 0:
            continue
        shape[axis] = pad
        sd[k] = torch.cat([v, v.new_zeros(*shape)], dim=axis)
        grown.append(f"  {k}: {tuple(v.shape)} -> {tuple(sd[k].shape)}")

    save_file(sd, str(dst / "model.safetensors"), metadata={"format": "pt"})
    print("grown tensors:" if grown else "nothing to grow (already wide enough)")
    print("\n".join(grown))
    print(f"wrote {dst}")


if __name__ == "__main__":
    main()