#!/usr/bin/env python
"""Convert the id-conditioned GR00T checkpoint to morphology conditioning.

Why this run exists: the physical_ai_ft checkpoint is `conditioning: id`, and a
learned id is a lookup table — it has no row for a robot that was never trained on,
and none can be synthesised. Measured on held-out LeKiwi: every trained id scores
1.18-1.33 on the absolute dims (worse than holding still) and a PURE RANDOM identity
vector does just as well. Simply switching the config to `morph` and grafting
tv2_C_scaled's morph encoder does not help either (1.30/1.11) — the expert spent 60k
steps learning to read an id token and has never seen morph tokens.

So the expert has to be taught to read the descriptor. Everything else is already
trained, so this only has to re-route one conditioning slot: 16 physical numbers ->
MLP -> 2 tokens, in place of the 1 id token.

Data is streamed from the hub exactly as the original run did (no disk), with the
six GR00T family descriptors from configs/groot_morphology.yaml.

The payoff is the project's central claim, measured rather than argued: if a
descriptor-conditioned model transfers to LeKiwi where the id-conditioned one cannot,
that is the difference between a code you can look up and a code you can compute.
"""

from __future__ import annotations

import argparse
import time
from pathlib import Path

import torch
import yaml


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--resume", default="/home/alexw/tinyvla/outputs/physical_ai_ft_morph")
    ap.add_argument("--out", default="/home/alexw/tinyvla/outputs/groot_morph")
    ap.add_argument("--spec", default="/home/alexw/tinyvla/configs/physical_ai_stream.yaml")
    ap.add_argument("--descriptors", default="/home/alexw/tinyvla/configs/groot_morphology.yaml")
    ap.add_argument("--steps", type=int, default=6000)
    ap.add_argument("--batch-size", type=int, default=32)
    ap.add_argument("--num-workers", type=int, default=6)
    ap.add_argument("--lr", type=float, default=1.0e-4)
    ap.add_argument("--backbone-lr-mult", type=float, default=0.1)
    ap.add_argument("--warmup", type=int, default=200)
    ap.add_argument("--save-freq", type=int, default=1000)
    ap.add_argument("--log-freq", type=int, default=50)
    args = ap.parse_args()

    from safetensors.torch import load_file
    from transformers import AutoTokenizer

    from tinyvla.data.hub_stream import HubEpisodeStream
    from tinyvla.modeling_tinyvla import TinyVLAPolicy
    from tinyvla.modules.embodiment import MORPH_FIELDS

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    pol = TinyVLAPolicy.from_pretrained(args.resume)
    pol.load_state_dict(load_file(f"{args.resume}/model.safetensors"), strict=True)
    cfg = pol.config
    assert cfg.conditioning == "morph", f"expected morph conditioning, got {cfg.conditioning}"
    pol = pol.cuda().train()
    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)
    print(f"resumed {args.resume} | conditioning={cfg.conditioning} "
          f"max_state={cfg.max_state_dim} max_action={cfg.max_action_dim}", flush=True)

    # same field scaling the rest of the project uses (see MORPH_FIELDS in embodiment.py)
    sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3,
          "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416,
          "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2}
    raw = yaml.safe_load(Path(args.descriptors).read_text())
    morph = {k: torch.tensor([v.get(f, 0) * sc.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32)
             for k, v in raw.items()}
    print("descriptors:", list(morph))

    specs = yaml.safe_load(Path(args.spec).read_text())["datasets"]
    missing = {s.get("morph_key") for s in specs} - set(morph)
    assert not missing, f"specs reference descriptors that do not exist: {missing}"

    # Fix the sampler bug the original run shipped with: HubEpisodeStream draws a
    # DATASET PER EPISODE, so a weight of sqrt(frames) is silently multiplied by
    # episode length. r1_pro episodes are ~37x longer, which turned an intended 49%
    # share into a realized 96.5% — and here it also starves the loader, because
    # every draw pulls a ~190 MB two-camera episode and the first batch never
    # arrives. Dividing by relative episode length restores the intended mixture.
    EP_LEN = {"r1_pro": 37.0, "gr1": 1.0, "panda_single": 1.0,
              "panda_bi_gripper": 1.0, "panda_bi_hand": 1.0, "g1": 0.6}
    for sp in specs:
        sp["weight"] = float(sp["weight"]) / EP_LEN.get(sp.get("morph_key"), 1.0)
    from collections import defaultdict as _dd
    _agg = _dd(float)
    for sp in specs:
        _agg[sp["morph_key"]] += sp["weight"]
    _tot = sum(_agg.values())
    print("reweighted family shares:", {k: round(v / _tot, 3) for k, v in sorted(_agg.items())})

    stream = HubEpisodeStream(specs, chunk=cfg.chunk_size, image_size=cfg.image_size,
                              max_state_dim=cfg.max_state_dim, max_action_dim=cfg.max_action_dim,
                              shuffle_buffer=64, seed=1234, morph_descriptors=morph)
    loader = torch.utils.data.DataLoader(
        stream, batch_size=args.batch_size, num_workers=args.num_workers,
        pin_memory=True, prefetch_factor=4, persistent_workers=True,
        multiprocessing_context="spawn",  # fork deadlocks the hub HTTP pool (B200.md)
    )

    backbone, head = [], []
    for n, p in pol.named_parameters():
        if not p.requires_grad:
            continue
        (backbone if n.startswith("semantic.") else head).append(p)
    opt = torch.optim.AdamW([
        {"params": head, "lr": args.lr},
        {"params": backbone, "lr": args.lr * args.backbone_lr_mult},
    ], weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / args.warmup))
    scaler_dtype = torch.bfloat16
    print(f"head {sum(p.numel() for p in head)/1e6:.1f}M | "
          f"backbone {sum(p.numel() for p in backbone)/1e6:.1f}M at {args.backbone_lr_mult}x", flush=True)

    step, t0 = 0, time.time()
    for batch in loader:
        t = tok(list(batch["task"]), padding="max_length", truncation=True,
                max_length=48, return_tensors="pt")
        b = {k: v.cuda(non_blocking=True) for k, v in batch.items()
             if isinstance(v, torch.Tensor)}
        b["observation.language.tokens"] = t["input_ids"].cuda()
        b["observation.language.attention_mask"] = t["attention_mask"].bool().cuda()
        with torch.autocast("cuda", scaler_dtype):
            loss, info = pol.forward(b)
        loss.backward()
        torch.nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]], 10.0)
        opt.step()
        sched.step()
        opt.zero_grad(set_to_none=True)
        step += 1
        if step % args.log_freq == 0:
            print(f"step {step}/{args.steps} loss {info['loss']:.4f} "
                  f"{step/(time.time()-t0):.2f} it/s", flush=True)
        if step % args.save_freq == 0 or step == args.steps:
            d = out / (f"step_{step}" if step < args.steps else "final")
            d.mkdir(parents=True, exist_ok=True)
            pol.save_pretrained(str(d))
            print(f"saved {d}", flush=True)
        if step >= args.steps:
            break
    print("DONE", flush=True)


if __name__ == "__main__":
    main()