"""Assemble Thor-deployable ckpt dir from a finished VQ-BeT yor run.

Output <run>/ckpt/:
    policy.pt    whole-module BehaviorTransformer pickle, CPU, eval mode
                 (VQ-BeT: no normalizer object; actions raw, act_scale=1.0)
    encoder.pt   whole-module DinoV2Encoder pickle, CPU, eval mode
                 (B,C,H,W)/255 RGB -> (B,P,E)
    card.yaml    views, img_size, obs_horizon, chunk, exec_n, action_hz,
                 act_fields (58, exact converter order), proprio_fields

action_hz = measured training rate = 1 / median(STRIDE * tick_dt),
computed from the raw episodes' t0_ns column.
"""
import argparse
import csv
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, "/home_shared/grail_charles/patch_policy")

EE_L = [f"left_ee_target_wxyz_xyz_{i}" for i in range(7)]
EE_R = [f"right_ee_target_wxyz_xyz_{i}" for i in range(7)]
HAND_L = [f"left_hand_qpos_{i}" for i in range(20)]
HAND_R = [f"right_hand_qpos_{i}" for i in range(20)]
ACTION_COLS = EE_L + EE_R + HAND_L + HAND_R + ["lift"] + \
    ["base_pose_target_0", "base_pose_target_1", "base_pose_target_2"]
assert len(ACTION_COLS) == 58


def measured_tick_dt(raw_root: Path) -> float:
    dts = []
    for ep in sorted(raw_root.glob("ep_*")):
        with open(ep / "robot_data.csv", newline="") as f:
            t0 = np.array([int(r["t0_ns"]) for r in csv.DictReader(f)], np.int64)
        d = np.diff(t0) / 1e9
        dts.append(d[(d > 0) & (d < 1.0)])
    return float(np.median(np.concatenate(dts)))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True, help="run dir with model_final.pt")
    ap.add_argument("--model", default="model_final.pt",
                    help="ckpt file name in the run dir (e.g. model_stopped.pt)")
    ap.add_argument("--raw-root",
                    default=str(Path.home() / "data/yor_episodes/5c829ca0"))
    ap.add_argument("--stride", type=int, default=3)
    ap.add_argument("--chunk", type=int, default=5)
    ap.add_argument("--exec-n", type=int, default=4)
    ap.add_argument("--feature-key", default="x_norm_patchtokens",
                    help="x_norm_patchtokens (P=256) or x_norm_clstoken (P=1)")
    args = ap.parse_args()

    import torch
    run = Path(args.run)
    out = run / "ckpt"
    out.mkdir(parents=True, exist_ok=True)

    tick_dt = measured_tick_dt(Path(args.raw_root))
    action_hz = 1.0 / (args.stride * tick_dt)
    print(f"tick_dt={tick_dt*1000:.1f}ms stride={args.stride} "
          f"-> action_hz={action_hz:.3f}")

    print("loading policy ...", flush=True)
    policy = torch.load(run / args.model, map_location="cpu",
                        weights_only=False)
    policy.to("cpu")
    policy.eval()  # statement: bet.train() returns None, do not chain
    torch.save(policy, out / "policy.pt")

    from models.encoder.dino import DinoV2Encoder
    print("building encoder ...", flush=True)
    n_patches = 1 if args.feature_key == "x_norm_clstoken" else 256
    enc = DinoV2Encoder(name="dinov2_vits14",
                        feature_key=args.feature_key,
                        output_dim=384, n_patches=n_patches)
    enc.to("cpu")
    enc.eval()
    for p in enc.parameters():
        p.requires_grad = False
    torch.save(enc, out / "encoder.pt")

    card = [
        "views: [odin, wl, wr]",
        "img_size: 224",
        "obs_horizon: 2",
        f"chunk: {args.chunk}",
        f"exec_n: {args.exec_n}",
        f"action_hz: {action_hz:.3f}",
        "act_fields: [" + ", ".join(ACTION_COLS) + "]",
        "proprio_fields: []",
    ]
    (out / "card.yaml").write_text("\n".join(card) + "\n")

    # verify: reload + forward with zeros
    print("verifying ...", flush=True)
    p2 = torch.load(out / "policy.pt", map_location="cpu", weights_only=False)
    p2.eval()
    e2 = torch.load(out / "encoder.pt", map_location="cpu", weights_only=False)
    e2.eval()
    with torch.no_grad():
        assert e2(torch.rand(2, 3, 224, 224)).shape == (2, n_patches, 384)
        pr, _, _ = p2(torch.zeros(1, 2, 768, 384), torch.ones(1, 2, 1, 1), None)
        assert pr.shape == (1, 2, args.chunk, 58), pr.shape
    print(f"OK -> {out}: policy{p2.__class__.__name__} "
          f"pred {tuple(pr.shape)}")


if __name__ == "__main__":
    main()