"""Scene-only DinoVolumeScene training loop — single arm, no wrist.

Per-step losses: volume CE (pixel + height) + gripper CE + K-means rotation CE.
Per-vis: PCA of refined features, GT-vs-pred keypoints overlay, per-T heatmap
grid, and per-(T, bin) heatmap grids for rotation / gripper / height.

Checkpoint payload (everything inference needs to unpack predictions):
  - ``model``         state_dict
  - ``opt``           state_dict
  - ``step``          int
  - ``args``          dict (CLI args)
  - ``dataset_cfg``   includes rot_centroids (K, 4) + bin-range configs

Run on puget (GPU)::

    ssh dev
    source ~/yam_para/.venv/bin/activate
    cd ~/lab/para/robot/yam
    CUDA_VISIBLE_DEVICES=0 nohup python train.py --name <run> \\
        --data_root ~/yam_para/data/<dataset> > /tmp/train.log 2>&1 &
"""
from __future__ import annotations

import argparse
import math
import sys
import time
from pathlib import Path

import cv2
import numpy as np
import torch
import wandb
from torch.utils.data import DataLoader

sys.path.insert(0, str(Path(__file__).resolve().parent))
from lib.dataset import YamScene
from lib.model import DinoVolumeScene, DinoVolumeScenePerVoxel, all_losses
from lib.viz.bin_grids import bin_grid_panel
from lib.viz.feature_pca import feature_pca_panel, shared_feature_pca_panels
from lib.viz.heatmap import marginal_heatmap_grid
from lib.viz.height_curve import height_curve_panel
from lib.viz.keypoints import keypoints_overlay


def _project_world_to_pix(xyz_w: np.ndarray, K: np.ndarray, T_w2c: np.ndarray,
                            img_size: int):
    """Project ``xyz_w`` (T, 3) through one cam → image-size pixel coords (T, 2).
    Caller should clip + treat as drawing coords; in-frustum mask is implicit
    in clipped values. Assumes K is already scaled to img_size×img_size."""
    ones = np.ones((xyz_w.shape[0], 1), dtype=xyz_w.dtype)
    xyz_h = np.concatenate([xyz_w, ones], axis=-1)              # (T, 4)
    cam = (T_w2c @ xyz_h.T).T[:, :3]                            # (T, 3)
    z = np.clip(cam[:, 2:3], 1e-3, None)
    uvw = (K @ cam.T).T                                         # (T, 3)
    return uvw[:, :2] / z


def _viz_batch_multiview(batch, out, img_size: int, vis_n: int):
    """Multi-view panel set: pred vs GT keypoints overlaid on EACH view's
    image side-by-side, plus per-view PCA, plus the existing grip/rot
    grids. Voxel-position lookup gives us the predicted world XYZ
    DIRECTLY (no ray-cast); we then project that to each camera."""
    panels = {}
    views = list(out["views"])
    vol = out["volume_logits"].detach().cpu().numpy()           # (B, T, NZ, P, P)
    voxel_pos = out["voxel_positions"].detach().cpu().numpy()    # (B, NZ, P, P, 3)
    grip = out["grip_logits"].detach().cpu().numpy()
    rot = out["rot_logits"].detach().cpu().numpy()
    view_feat_maps = out["view_feat_maps"]                       # list of (B, F, P, P)
    view_feats_np = [v.detach().cpu().numpy() for v in view_feat_maps]

    bgrs_scene = batch["_bgr_src"]
    bgrs_wrist = batch.get("_wrist_bgr_src")
    target_xyz = batch["target_xyz"].cpu().numpy()              # (B, T, 3)
    target_grip = batch["target_grip"].cpu().numpy()
    target_rot = batch["target_rot"].cpu().numpy()
    K_scene = batch["_K_in"].cpu().numpy()
    T_w2c_scene = batch["_T_w2c"].cpu().numpy()
    K_wrist_b = batch["_K_wrist"].cpu().numpy() if "_K_wrist" in batch else None
    T_w2c_wrist_b = batch["_T_w2c_wrist"].cpu().numpy() if "_T_w2c_wrist" in batch else None

    B, T, NZ, P, _ = vol.shape

    for b in range(min(vis_n, B)):
        # Per-view backgrounds, sized to img_size (square — matches K).
        view_bgrs = []
        view_Ks = []
        view_Ts = []
        for v in views:
            if v == "scene":
                view_bgrs.append(cv2.resize(bgrs_scene[b], (img_size, img_size),
                                              interpolation=cv2.INTER_AREA))
                view_Ks.append(K_scene[b])
                view_Ts.append(T_w2c_scene[b])
            elif v == "wrist":
                wbgr = bgrs_wrist[b] if bgrs_wrist is not None else None
                if wbgr is None or K_wrist_b is None:
                    continue
                view_bgrs.append(cv2.resize(wbgr, (img_size, img_size),
                                              interpolation=cv2.INTER_AREA))
                view_Ks.append(K_wrist_b[b])
                view_Ts.append(T_w2c_wrist_b[b])
        if not view_bgrs:
            continue

        # Predicted world XYZ per timestep — argmax over (NZ, P, P), look up voxel_positions.
        flat = vol[b].reshape(T, -1)
        idx_joint = flat.argmax(axis=1)
        nz_idx = idx_joint // (P * P)
        hw = idx_joint % (P * P)
        y_idx = hw // P; x_idx = hw % P
        pred_xyz_world = voxel_pos[b, nz_idx, y_idx, x_idx, :]   # (T, 3)
        gt_xyz_world = target_xyz[b]                              # (T, 3)

        # Build the keypoint-overlay panel per view, then concat horizontally.
        kp_strips = []
        feats_per_view = []
        for k, v in enumerate(views[:len(view_bgrs)]):
            bg = view_bgrs[k]
            gt_pix = _project_world_to_pix(gt_xyz_world, view_Ks[k], view_Ts[k], img_size)
            pred_pix = _project_world_to_pix(pred_xyz_world, view_Ks[k], view_Ts[k], img_size)
            kp = keypoints_overlay(bg.copy(), gt_pix,
                                     filled=False, draw_line=False,
                                     color_override=(255, 255, 255), radius_px=10)
            kp = keypoints_overlay(kp, pred_pix, draw_line=True, radius_px=6)
            # Title strip on top.
            bar = np.full((24, kp.shape[1], 3), 30, dtype=np.uint8)
            cv2.putText(bar, f"{v}", (8, 17), cv2.FONT_HERSHEY_SIMPLEX,
                         0.55, (240, 240, 240), 1, cv2.LINE_AA)
            kp_strips.append(np.vstack([bar, kp]))
            # Stash (P, P, F) feature map for joint PCA below.
            feats_per_view.append(view_feats_np[k][b].transpose(1, 2, 0))

        # Joint PCA: fit one basis on ALL views' pixel features stacked,
        # then paint each view with the SAME basis so colours are comparable.
        pca_imgs = shared_feature_pca_panels(
            feats_per_view, out_wh=(img_size, img_size))
        pca_strips = []
        for k, (v, pca) in enumerate(zip(views[:len(view_bgrs)], pca_imgs)):
            bar2 = np.full((24, pca.shape[1], 3), 30, dtype=np.uint8)
            cv2.putText(bar2, f"{v} PCA (joint)", (8, 17), cv2.FONT_HERSHEY_SIMPLEX,
                         0.55, (240, 240, 240), 1, cv2.LINE_AA)
            pca_strips.append(np.vstack([bar2, pca]))

        kp_concat = np.hstack(kp_strips)
        pca_concat = np.hstack(pca_strips)
        panels[f"viz/kp_gt_vs_pred_{b}"] = wandb.Image(kp_concat[:, :, ::-1])
        panels[f"viz/pca_dino_mlp_{b}"] = wandb.Image(pca_concat[:, :, ::-1])

        # Grip + rot bin grids (same as single-view).
        panels[f"viz/grip_grid_{b}"] = wandb.Image(
            bin_grid_panel(grip[b], target=target_grip[b])[:, :, ::-1])
        panels[f"viz/rot_grid_{b}"] = wandb.Image(
            bin_grid_panel(rot[b], target=target_rot[b])[:, :, ::-1])
    return panels


def _viz_batch_v3(batch, out, img_size: int, vis_n: int):
    """v3 panel set: per-view marginal heatmap (t=0), per-view PCA, pred
    vs GT keypoints projected to each view, grip/rot bin grids, and a
    per-timestep view-dominance bar.
    """
    from lib.model_volumetric_v3 import (v3_per_view_heatmaps,
                                            v3_view_dominance,
                                            v3_z_marginal)
    panels = {}
    views = list(out["views"])
    # Use the argmax-based pred for viz (chosen_xyz is GT-nearest in train).
    chosen_xyz = out["pred_xyz_world"].detach().cpu().numpy()      # (B, T, 3)
    grip = out["grip_logits"].detach().cpu().numpy()
    rot = out["rot_logits"].detach().cpu().numpy()
    view_feats_np = [v.detach().cpu().numpy() for v in out["view_feat_maps"]]

    bgrs_scene = batch["_bgr_src"]
    bgrs_wrist = batch.get("_wrist_bgr_src")
    target_xyz = batch["target_xyz"].cpu().numpy()
    target_grip = batch["target_grip"].cpu().numpy()
    target_rot = batch["target_rot"].cpu().numpy()
    K_scene = batch["_K_in"].cpu().numpy()
    T_w2c_scene = batch["_T_w2c"].cpu().numpy()
    K_wrist_b = batch["_K_wrist"].cpu().numpy() if "_K_wrist" in batch else None
    T_w2c_wrist_b = batch["_T_w2c_wrist"].cpu().numpy() if "_T_w2c_wrist" in batch else None

    B = chosen_xyz.shape[0]
    T_ = chosen_xyz.shape[1]

    # Per-view (B, P, P) log-marginals at t=0.
    log_p_per_view_t0 = [hp.detach().cpu().numpy()
                          for hp in v3_per_view_heatmaps(out, t_index=0)]
    # Per-(B, N, T) view dominance log-marginal.
    dom = v3_view_dominance(out).detach().cpu().numpy()             # (B, N, T)

    for b in range(min(vis_n, B)):
        view_bgrs = []
        view_Ks = []
        view_Ts = []
        for v in views:
            if v == "scene":
                view_bgrs.append(cv2.resize(bgrs_scene[b], (img_size, img_size),
                                              interpolation=cv2.INTER_AREA))
                view_Ks.append(K_scene[b]); view_Ts.append(T_w2c_scene[b])
            elif v == "wrist":
                wbgr = bgrs_wrist[b] if bgrs_wrist is not None else None
                if wbgr is None or K_wrist_b is None:
                    continue
                view_bgrs.append(cv2.resize(wbgr, (img_size, img_size),
                                              interpolation=cv2.INTER_AREA))
                view_Ks.append(K_wrist_b[b]); view_Ts.append(T_w2c_wrist_b[b])
        if not view_bgrs:
            continue

        pred_xyz_world = chosen_xyz[b]                              # (T, 3)
        gt_xyz_world = target_xyz[b]                                # (T, 3)

        kp_strips = []
        feats_per_view = []
        heatmap_strips = []
        for k, v in enumerate(views[:len(view_bgrs)]):
            bg = view_bgrs[k]
            gt_pix = _project_world_to_pix(gt_xyz_world, view_Ks[k], view_Ts[k], img_size)
            pred_pix = _project_world_to_pix(pred_xyz_world, view_Ks[k], view_Ts[k], img_size)
            kp = keypoints_overlay(bg.copy(), gt_pix,
                                     filled=False, draw_line=False,
                                     color_override=(255, 255, 255), radius_px=10)
            kp = keypoints_overlay(kp, pred_pix, draw_line=True, radius_px=6)
            bar = np.full((24, kp.shape[1], 3), 30, dtype=np.uint8)
            cv2.putText(bar, f"{v}", (8, 17), cv2.FONT_HERSHEY_SIMPLEX,
                         0.55, (240, 240, 240), 1, cv2.LINE_AA)
            kp_strips.append(np.vstack([bar, kp]))
            feats_per_view.append(view_feats_np[k][b].transpose(1, 2, 0))

            # Per-view heatmap at t=0 — exp the log-prob, normalize for viz.
            hp = log_p_per_view_t0[k][b]                            # (P, P)
            hp_lin = np.exp(hp - hp.max())
            hp_lin /= max(hp_lin.max(), 1e-9)
            hp_img = (hp_lin * 255.0).astype(np.uint8)
            hp_color = cv2.applyColorMap(hp_img, cv2.COLORMAP_INFERNO)
            hp_color = cv2.resize(hp_color, (img_size, img_size),
                                    interpolation=cv2.INTER_LINEAR)
            blend = cv2.addWeighted(bg, 0.4, hp_color, 0.6, 0)
            bar2 = np.full((24, blend.shape[1], 3), 30, dtype=np.uint8)
            cv2.putText(bar2, f"{v} heatmap (t=0)", (8, 17),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.55,
                         (240, 240, 240), 1, cv2.LINE_AA)
            heatmap_strips.append(np.vstack([bar2, blend]))

        pca_imgs = shared_feature_pca_panels(
            feats_per_view, out_wh=(img_size, img_size))
        pca_strips = []
        for k, (v, pca) in enumerate(zip(views[:len(view_bgrs)], pca_imgs)):
            bar3 = np.full((24, pca.shape[1], 3), 30, dtype=np.uint8)
            cv2.putText(bar3, f"{v} PCA (joint)", (8, 17),
                         cv2.FONT_HERSHEY_SIMPLEX, 0.55,
                         (240, 240, 240), 1, cv2.LINE_AA)
            pca_strips.append(np.vstack([bar3, pca]))

        panels[f"viz/kp_gt_vs_pred_{b}"] = wandb.Image(np.hstack(kp_strips)[:, :, ::-1])
        panels[f"viz/pca_dino_mlp_{b}"] = wandb.Image(np.hstack(pca_strips)[:, :, ::-1])
        panels[f"viz/heatmap_per_view_t0_{b}"] = wandb.Image(
            np.hstack(heatmap_strips)[:, :, ::-1])

        # Dominance bar: each timestep gets a stacked column showing
        # softmax over views. Stack views top-to-bottom by their score.
        dom_t = dom[b]                                              # (N, T)
        # Softmax over N.
        dom_t -= dom_t.max(axis=0, keepdims=True)
        dom_p = np.exp(dom_t)
        dom_p /= np.maximum(dom_p.sum(axis=0, keepdims=True), 1e-9)
        # Render: H=140, W=24*T_, each col a stacked bar (scene cyan, wrist magenta).
        col_w = 24; H = 140
        bar_w = col_w * T_
        bar_img = np.full((H, bar_w, 3), 30, dtype=np.uint8)
        palette = [(220, 200, 60), (200, 80, 220), (80, 220, 120), (220, 120, 80)]
        for t in range(T_):
            x0 = t * col_w + 2; x1 = (t + 1) * col_w - 2
            y_cur = H
            for k in range(dom_p.shape[0]):
                h_k = int(round(dom_p[k, t] * (H - 4)))
                if h_k < 1:
                    continue
                col = palette[k % len(palette)]
                cv2.rectangle(bar_img, (x0, y_cur - h_k), (x1, y_cur),
                               col, thickness=-1)
                y_cur -= h_k
        # Legend strip on top.
        leg = np.full((26, bar_w, 3), 30, dtype=np.uint8)
        x = 8
        for k, v in enumerate(views[:dom_p.shape[0]]):
            col = palette[k % len(palette)]
            cv2.rectangle(leg, (x, 6), (x + 14, 20), col, thickness=-1)
            cv2.putText(leg, v, (x + 18, 18), cv2.FONT_HERSHEY_SIMPLEX,
                         0.45, (240, 240, 240), 1, cv2.LINE_AA)
            x += 18 + 12 + 8 * len(v)
        panels[f"viz/view_dominance_{b}"] = wandb.Image(
            np.vstack([leg, bar_img])[:, :, ::-1])

        panels[f"viz/grip_grid_{b}"] = wandb.Image(
            bin_grid_panel(grip[b], target=target_grip[b])[:, :, ::-1])
        panels[f"viz/rot_grid_{b}"] = wandb.Image(
            bin_grid_panel(rot[b], target=target_rot[b])[:, :, ::-1])

    # Per-t marginal over Z bins — compute one item at a time to bound
    # the peak (B, P, P, Z) intermediate. Only slice batched tensors;
    # PEs and scalars are shared across items.
    _no_batch = {"pe_z", "pe_t", "n_z", "pred_size", "n_views", "views"}
    target_z = batch["target_z"].cpu().numpy()                      # (B, T)
    for b in range(min(vis_n, B)):
        out_b = {}
        for k, v in out.items():
            if k in _no_batch:
                out_b[k] = v
            elif torch.is_tensor(v):
                out_b[k] = v[b:b + 1]
            elif isinstance(v, list) and v and torch.is_tensor(v[0]):
                out_b[k] = [x[b:b + 1] for x in v]
            else:
                out_b[k] = v
        z_log = v3_z_marginal(out_b)[0].detach().cpu().numpy()      # (Z, T)
        panels[f"viz/z_grid_{b}"] = wandb.Image(
            bin_grid_panel(z_log.T, target=target_z[b])[:, :, ::-1])
    return panels


def _viz_batch(batch, out, img_size: int, vis_n: int = 1,
                 per_voxel: bool = False, z_lo: float = 0.0, z_hi: float = 0.4):
    """Per-batch wandb panels. ``per_voxel=True`` swaps the marginal-Z bin
    grid for a pred-vs-GT height-over-time line plot (the per-voxel model
    doesn't have a separate height head).

    Multi-view models (``out["views"]`` set) take a different path that
    renders pred-vs-GT keypoints + PCA on EACH view's image, concatenated
    horizontally. Required batch fields for multi-view: ``_bgr_src``,
    ``_wrist_bgr_src``, ``_K_in``, ``_T_w2c``, ``_K_wrist``,
    ``_T_w2c_wrist``, ``target_xyz``.
    """
    # v3 detection: out lacks ``volume_logits`` and carries ``view_Fz``.
    if "view_Fz" in out and "volume_logits" not in out:
        return _viz_batch_v3(batch, out, img_size, vis_n)
    if isinstance(out.get("views"), (list, tuple)) and len(out["views"]) >= 1:
        return _viz_batch_multiview(batch, out, img_size, vis_n)
    panels = {}
    vol = out["volume_logits"].detach().cpu().numpy()           # (B, T, Z, P, P)
    grip = out["grip_logits"].detach().cpu().numpy()             # (B, T, n_grip)
    rot = out["rot_logits"].detach().cpu().numpy()               # (B, T, n_rot)
    feats_ref = out["F_refined"].detach().cpu().numpy()          # (B, C[*Z], P, P)
    bgrs = batch["_bgr_src"]
    target_yx = batch["target_yx"].cpu().numpy()                 # (B, T, 2)
    target_z = batch["target_z"].cpu().numpy()                   # (B, T)
    target_grip = batch["target_grip"].cpu().numpy()             # (B, T)
    target_rot = batch["target_rot"].cpu().numpy()               # (B, T)
    B, T, Z, P, _ = vol.shape
    scale = img_size / P

    for b in range(min(vis_n, B)):
        bg = cv2.resize(bgrs[b], (img_size, img_size), interpolation=cv2.INTER_AREA)
        # PCA over refined features (C or C*Z channels — flatten works either way).
        pca = feature_pca_panel(feats_ref[b].transpose(1, 2, 0))
        panels[f"viz/pca_dino_mlp_{b}"] = wandb.Image(pca[:, :, ::-1])

        # Joint argmax over (Z, P, P) — the cell of the most-likely voxel.
        # GT is white open circles; pred is rainbow filled + polyline.
        gt_pix = (target_yx[b][:, ::-1].astype(np.float32)) * scale
        flat = vol[b].reshape(T, -1)                              # (T, Z*P*P)
        idx_joint = flat.argmax(axis=1)
        pred_z_bin = idx_joint // (P * P)
        hw = idx_joint % (P * P)
        pred_pix = np.stack([(hw % P).astype(np.float32) * scale,
                              (hw // P).astype(np.float32) * scale], axis=1)
        kp = keypoints_overlay(bg.copy(), gt_pix,
                                filled=False, draw_line=False,
                                color_override=(255, 255, 255), radius_px=10)
        kp = keypoints_overlay(kp, pred_pix, draw_line=True, radius_px=6)
        panels[f"viz/kp_gt_vs_pred_{b}"] = wandb.Image(kp[:, :, ::-1])

        # Per-T heatmap grid (one cell per timestep, RGB underneath).
        grid = marginal_heatmap_grid(bg, vol[b], cell_size=P, alpha=0.5)
        panels[f"viz/heatmap_grid_{b}"] = wandb.Image(grid[:, :, ::-1])

        if per_voxel:
            # Per-voxel model: no separate height head — log pred-vs-GT
            # height line plot from the volume's joint argmax z-bin.
            hc = height_curve_panel(pred_z_bins=pred_z_bin,
                                     gt_z_bins=target_z[b],
                                     z_lo=z_lo, z_hi=z_hi, n_z=Z,
                                     out_wh=(512, 256))
            panels[f"viz/height_curve_{b}"] = wandb.Image(hc[:, :, ::-1])
        else:
            # Factorized model: per-T height-bin marginal (T x Z grid).
            flat_norm = flat - flat.max(axis=1, keepdims=True)
            e = np.exp(flat_norm); e /= e.sum(axis=1, keepdims=True)
            probs_z = e.reshape(T, Z, P, P).sum(axis=(2, 3))
            panels[f"viz/z_grid_{b}"] = wandb.Image(
                bin_grid_panel(np.log(probs_z + 1e-12), target=target_z[b])[:, :, ::-1])

        # Per-T gripper-bin grid.
        panels[f"viz/grip_grid_{b}"] = wandb.Image(
            bin_grid_panel(grip[b], target=target_grip[b])[:, :, ::-1])
        # Per-T rotation-bin grid (K=64 cells wide).
        panels[f"viz/rot_grid_{b}"] = wandb.Image(
            bin_grid_panel(rot[b], target=target_rot[b])[:, :, ::-1])

    return panels


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--name", required=True)
    p.add_argument("--data_root", required=True)
    p.add_argument("--ckpt_root", default=str(Path.home() / "yam_para/checkpoints"))
    p.add_argument("--wandb_project", default="yam_train_simple")
    p.add_argument("--batch_size", type=int, default=40)
    p.add_argument("--num_workers", type=int, default=8,
                   help="Default bumped 4→8 since training is CPU-bound on "
                        "JPEG decode at batch 40.")
    p.add_argument("--lr", type=float, default=3e-4,
                   help="Head LR (and backbone LR if --backbone_lr_mult==1).")
    p.add_argument("--lr_min", type=float, default=None,
                   help="Terminal head LR for cosine schedule. Default: same as --lr "
                        "(constant LR). Set to e.g. 1e-5 with --lr_schedule cosine.")
    p.add_argument("--backbone_lr_mult", type=float, default=1.0,
                   help="Backbone LR = --lr * this. <1.0 protects pretrained DINO "
                        "weights from being damaged by the fresh-head LR. The "
                        "ablation experiment showed lr=2e-5 was the magic value "
                        "for from-scratch — i.e. 0.1× of head-friendly 2e-4. "
                        "Backbone min LR scales by the same factor.")
    p.add_argument("--lr_warmup", type=int, default=0,
                   help="Linear-warmup iters for both groups (0 = no warmup). "
                        "Recommended ~5%% of --max_iters for from-scratch runs.")
    p.add_argument("--lr_schedule", default="constant",
                   choices=["constant", "cosine"],
                   help="LR schedule after warmup. cosine = decay from lr → lr_min.")
    p.add_argument("--n_window", type=int, default=16)
    p.add_argument("--frame_stride", type=int, default=8)
    p.add_argument("--n_height_bins", type=int, default=32)
    p.add_argument("--n_gripper_bins", type=int, default=32)
    p.add_argument("--n_rot_clusters", type=int, default=256)
    p.add_argument("--pred_size", type=int, default=128)
    p.add_argument("--img_size", type=int, default=405)
    p.add_argument("--feat_dim", type=int, default=48)
    p.add_argument("--d_model", type=int, default=256)
    p.add_argument("--n_blocks", type=int, default=5)
    p.add_argument("--mlp_hidden", type=int, default=128)
    p.add_argument("--z_lo", type=float, default=None,
                   help="Lower height bin edge (m). Default = auto-derive from "
                        "dataset (min EE z − pad). Finetunes override with "
                        "the source ckpt's range automatically.")
    p.add_argument("--z_hi", type=float, default=None,
                   help="Upper height bin edge (m). Default = auto-derive from "
                        "dataset (max EE z + pad).")
    p.add_argument("--grip_lo", type=float, default=-0.2)
    p.add_argument("--grip_hi", type=float, default=0.8)
    p.add_argument("--max_iters", type=int, default=8000,
                   help="Default sized so a fresh run finishes in ~70 min "
                        "at 2 it/s (avoids overfitting on small datasets).")
    p.add_argument("--ckpt_every_iters", type=int, default=100)
    p.add_argument("--vis_every_iters", type=int, default=50)
    p.add_argument("--device", default="cuda")
    p.add_argument("--init_from", default=None,
                   help="Path to a checkpoint to initialize model weights from "
                        "(finetune mode). Also reuses the source ckpt's K-means "
                        "rot centroids so the rot_head stays calibrated.")
    p.add_argument("--per_voxel", action="store_true",
                   help="Use DinoVolumeScenePerVoxel instead of DinoVolumeScene "
                        "(per-voxel features, no factorized scoring, CLS-only "
                        "input). Saves model_kind='dino_volume_per_voxel' in ckpt.")
    p.add_argument("--da3", action="store_true",
                   help="Use DA3VolumeScene (Depth-Anything-3 backbone + camera "
                        "token + parallel feat DPT branch deep-copied from depth). "
                        "Same factored YX×Z×T scoring as DinoVolumeScene. "
                        "Defaults swap to n_height_bins=128, feat_dim=32. "
                        "Saves model_kind='da3_volume_scene' in ckpt. "
                        "Mutually exclusive with --per_voxel.")
    p.add_argument("--hires_factorized", action="store_true",
                   help="Use DinoVolumeSceneHires — DINOv3 backbone + bilinear "
                        "up + 5-layer CNN at pred_size=256. Factorized scoring. "
                        "Defaults: n_height_bins=128, feat_dim=48. "
                        "Saves model_kind='dino_volume_scene_hires'.")
    p.add_argument("--hires_per_voxel", action="store_true",
                   help="Use DinoVolumeScenePerVoxelHires — DINOv3 + 5-layer "
                        "CNN at pred_size=256, then 1×1 conv expansion to "
                        "feat_dim*n_z and reshape to per-voxel volume. "
                        "Defaults: n_height_bins=64, feat_dim=16. "
                        "Saves model_kind='dino_volume_per_voxel_hires'.")
    p.add_argument("--hires_flat", action="store_true",
                   help="Use DinoVolumeSceneHiresFlat — same backbone + CNN "
                        "as DinoVolumeSceneHires, but the per-T queries come "
                        "from a single flat MLP output (output dim = T × "
                        "per_T_dim) instead of T independent AdaLN-Zero "
                        "forwards. Better temporal coherence by construction.")
    p.add_argument("--volumetric", action="store_true",
                   help="Use DinoVolumeSceneVolumetric — unproject image "
                        "features into a P×P×Z world-space voxel grid, "
                        "concat sin/cos height encoding, 3-layer per-voxel "
                        "MLP, then dot-product against per-T queries from a "
                        "flat global MLP. Adds a +1 GT-voxel slot at training "
                        "for precise (unquantized) supervision. Defaults: "
                        "pred_size=128, n_height_bins=64, feat_dim=32.")
    p.add_argument("--cls_fusion", default="concat", choices=["concat", "primary"],
                   help="Multi-view CLS fusion mode for the grip/rot head. "
                        "'concat' (new default) concatenates every view's CLS "
                        "before the global MLP; 'primary' uses only views[0]'s "
                        "CLS (matches the original 2026-06-09 ckpts).")
    p.add_argument("--cross_view_layers", type=int, default=0,
                   help="DA3-style alternating within/cross-view self-attention "
                        "blocks over the native DINO patch grid (~31x31). 0 disables. "
                        "4-5 is the sweet spot per Depth-Anything-3 (arXiv:2511.10647).")
    p.add_argument("--cross_view_heads", type=int, default=6)
    p.add_argument("--v4", action="store_true",
                   help="v4 = v2 multi-view MLP fusion + sin/cos height PE "
                        "+ per-view adaptive z range. Uses "
                        "DinoVolumeSceneV4. Keeps the same out dict shape "
                        "as the v2 multiview model so multiview_losses + "
                        "viz pipeline reuse unchanged. Requires --multiview.")
    p.add_argument("--v4_da3", action="store_true",
                   help="v4 head + DA3 backbone (multi-view ViT with "
                        "built-in cross-attention + camera token; NO DPT "
                        "head). Requires --multiview.")
    p.add_argument("--v4_vlm", action="store_true",
                   help="v4 head + InternVL2.5-1B vision tower (LLM frozen, "
                        "vision trainable, text-conditioned via task name). "
                        "Requires --multiview.")
    p.add_argument("--v4_dynaflip", action="store_true",
                   help="v4 head + DynaFLIP-base vision backbone "
                        "(jlee-larr/dynaflip-base, ~0.2 B params). "
                        "Per-view loop, no backbone-level cross-view fusion. "
                        "Requires --multiview.")
    p.add_argument("--v4_paligemma", action="store_true",
                   help="v4 head + PaliGemma vision tower + Gemma LLM "
                        "(google/paligemma-3b-pt-448 by default). Mirrors "
                        "v4_vlm but with the PaliGemma stack instead of "
                        "InternVL. Requires --multiview.")
    p.add_argument("--v4_pi0", action="store_true",
                   help="v4 head + pi0's PaliGemma backbone (LeRobot port "
                        "lerobot/pi0_base, ~14 GB safetensors). PaliGemma "
                        "weights replaced by pi0's robot-data pretrain; "
                        "the flow-matching action expert is discarded. "
                        "Requires --multiview.")
    p.add_argument("--pi0_safetensors_path", default=None,
                   help="Optional override for the pi0_base safetensors "
                        "file. Default = pull from HF cache.")
    p.add_argument("--paligemma_model_name",
                   default="google/paligemma-3b-pt-448",
                   help="HF model id for the PaliGemma variant. The 448 "
                        "variant fits puget 24 GB at B=2 with grad-ckpt + "
                        "frozen LLM; for unfrozen-LLM training run on yukon "
                        "(96 GB).")
    p.add_argument("--vlm_task_text", default="pick and place",
                   help="Task prompt for the VLM (e.g. dataset name).")
    p.add_argument("--vlm_freeze_llm", action=argparse.BooleanOptionalAction,
                   default=False,
                   help="Default: train the LLM (gradient checkpointing on). "
                        "--no-vlm_freeze_llm keeps LLM frozen.")
    p.add_argument("--cam_margin_m", type=float, default=0.01,
                   help="v4: each view's z_hi is capped at cam_z - this "
                        "margin (metres). Wrist cam → adaptive small z "
                        "range; scene cam → unchanged (sanity-checked to "
                        "stay above global z_hi).")
    p.add_argument("--v3", action="store_true",
                   help="v3 factorized dot-product model "
                        "(DinoVolumeSceneV3). Per-pixel feature splits "
                        "into [F_z|F_t|F_s] which dot-product against "
                        "[PE_z|PE_t|spatial_q]. No per-voxel MLP. Factored "
                        "loss never materialises the (B, T, V) volume — "
                        "supports P=Z=256 cheaply.")
    p.add_argument("--v3_F_z", type=int, default=16)
    p.add_argument("--v3_F_t", type=int, default=16)
    p.add_argument("--v3_F_s", type=int, default=32)
    p.add_argument("--per_voxel_grip_rot", action=argparse.BooleanOptionalAction,
                   default=True,
                   help="Predict grip + rot from the GT-nearest voxel feature "
                        "(teacher-forced at train, argmax at infer) via a small "
                        "1D-CNN over T, instead of from the global MLP query. "
                        "Same idea as production PARA's pixel-indexed heads. "
                        "Default ON (2026-06-10 — Cameron made this canonical). "
                        "Pass --no-per_voxel_grip_rot to fall back to the global head.")
    p.add_argument("--temporal_head", default="flat_mlp",
                   choices=["flat_mlp", "temporal_cnn"],
                   help="Per-voxel grip/rot temporal head. 'flat_mlp' (new "
                        "default 2026-06-10): 4-layer dense MLP over the "
                        "T·F flattened gathered features → full receptive "
                        "field, implicit absolute positional info. "
                        "'temporal_cnn': 3-layer 1D conv over T — the old "
                        "default; receptive field ~7 timesteps, no time PE.")
    p.add_argument("--temporal_cnn_layers", type=int, default=3,
                   help="1D CNN depth (only with --temporal_head temporal_cnn).")
    p.add_argument("--temporal_cnn_kernel", type=int, default=3)
    p.add_argument("--temporal_mlp_hidden", type=int, default=512,
                   help="Hidden dim of the flat MLP temporal head.")
    p.add_argument("--mv_past_n", type=int, default=0,
                   help="(multi-view only) Past N frames of (grip, z) fed "
                        "as proprio conditioning to the global query MLP. "
                        "0 disables. 8 was Cameron's 2026-06-10 default — "
                        "matches frame_stride * 8 = 32 raw frames of recent "
                        "EE state, enough to disambiguate 'just closed' vs "
                        "'closed for a while → ready to release'.")
    p.add_argument("--mv_past_enc_dim", type=int, default=64,
                   help="Encoder output dim for the past-trajectory MLP.")
    p.add_argument("--multiview", default=None,
                   help="Comma-separated views for the multi-view volumetric "
                        "model — e.g. 'scene', 'wrist', or 'scene,wrist'. "
                        "Builds one P×P×Z sub-volume per view; each voxel gets "
                        "features bilinear-sampled from ALL listed views' "
                        "feature maps, concatenated, then fused by a 3-layer "
                        "MLP. Loss is nearest-3D-point across the full "
                        "N×P×P×Z combined volume. Mutex with --volumetric.")
    p.add_argument("--height_enc_dim", type=int, default=16,
                   help="Volumetric model only: sin/cos PE dim for the "
                        "per-voxel height encoding.")
    p.add_argument("--voxel_mlp_hidden", type=int, default=64,
                   help="Volumetric model only: per-voxel MLP hidden width.")
    p.add_argument("--refit_z_range", action="store_true",
                   help="When finetuning, ignore the source ckpt's z_range and "
                        "let YamScene auto-fit a fresh range on the new dataset. "
                        "Also drops ``height_emb.weight`` from the source state_dict "
                        "(the bin-index → physical-Z mapping changes, so the "
                        "embedding is stale). Required when the new task's EE-Z "
                        "extent differs materially from the source's (e.g. "
                        "table-level manipulation finetuned from a pickplace cup).")
    p.add_argument("--refit_rot_centroids", action="store_true",
                   help="Used with --init_from: refit K-means rot centroids on "
                        "the NEW dataset instead of inheriting from the source "
                        "ckpt, AND skip loading rot_head.* weights (since their "
                        "output indices map to the OLD centroid IDs). Use when "
                        "the new dataset has a different rotation distribution "
                        "from the source ckpt's training data.")
    # Ablation flags (consumed only by --hires_ablate model).
    p.add_argument("--hires_ablate", action="store_true",
                   help="Use DinoVolumeSceneHiresAblate (factorized hires with "
                        "configurable past-encoding / MLP / backbone-fusion "
                        "ablation knobs). Mutex with the other model flags.")
    p.add_argument("--past_n", type=int, default=0,
                   help="Past timesteps to encode (0 = no past).")
    p.add_argument("--past_z_delta", action="store_true",
                   help="Encode past_z as (past_z − current_z) instead of "
                        "absolute height. Pairs naturally with --past_uv_delta.")
    p.add_argument("--past_use_uv", action="store_true",
                   help="Concat past pixel coords (img_size px → [-1, 1]).")
    p.add_argument("--past_uv_delta", action="store_true",
                   help="With --past_use_uv: subtract current start_pix.")
    p.add_argument("--past_use_xyz", action="store_true",
                   help="Concat past world XYZ (m, centered ~0.3).")
    p.add_argument("--past_xyz_delta", action="store_true",
                   help="With --past_use_xyz: also concat past_xyz - current_xyz.")
    p.add_argument("--mlp_layers", type=int, default=1,
                   help="Depth of the global MLP (1=plain Linear, ≥2=MLP).")
    p.add_argument("--mlp_residual", action="store_true",
                   help="Residual MLP (skip every block; needs --mlp_layers ≥3).")
    p.add_argument("--mlp_hidden_mult", type=float, default=2.0,
                   help="MLP hidden = d_model × this (mlp_layers > 1 only).")
    p.add_argument("--backbone_layer_fusion", action="store_true",
                   help="Concat DINOv3 patch tokens from layers [3,6,9,11] "
                        "before feat_head (4× embed_dim per patch).")
    args = p.parse_args()

    ckpt_dir = Path(args.ckpt_root) / args.name
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # If finetuning, load the source ckpt up front so we can reuse its
    # K-means rot centroids AND z_range in the dataset constructor.
    # Refitting either would scramble the volume / rot heads.
    init_ck = None
    if args.init_from:
        init_ck = torch.load(args.init_from, map_location="cpu", weights_only=False)
        print(f"  finetune from {args.init_from}  "
              f"(source step={init_ck['step']})")
        # Source ckpt's z_range wins on finetune UNLESS --refit_z_range.
        if not args.refit_z_range:
            src_z = init_ck["dataset_cfg"]["z_range"]
            args.z_lo, args.z_hi = float(src_z[0]), float(src_z[1])
        else:
            print("  --refit_z_range: ignoring source ckpt's z_range; "
                  "auto-fitting fresh range on new dataset.")

    # z_range: explicit (CLI or finetune source) → use it; else None → YamScene auto.
    z_range_arg = ((args.z_lo, args.z_hi)
                    if (args.z_lo is not None and args.z_hi is not None)
                    else None)
    # rot_centroids: inherit from source ckpt unless --refit_rot_centroids
    # was passed (in which case the new dataset's rotation distribution drives
    # a fresh K-means fit, and we'll also skip loading rot_head weights below).
    rot_c_from_src = (
        init_ck["dataset_cfg"]["rot_centroids"]
        if (init_ck is not None and not args.refit_rot_centroids)
        else None)
    if args.refit_rot_centroids and init_ck is not None:
        print("  --refit_rot_centroids: ignoring source ckpt's rot_centroids; "
              "refitting K-means on new dataset.")
    ds = YamScene(
        root=Path(args.data_root),
        n_window=args.n_window, frame_stride=args.frame_stride,
        pred_size=args.pred_size, img_size=args.img_size,
        n_z=args.n_height_bins, n_gripper_bins=args.n_gripper_bins,
        n_rot_clusters=args.n_rot_clusters,
        z_range=z_range_arg, grip_range=(args.grip_lo, args.grip_hi),
        rot_centroids=rot_c_from_src,
        past_n=(args.past_n if args.hires_ablate
                else args.mv_past_n if args.multiview
                else 0),
    )
    # Reflect the chosen range into args so ckpt dump + viz panels see it.
    args.z_lo, args.z_hi = ds.z_lo, ds.z_hi
    dataset_cfg = ds.config()

    _past_keys = ("past_z", "past_grip", "past_uv", "past_xyz", "current_xyz")

    def _collate(items):
        out = {}
        for k in ("rgb", "start_pix", "target_yx", "target_z",
                  "target_grip", "target_rot"):
            out[k] = torch.stack([x[k] for x in items])
        out["_bgr_src"] = [x["_bgr_src"] for x in items]
        if "_wrist_bgr_src" in items[0]:
            out["_wrist_bgr_src"] = [x["_wrist_bgr_src"] for x in items]
        # Volumetric + multi-view models need target_xyz + K + T_w2c per
        # sample (per view). Other models ignore these; cheap to include
        # when available.
        for k in ("target_xyz", "_K_in", "_T_w2c",
                  "wrist_rgb", "_K_wrist", "_T_w2c_wrist"):
            if k in items[0]:
                out[k] = torch.stack([x[k] for x in items])
        if all(k in items[0] for k in _past_keys):
            for k in _past_keys:
                out[k] = torch.stack([x[k] for x in items])
        return out

    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
                         num_workers=args.num_workers, pin_memory=True,
                         drop_last=True, collate_fn=_collate)

    mutex_flags = sum([args.da3, args.per_voxel,
                       args.hires_factorized, args.hires_per_voxel,
                       args.hires_ablate, args.hires_flat,
                       args.volumetric, bool(args.multiview)])
    if mutex_flags > 1:
        raise ValueError("--da3 / --per_voxel / --hires_factorized / "
                         "--hires_per_voxel / --hires_ablate / --hires_flat / "
                         "--volumetric / --multiview are mutually exclusive")
    multi_views = (args.multiview.split(",") if args.multiview else None)
    if args.da3:
        from lib.model_da3 import DA3VolumeScene
        model_cls = DA3VolumeScene
    elif args.hires_factorized:
        from lib.model_hires import DinoVolumeSceneHires
        model_cls = DinoVolumeSceneHires
    elif args.hires_per_voxel:
        from lib.model_hires import DinoVolumeScenePerVoxelHires
        model_cls = DinoVolumeScenePerVoxelHires
    elif args.hires_ablate:
        from lib.model_hires import DinoVolumeSceneHiresAblate
        model_cls = DinoVolumeSceneHiresAblate
    elif args.hires_flat:
        from lib.model_hires import DinoVolumeSceneHiresFlat
        model_cls = DinoVolumeSceneHiresFlat
    elif args.volumetric:
        from lib.model_volumetric import DinoVolumeSceneVolumetric
        model_cls = DinoVolumeSceneVolumetric
    elif args.v4_da3:
        assert multi_views, "--v4_da3 requires --multiview to be set."
        from lib.model_volumetric_v4_da3 import DinoVolumeSceneV4DA3
        model_cls = DinoVolumeSceneV4DA3
    elif args.v4_vlm:
        assert multi_views, "--v4_vlm requires --multiview to be set."
        from lib.model_volumetric_v4_vlm import DinoVolumeSceneV4VLM
        model_cls = DinoVolumeSceneV4VLM
    elif args.v4_dynaflip:
        assert multi_views, "--v4_dynaflip requires --multiview to be set."
        from lib.model_volumetric_v4_dynaflip import DinoVolumeSceneV4DynaFLIP
        model_cls = DinoVolumeSceneV4DynaFLIP
    elif args.v4_paligemma:
        assert multi_views, "--v4_paligemma requires --multiview to be set."
        from lib.model_volumetric_v4_paligemma import DinoVolumeSceneV4PaliGemma
        model_cls = DinoVolumeSceneV4PaliGemma
    elif args.v4_pi0:
        assert multi_views, "--v4_pi0 requires --multiview to be set."
        from lib.model_volumetric_v4_pi0 import DinoVolumeSceneV4Pi0
        model_cls = DinoVolumeSceneV4Pi0
    elif args.v4:
        assert multi_views, "--v4 requires --multiview to be set."
        from lib.model_volumetric_v4 import DinoVolumeSceneV4
        model_cls = DinoVolumeSceneV4
    elif args.v3:
        assert multi_views, "--v3 requires --multiview to be set."
        from lib.model_volumetric_v3 import DinoVolumeSceneV3
        model_cls = DinoVolumeSceneV3
    elif multi_views:
        from lib.model_volumetric_multiview import DinoVolumeSceneVolumetricMultiView
        model_cls = DinoVolumeSceneVolumetricMultiView
    elif args.per_voxel:
        model_cls = DinoVolumeScenePerVoxel
    else:
        model_cls = DinoVolumeScene
    # Per-voxel + DA3 + hires_per_voxel default feat_dim differs from factorized's 48.
    # Caller can override with --feat_dim explicitly.
    if args.da3 and args.feat_dim == 48:
        model_feat_dim = 32
    elif args.hires_per_voxel and args.feat_dim == 48:
        model_feat_dim = 16
    elif args.per_voxel and args.feat_dim == 48:
        model_feat_dim = 32
    elif args.volumetric and args.feat_dim == 48:
        model_feat_dim = 32
    elif args.v3 and args.feat_dim == 48:
        model_feat_dim = args.v3_F_z + args.v3_F_t + args.v3_F_s
    elif multi_views and args.feat_dim == 48:
        model_feat_dim = 32
    else:
        model_feat_dim = args.feat_dim
    args.feat_dim = model_feat_dim
    # Build kwargs; ablate + volumetric models take extra keys.
    if args.volumetric:
        # Cameron's spec: P=128, Z=64, feat=32, height_enc=16.
        model_kwargs = dict(
            n_window=args.n_window,
            n_height_bins=args.n_height_bins,
            pred_size=args.pred_size,
            feat_dim=model_feat_dim,
            height_enc_dim=args.height_enc_dim,
            voxel_mlp_hidden=args.voxel_mlp_hidden,
            d_model=args.d_model,
            mlp_layers=args.mlp_layers,
            mlp_hidden_mult=args.mlp_hidden_mult,
            n_gripper_bins=args.n_gripper_bins,
            n_rot_clusters=args.n_rot_clusters,
            img_size=args.img_size,
            z_lo=args.z_lo, z_hi=args.z_hi,
        )
    elif args.v4_da3 or args.v4_vlm or args.v4_dynaflip or args.v4_paligemma:
        # Same kwargs as v4 + per-backbone extras. The model class
        # itself loads the swapped backbone in __init__.
        model_kwargs = dict(
            views=multi_views,
            n_window=args.n_window,
            n_height_bins=args.n_height_bins,
            pred_size=args.pred_size,
            feat_dim=model_feat_dim,
            height_enc_dim=args.height_enc_dim,
            voxel_mlp_hidden=args.voxel_mlp_hidden,
            d_model=args.d_model,
            mlp_layers=args.mlp_layers,
            mlp_hidden_mult=args.mlp_hidden_mult,
            n_gripper_bins=args.n_gripper_bins,
            n_rot_clusters=args.n_rot_clusters,
            img_size=args.img_size,
            z_lo=args.z_lo, z_hi=args.z_hi,
            cls_fusion=args.cls_fusion,
            cross_view_layers=args.cross_view_layers,
            cross_view_heads=args.cross_view_heads,
            per_voxel_grip_rot=args.per_voxel_grip_rot,
            temporal_head=args.temporal_head,
            temporal_cnn_layers=args.temporal_cnn_layers,
            temporal_cnn_kernel=args.temporal_cnn_kernel,
            temporal_mlp_hidden=args.temporal_mlp_hidden,
            past_n=args.mv_past_n,
            past_enc_dim=args.mv_past_enc_dim,
            cam_margin_m=args.cam_margin_m,
        )
        if args.v4_vlm:
            model_kwargs["task_text_default"] = args.vlm_task_text
            model_kwargs["freeze_llm"] = args.vlm_freeze_llm
        if args.v4_paligemma:
            model_kwargs["paligemma_model_name"] = args.paligemma_model_name
            model_kwargs["task_text_default"] = args.vlm_task_text
            model_kwargs["freeze_llm"] = args.vlm_freeze_llm
        if args.v4_pi0:
            model_kwargs["paligemma_model_name"] = args.paligemma_model_name
            model_kwargs["task_text_default"] = args.vlm_task_text
            model_kwargs["freeze_llm"] = args.vlm_freeze_llm
            if args.pi0_safetensors_path is not None:
                model_kwargs["pi0_safetensors_path"] = args.pi0_safetensors_path
    elif args.v4:
        # Same kwargs as the v2 multiview model + cam_margin_m. The v4
        # model class internally swaps height_emb → sin/cos PE and
        # computes per-view adaptive z range.
        model_kwargs = dict(
            views=multi_views,
            n_window=args.n_window,
            n_height_bins=args.n_height_bins,
            pred_size=args.pred_size,
            feat_dim=model_feat_dim,
            height_enc_dim=args.height_enc_dim,
            voxel_mlp_hidden=args.voxel_mlp_hidden,
            d_model=args.d_model,
            mlp_layers=args.mlp_layers,
            mlp_hidden_mult=args.mlp_hidden_mult,
            n_gripper_bins=args.n_gripper_bins,
            n_rot_clusters=args.n_rot_clusters,
            img_size=args.img_size,
            z_lo=args.z_lo, z_hi=args.z_hi,
            cls_fusion=args.cls_fusion,
            cross_view_layers=args.cross_view_layers,
            cross_view_heads=args.cross_view_heads,
            per_voxel_grip_rot=args.per_voxel_grip_rot,
            temporal_head=args.temporal_head,
            temporal_cnn_layers=args.temporal_cnn_layers,
            temporal_cnn_kernel=args.temporal_cnn_kernel,
            temporal_mlp_hidden=args.temporal_mlp_hidden,
            past_n=args.mv_past_n,
            past_enc_dim=args.mv_past_enc_dim,
            cam_margin_m=args.cam_margin_m,
        )
    elif args.v3:
        model_kwargs = dict(
            views=multi_views,
            n_window=args.n_window,
            n_height_bins=args.n_height_bins,
            pred_size=args.pred_size,
            d_model=args.d_model,
            mlp_layers=args.mlp_layers,
            mlp_hidden_mult=args.mlp_hidden_mult,
            n_gripper_bins=args.n_gripper_bins,
            n_rot_clusters=args.n_rot_clusters,
            img_size=args.img_size,
            z_lo=args.z_lo, z_hi=args.z_hi,
            F_z=args.v3_F_z,
            F_t=args.v3_F_t,
            F_s=args.v3_F_s,
            past_n=args.mv_past_n,
            past_enc_dim=args.mv_past_enc_dim,
        )
    elif multi_views:
        model_kwargs = dict(
            views=multi_views,
            n_window=args.n_window,
            n_height_bins=args.n_height_bins,
            pred_size=args.pred_size,
            feat_dim=model_feat_dim,
            height_enc_dim=args.height_enc_dim,
            voxel_mlp_hidden=args.voxel_mlp_hidden,
            d_model=args.d_model,
            mlp_layers=args.mlp_layers,
            mlp_hidden_mult=args.mlp_hidden_mult,
            n_gripper_bins=args.n_gripper_bins,
            n_rot_clusters=args.n_rot_clusters,
            img_size=args.img_size,
            z_lo=args.z_lo, z_hi=args.z_hi,
            cls_fusion=args.cls_fusion,
            cross_view_layers=args.cross_view_layers,
            cross_view_heads=args.cross_view_heads,
            per_voxel_grip_rot=args.per_voxel_grip_rot,
            temporal_head=args.temporal_head,
            temporal_cnn_layers=args.temporal_cnn_layers,
            temporal_cnn_kernel=args.temporal_cnn_kernel,
            temporal_mlp_hidden=args.temporal_mlp_hidden,
            past_n=args.mv_past_n,
            past_enc_dim=args.mv_past_enc_dim,
        )
    else:
        model_kwargs = dict(
            n_window=args.n_window, n_height_bins=args.n_height_bins,
            pred_size=args.pred_size, feat_dim=model_feat_dim,
            d_model=args.d_model, n_blocks=args.n_blocks,
            mlp_hidden=args.mlp_hidden,
            n_gripper_bins=args.n_gripper_bins,
            n_rot_clusters=args.n_rot_clusters,
        )
    if args.hires_ablate:
        model_kwargs.update(
            img_size=args.img_size,
            past_n=args.past_n,
            past_z_delta=args.past_z_delta,
            past_use_uv=args.past_use_uv,
            past_uv_delta=args.past_uv_delta,
            past_use_xyz=args.past_use_xyz,
            past_xyz_delta=args.past_xyz_delta,
            mlp_layers=args.mlp_layers,
            mlp_residual=args.mlp_residual,
            mlp_hidden_mult=args.mlp_hidden_mult,
            backbone_layer_fusion=args.backbone_layer_fusion,
        )
    model = model_cls(**model_kwargs).to(args.device)
    print(f"  model = {model_cls.__name__} (feat_dim={model_feat_dim}, "
          f"n_height_bins={args.n_height_bins})")
    if init_ck is not None:
        src_sd = init_ck["model"]
        if args.refit_rot_centroids:
            dropped = [k for k in src_sd if k.startswith("rot_head.")]
            src_sd = {k: v for k, v in src_sd.items() if not k.startswith("rot_head.")}
            print(f"  --refit_rot_centroids: dropped {len(dropped)} rot_head "
                  "keys from source ckpt (keeping fresh random init).")
        if args.refit_z_range:
            # bin_centers_z is registered as a BUFFER (not a parameter), so
            # load_state_dict will silently overwrite it back to the source's
            # cup-derived bins unless we drop it here. height_emb is the
            # learned mapping per bin and must also be reinitialized.
            drop_prefixes = ("height_emb.", "bin_centers_z")
            dropped_z = [k for k in src_sd
                          if any(k == p or k.startswith(p) for p in drop_prefixes)]
            src_sd = {k: v for k, v in src_sd.items() if k not in dropped_z}
            print(f"  --refit_z_range: dropped {len(dropped_z)} z-axis keys "
                  f"from source ckpt ({dropped_z}); bin-index → physical-Z "
                  "mapping changed, keeping fresh random init for height_emb "
                  "and the model's freshly-computed bin_centers_z buffer.")
        # Shape-aware filter: skip keys whose shape changed (e.g., t_sin if
        # n_window doubled, z_sin if n_height_bins changed). These are
        # buffers/embeddings re-derivable at init from the new hparams.
        model_sd = model.state_dict()
        shape_skipped = []
        filtered_sd = {}
        for k, v in src_sd.items():
            if k in model_sd and model_sd[k].shape != v.shape:
                shape_skipped.append((k, tuple(v.shape), tuple(model_sd[k].shape)))
            else:
                filtered_sd[k] = v
        if shape_skipped:
            print(f"  shape-skip {len(shape_skipped)} keys (n_window or "
                  "n_height_bins probably changed):")
            for k, src_shape, dst_shape in shape_skipped[:5]:
                print(f"    {k}: src {src_shape} → dst {dst_shape}")
        missing, unexpected = model.load_state_dict(filtered_sd, strict=False)
        if missing and not args.refit_rot_centroids:
            print(f"  load_state_dict missing keys: {missing[:5]}")
        if unexpected:
            print(f"  load_state_dict unexpected keys: {unexpected[:5]}")
        print(f"  loaded model weights from {args.init_from}")
    # Split params into backbone (pretrained) vs head (everything else)
    # so backbone can run at a lower LR than the fresh head. The ablation
    # experiment showed lr=2e-5 was the safe DINO LR for from-scratch.
    #
    # Prefixes for every pretrained backbone attribute name we use:
    #   ``backbone.``  — DINOv3 (v4 / v4_da3 / hires / volumetric models)
    #   ``dynaflip.``  — DynaFLIP (v4_dynaflip)
    #   ``vlm.``       — InternVL (v4_vlm) and PaliGemma (v4_paligemma)
    # Bug fixed 2026-06-15: v4_vlm / v4_paligemma / v4_dynaflip were
    # incorrectly routed entirely to the head group at the high LR, which
    # destabilised PaliGemma's 3 B LLM. Each new backbone must register its
    # prefix here.
    _BB_PREFIXES = ("backbone.", "dynaflip.", "vlm.")
    backbone_params = [
        p for n, p in model.named_parameters()
        if p.requires_grad and any(n.startswith(pre) for pre in _BB_PREFIXES)
    ]
    bb_ids = {id(p) for p in backbone_params}
    head_params = [p for p in model.parameters()
                   if p.requires_grad and id(p) not in bb_ids]
    head_lr_max = args.lr
    head_lr_min = args.lr_min if args.lr_min is not None else args.lr
    bb_lr_max = args.lr * args.backbone_lr_mult
    bb_lr_min = head_lr_min * args.backbone_lr_mult
    print(f"  param groups: head={len(head_params)} ({head_lr_max:.1e}→{head_lr_min:.1e}) "
          f"backbone={len(backbone_params)} ({bb_lr_max:.1e}→{bb_lr_min:.1e}) "
          f"warmup={args.lr_warmup} schedule={args.lr_schedule}")
    opt = torch.optim.AdamW([
        {"params": head_params,     "lr": head_lr_max, "name": "head",
         "lr_max": head_lr_max, "lr_min": head_lr_min},
        {"params": backbone_params, "lr": bb_lr_max,   "name": "backbone",
         "lr_max": bb_lr_max,   "lr_min": bb_lr_min},
    ])

    def _lr_at(it: int, lr_max: float, lr_min: float) -> float:
        if args.lr_warmup > 0 and it < args.lr_warmup:
            return lr_max * (it + 1) / args.lr_warmup
        if args.lr_schedule == "cosine":
            denom = max(1, args.max_iters - args.lr_warmup)
            progress = min(1.0, max(0.0, (it - args.lr_warmup) / denom))
            return lr_min + 0.5 * (lr_max - lr_min) * (1.0 + math.cos(math.pi * progress))
        return lr_max

    wandb.init(project=args.wandb_project, name=args.name, config=vars(args))
    print(f"wandb: {wandb.run.url if wandb.run else '(disabled)'}")
    print(f"ckpts → {ckpt_dir}")

    def _past_kw(batch):
        """Build past-feature kwargs for model.forward() — empty dict if
        --hires_ablate isn't set or no past keys are in the batch."""
        if not args.hires_ablate or "past_z" not in batch:
            return {}
        return {
            "past_z": batch["past_z"].to(args.device, non_blocking=True),
            "past_grip": batch["past_grip"].to(args.device, non_blocking=True),
            "past_uv": batch["past_uv"].to(args.device, non_blocking=True),
            "past_xyz": batch["past_xyz"].to(args.device, non_blocking=True),
            "current_xyz": batch["current_xyz"].to(args.device, non_blocking=True),
        }

    def _volumetric_kw(batch: dict) -> dict:
        """K_in, T_w2c kwargs the volumetric model needs at forward."""
        if not args.volumetric:
            return {}
        return {
            "K_in": batch["_K_in"].to(args.device, non_blocking=True).float(),
            "T_w2c": batch["_T_w2c"].to(args.device, non_blocking=True).float(),
        }

    if args.volumetric:
        from lib.model_volumetric import volumetric_losses as _vol_losses
    else:
        _vol_losses = None

    def _multiview_call(model, batch):
        """Forward call for the multi-view volumetric model. Builds the
        (rgb, K, T_w2c) lists in `multi_views` order from the batch."""
        rgb_list = []
        K_list = []
        T_list = []
        for v in multi_views:
            if v == "scene":
                rgb_list.append(batch["rgb"].to(args.device, non_blocking=True).float())
                K_list.append(batch["_K_in"].to(args.device, non_blocking=True).float())
                T_list.append(batch["_T_w2c"].to(args.device, non_blocking=True).float())
            elif v == "wrist":
                rgb_list.append(batch["wrist_rgb"].to(args.device, non_blocking=True).float())
                K_list.append(batch["_K_wrist"].to(args.device, non_blocking=True).float())
                T_list.append(batch["_T_w2c_wrist"].to(args.device, non_blocking=True).float())
            else:
                raise ValueError(f"unknown view: {v!r}")
        sp = batch["start_pix"].to(args.device, non_blocking=True).float()
        # v3 always uses target_xyz for chosen-voxel feature gather; the
        # multiview model only does so when per_voxel_grip_rot is on.
        want_tgt = args.v3 or args.per_voxel_grip_rot
        tgt = (batch["target_xyz"].to(args.device, non_blocking=True).float()
                if want_tgt else None)
        # Past-trajectory conditioning (only when --mv_past_n > 0).
        past_grip = past_z = None
        if args.mv_past_n > 0:
            past_grip = batch["past_grip"].to(args.device, non_blocking=True).float()
            past_z = batch["past_z"].to(args.device, non_blocking=True).float()
        if args.v3:
            return model(rgb_list, K_in=K_list, T_w2c=T_list, target_xyz=tgt,
                          past_grip=past_grip, past_z=past_z)
        return model(rgb_list, sp, K_in=K_list, T_w2c=T_list, target_xyz=tgt,
                      past_grip=past_grip, past_z=past_z)

    if args.v3:
        from lib.model_volumetric_v3 import v3_losses as _mv_losses
    elif multi_views:
        from lib.model_volumetric_multiview import multiview_losses as _mv_losses
    else:
        _mv_losses = None

    # Step-0 baseline viz.
    with torch.no_grad():
        peek = next(iter(loader))
        rgb0 = peek["rgb"].to(args.device, non_blocking=True)
        sp0 = peek["start_pix"].to(args.device, non_blocking=True)
        if multi_views:
            out0 = _multiview_call(model, peek)
        else:
            out0 = model(rgb0, sp0, **_past_kw(peek), **_volumetric_kw(peek))
        wandb.log(_viz_batch(peek, out0, img_size=args.img_size,
                            per_voxel=(args.per_voxel or args.hires_per_voxel or args.volumetric), z_lo=args.z_lo, z_hi=args.z_hi), step=0)
        print("  step 0 viz logged (pre-training baseline)")

    step = 0
    t_loop = time.perf_counter()
    while step < args.max_iters:
        for batch in loader:
            rgb = batch["rgb"].to(args.device, non_blocking=True)
            sp = batch["start_pix"].to(args.device, non_blocking=True)
            tgt_yx = batch["target_yx"].to(args.device, non_blocking=True)
            tgt_z = batch["target_z"].to(args.device, non_blocking=True)
            tgt_g = batch["target_grip"].to(args.device, non_blocking=True)
            tgt_r = batch["target_rot"].to(args.device, non_blocking=True)
            for pg in opt.param_groups:
                pg["lr"] = _lr_at(step, pg["lr_max"], pg["lr_min"])
            if multi_views:
                out = _multiview_call(model, batch)
                tgt_xyz = batch["target_xyz"].to(args.device, non_blocking=True).float()
                if args.v3:
                    losses = _mv_losses(out, tgt_xyz, tgt_g, tgt_r)
                else:
                    losses = _mv_losses(out, tgt_yx, tgt_z, tgt_g, tgt_r, tgt_xyz)
            else:
                out = model(rgb, sp, **_past_kw(batch), **_volumetric_kw(batch))
                if _vol_losses is not None:
                    tgt_xyz = batch["target_xyz"].to(args.device, non_blocking=True).float()
                    losses = _vol_losses(out, tgt_yx, tgt_z, tgt_g, tgt_r, tgt_xyz)
                else:
                    losses = all_losses(out, tgt_yx, tgt_z, tgt_g, tgt_r)
            opt.zero_grad(set_to_none=True)
            losses["loss/total"].backward()
            opt.step()
            step += 1

            now = time.perf_counter()
            its = 1.0 / (now - t_loop)
            t_loop = now

            log = {"iter": step, "it_per_s": its}
            for pg in opt.param_groups:
                log[f"lr/{pg.get('name', 'group')}"] = pg["lr"]
            log["lr"] = opt.param_groups[0]["lr"]  # head LR (legacy key)
            for k, v in losses.items():
                log[k] = v.item()
            if step % args.vis_every_iters == 0:
                log.update(_viz_batch(batch, out, img_size=args.img_size,
                                       per_voxel=(args.per_voxel or args.hires_per_voxel or args.volumetric),
                                       z_lo=args.z_lo, z_hi=args.z_hi))
            wandb.log(log, step=step)

            if step % args.ckpt_every_iters == 0:
                tmp = ckpt_dir / "latest.pth.tmp"
                torch.save({
                    "model": model.state_dict(),
                    "opt": opt.state_dict(),
                    "step": step,
                    "args": vars(args),
                    "dataset_cfg": dataset_cfg,
                    "model_kind": (
                        "dino_volume_scene_v4_da3" if args.v4_da3
                        else "dino_volume_scene_v4_vlm" if args.v4_vlm
                        else "dino_volume_scene_v4_dynaflip" if args.v4_dynaflip
                        else "dino_volume_scene_v4_paligemma" if args.v4_paligemma
                        else "dino_volume_scene_v4_pi0" if args.v4_pi0
                        else "dino_volume_scene_v4" if args.v4
                        else "dino_volume_scene_v3" if args.v3
                        else "da3_volume_scene" if args.da3
                        else "dino_volume_scene_hires_ablate" if args.hires_ablate
                        else "dino_volume_scene_hires_flat" if args.hires_flat
                        else "dino_volume_scene_hires" if args.hires_factorized
                        else "dino_volume_per_voxel_hires" if args.hires_per_voxel
                        else "dino_volume_per_voxel" if args.per_voxel
                        else "dino_volume_scene_volumetric_multiview" if args.multiview
                        else "dino_volume_scene_volumetric" if args.volumetric
                        else "dino_volume_scene"),
                }, tmp)
                tmp.replace(ckpt_dir / "latest.pth")
                print(f"  ckpt @ {step}  total={losses['loss/total'].item():.4f} "
                      f"vol={losses['loss/volume'].item():.4f} "
                      f"grip={losses['loss/grip'].item():.4f} "
                      f"rot={losses['loss/rot'].item():.4f}  {its:.2f} it/s",
                      flush=True)

            if step >= args.max_iters:
                break


if __name__ == "__main__":
    main()