"""Inference-time prediction decoder.

A checkpoint saved by ``train.py`` carries everything needed to decode
the model's argmax outputs back to physical quantities. Use::

    ck = load_checkpoint("…/latest.pth", device="cuda")
    pred = decode(ck["model"], rgb, start_pix, ck["cfg"])
    # pred["pix_uv"]   (T, 2) — pixel in img_size (squish) coords
    # pred["height_m"] (T,)   — predicted EE Z in meters (left_arm_base)
    # pred["grip"]     (T,)   — predicted gripper in original units
    # pred["quat"]     (T, 4) — predicted EE rotation quaternion (w, x, y, z)
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import torch

from .model import DinoVolumeScene


def load_checkpoint(path: Path, device: str = "cuda"):
    """Load either ``DinoVolumeScene`` or ``DinoCLSXYZBaseline`` ckpt.

    Auto-detects via the ``model_kind`` key (written by ``train_baseline.py``
    as ``"baseline_cls_xyz"``; volume ckpts default to ``"dino_volume_scene"``).
    Returns ``kind`` so callers can branch their decode + viz paths.
    """
    ck = torch.load(str(path), map_location=device, weights_only=False)
    args = ck["args"]
    cfg = ck["dataset_cfg"]
    kind = ck.get("model_kind", "dino_volume_scene")
    # Backfill for ckpts written before train.py learned to emit the
    # volumetric model_kind — detect from saved args + state_dict.
    if kind == "dino_volume_scene" and args.get("volumetric"):
        kind = "dino_volume_scene_volumetric"
    if kind == "dino_volume_scene" and args.get("multiview"):
        kind = "dino_volume_scene_volumetric_multiview"
    if kind == "dino_volume_scene" and args.get("v3"):
        kind = "dino_volume_scene_v3"
    if kind == "dino_volume_scene" and args.get("v4"):
        kind = "dino_volume_scene_v4"
    if kind == "dino_volume_scene_volumetric_multiview" and args.get("v4"):
        kind = "dino_volume_scene_v4"
    if kind == "dino_volume_scene" and args.get("v4_da3"):
        kind = "dino_volume_scene_v4_da3"
    if kind == "dino_volume_scene" and args.get("v4_vlm"):
        kind = "dino_volume_scene_v4_vlm"
    if kind == "dino_volume_scene" and args.get("v4_dynaflip"):
        kind = "dino_volume_scene_v4_dynaflip"
    if kind == "dino_volume_scene" and args.get("v4_paligemma"):
        kind = "dino_volume_scene_v4_paligemma"
    if kind == "dino_volume_scene" and args.get("v4_pi0"):
        kind = "dino_volume_scene_v4_pi0"
    if kind == "dino_volume_scene" and args.get("v4_radio"):
        kind = "dino_volume_scene_v4_radio"
    if kind == "dino_volume_scene" and args.get("v4_clip"):
        kind = "dino_volume_scene_v4_clip"
    if kind == "baseline_cls_xyz":
        from .model_baseline import DinoCLSXYZBaseline
        model = DinoCLSXYZBaseline(
            n_window=args["n_window"], d_model=args["d_model"],
            n_blocks=args["n_blocks"],
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
        ).to(device).eval()
    elif kind == "baseline_cls_xyz_diffusion":
        from .model_baseline_diffusion import DinoCLSXYZBaselineDiffusion
        model = DinoCLSXYZBaselineDiffusion(
            n_window=args["n_window"], d_model=args["d_model"],
            n_blocks=args["n_blocks"],
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
            n_diffusion_steps_train=args.get("n_diffusion_steps", 100),
            n_diffusion_steps_infer=args.get("n_diffusion_steps", 100),
        ).to(device).eval()
    elif kind == "baseline_cls_uvz_motion_tracks":
        from .model_baseline_motion_tracks import DinoCLSUVZBaseline
        model = DinoCLSUVZBaseline(
            n_window=args["n_window"], d_model=args["d_model"],
            n_blocks=args["n_blocks"],
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
        ).to(device).eval()
    elif kind == "baseline_cls_xyz_multiview":
        from .model_baseline import DinoCLSXYZBaselineMultiView
        views_csv = args.get("multiview") or "scene"
        views = [v.strip() for v in views_csv.split(",") if v.strip()]
        model = DinoCLSXYZBaselineMultiView(
            views=views,
            n_window=args["n_window"], d_model=args["d_model"],
            n_blocks=args["n_blocks"],
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
            cross_view_layers=args.get("cross_view_layers", 0),
            cross_view_heads=args.get("cross_view_heads", 6),
        ).to(device).eval()
    elif kind == "dino_volume_per_voxel":
        from .model import DinoVolumeScenePerVoxel
        # Auto-detect feat_dim from the saved weights — feat_head.weight has
        # shape (feat_dim * n_height_bins, embed_dim, 1, 1). Robust to the
        # train.py args.feat_dim override bug (overrode local var but didn't
        # update args, so ckpt's args["feat_dim"] is the unhelpful CLI default).
        fh_out = ck["model"]["feat_head.weight"].shape[0]
        feat_dim_detected = fh_out // args["n_height_bins"]
        model = DinoVolumeScenePerVoxel(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=feat_dim_detected,
            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"],
        ).to(device).eval()
    elif kind == "da3_volume_scene":
        from .model_da3 import DA3VolumeScene
        model = DA3VolumeScene(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=args["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"],
        ).to(device).eval()
    elif kind == "dino_volume_scene_hires":
        from .model_hires import DinoVolumeSceneHires
        model = DinoVolumeSceneHires(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=args["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"],
        ).to(device).eval()
    elif kind == "dino_volume_per_voxel_hires":
        from .model_hires import DinoVolumeScenePerVoxelHires
        model = DinoVolumeScenePerVoxelHires(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=args["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"],
        ).to(device).eval()
    elif kind == "dino_volume_scene_hires_flat":
        from .model_hires import DinoVolumeSceneHiresFlat
        model = DinoVolumeSceneHiresFlat(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=args["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"],
        ).to(device).eval()
    elif kind == "dino_volume_scene_volumetric":
        from .model_volumetric import DinoVolumeSceneVolumetric
        # z_lo/z_hi come from the dataset_cfg's auto-derived range, persisted
        # at train time. Must pass them for the height embedding to map back
        # to physical heights at decode time.
        z_lo, z_hi = cfg["z_range"]
        model = DinoVolumeSceneVolumetric(
            n_window=args["n_window"],
            n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"],
            feat_dim=args["feat_dim"],
            height_enc_dim=args.get("height_enc_dim", 16),
            voxel_mlp_hidden=args.get("voxel_mlp_hidden", 64),
            d_model=args["d_model"],
            mlp_layers=args.get("mlp_layers", 3),
            mlp_hidden_mult=args.get("mlp_hidden_mult", 2.0),
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
            z_lo=z_lo, z_hi=z_hi,
            img_size=args.get("img_size", 504),
        ).to(device).eval()
    elif kind == "dino_volume_scene_volumetric_multiview":
        from .model_volumetric_multiview import DinoVolumeSceneVolumetricMultiView
        z_lo, z_hi = cfg["z_range"]
        views_csv = args.get("multiview") or "scene"
        views = [v.strip() for v in views_csv.split(",") if v.strip()]
        model = DinoVolumeSceneVolumetricMultiView(
            views=views,
            n_window=args["n_window"],
            n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"],
            feat_dim=args["feat_dim"],
            height_enc_dim=args.get("height_enc_dim", 16),
            voxel_mlp_hidden=args.get("voxel_mlp_hidden", 64),
            d_model=args["d_model"],
            mlp_layers=args.get("mlp_layers", 3),
            mlp_hidden_mult=args.get("mlp_hidden_mult", 2.0),
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
            z_lo=z_lo, z_hi=z_hi,
            img_size=args.get("img_size", 504),
            # Old ckpts (pre-2026-06-10) trained with primary-only CLS for the
            # grip/rot head; new default is concat-all-views.
            cls_fusion=args.get("cls_fusion", "primary"),
            cross_view_layers=args.get("cross_view_layers", 0),
            cross_view_heads=args.get("cross_view_heads", 6),
            per_voxel_grip_rot=args.get("per_voxel_grip_rot", False),
            # Old ckpts trained before the flat-MLP variant existed used the
            # 1D CNN. Default backfill respects that; new ckpts opt into flat.
            temporal_head=args.get("temporal_head", "temporal_cnn"),
            temporal_cnn_layers=args.get("temporal_cnn_layers", 3),
            temporal_cnn_kernel=args.get("temporal_cnn_kernel", 3),
            temporal_mlp_hidden=args.get("temporal_mlp_hidden", 512),
            past_n=args.get("mv_past_n", 0),
            past_enc_dim=args.get("mv_past_enc_dim", 64),
        ).to(device).eval()
    elif kind in ("dino_volume_scene_v4", "dino_volume_scene_v4_da3",
                    "dino_volume_scene_v4_vlm",
                    "dino_volume_scene_v4_dynaflip",
                    "dino_volume_scene_v4_paligemma",
                    "dino_volume_scene_v4_pi0",
                    "dino_volume_scene_v4_radio",
                    "dino_volume_scene_v4_clip"):
        z_lo, z_hi = cfg["z_range"]
        views_csv = args.get("multiview") or "scene,wrist"
        views = [v.strip() for v in views_csv.split(",") if v.strip()]
        v4_kwargs = dict(
            views=views,
            n_window=args["n_window"],
            n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"],
            feat_dim=args["feat_dim"],
            height_enc_dim=args.get("height_enc_dim", 16),
            voxel_mlp_hidden=args.get("voxel_mlp_hidden", 64),
            d_model=args["d_model"],
            mlp_layers=args.get("mlp_layers", 3),
            mlp_hidden_mult=args.get("mlp_hidden_mult", 2.0),
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
            z_lo=z_lo, z_hi=z_hi,
            img_size=args.get("img_size", 504),
            cls_fusion=args.get("cls_fusion", "concat"),
            cross_view_layers=args.get("cross_view_layers", 0),
            cross_view_heads=args.get("cross_view_heads", 6),
            per_voxel_grip_rot=args.get("per_voxel_grip_rot", True),
            temporal_head=args.get("temporal_head", "flat_mlp"),
            temporal_cnn_layers=args.get("temporal_cnn_layers", 3),
            temporal_cnn_kernel=args.get("temporal_cnn_kernel", 3),
            temporal_mlp_hidden=args.get("temporal_mlp_hidden", 512),
            past_n=args.get("mv_past_n", 8),
            past_enc_dim=args.get("mv_past_enc_dim", 64),
            cam_margin_m=args.get("cam_margin_m", 0.01),
        )
        if kind == "dino_volume_scene_v4_da3":
            from .model_volumetric_v4_da3 import DinoVolumeSceneV4DA3
            model = DinoVolumeSceneV4DA3(**v4_kwargs).to(device).eval()
        elif kind == "dino_volume_scene_v4_vlm":
            from .model_volumetric_v4_vlm import DinoVolumeSceneV4VLM
            v4_kwargs["task_text_default"] = args.get(
                "vlm_task_text", "pick and place")
            model = DinoVolumeSceneV4VLM(**v4_kwargs).to(device).eval()
        elif kind == "dino_volume_scene_v4_dynaflip":
            from .model_volumetric_v4_dynaflip import DinoVolumeSceneV4DynaFLIP
            model = DinoVolumeSceneV4DynaFLIP(**v4_kwargs).to(device).eval()
        elif kind == "dino_volume_scene_v4_paligemma":
            from .model_volumetric_v4_paligemma import DinoVolumeSceneV4PaliGemma
            v4_kwargs["paligemma_model_name"] = args.get(
                "paligemma_model_name", "google/paligemma-3b-pt-448")
            v4_kwargs["task_text_default"] = args.get(
                "vlm_task_text", "pick and place")
            model = DinoVolumeSceneV4PaliGemma(**v4_kwargs).to(device).eval()
        elif kind == "dino_volume_scene_v4_pi0":
            from .model_volumetric_v4_pi0 import DinoVolumeSceneV4Pi0
            # pi0 MUST use the 224² PaliGemma — position-embed table has
            # 256=16² entries and lerobot/pi0_base was fine-tuned on that
            # geometry. Ignore ckpt-saved `paligemma_model_name` (paligemma
            # trainer defaults to 448 which mismatches on load).
            v4_kwargs["paligemma_model_name"] = "google/paligemma-3b-pt-224"
            v4_kwargs["task_text_default"] = args.get(
                "vlm_task_text", "pick and place")
            model = DinoVolumeSceneV4Pi0(**v4_kwargs).to(device).eval()
        elif kind == "dino_volume_scene_v4_radio":
            from .model_volumetric_v4_radio import DinoVolumeSceneV4Radio
            model = DinoVolumeSceneV4Radio(**v4_kwargs).to(device).eval()
        elif kind == "dino_volume_scene_v4_clip":
            from .model_volumetric_v4_clip import DinoVolumeSceneV4Clip
            model = DinoVolumeSceneV4Clip(**v4_kwargs).to(device).eval()
        else:
            from .model_volumetric_v4 import DinoVolumeSceneV4
            model = DinoVolumeSceneV4(**v4_kwargs).to(device).eval()
    elif kind == "dino_volume_scene_v3":
        from .model_volumetric_v3 import DinoVolumeSceneV3
        z_lo, z_hi = cfg["z_range"]
        views_csv = args.get("multiview") or "scene,wrist"
        views = [v.strip() for v in views_csv.split(",") if v.strip()]
        model = DinoVolumeSceneV3(
            views=views,
            n_window=args["n_window"],
            n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"],
            F_z=args.get("v3_F_z", 16),
            F_t=args.get("v3_F_t", 16),
            F_s=args.get("v3_F_s", 32),
            d_model=args["d_model"],
            mlp_layers=args.get("mlp_layers", 3),
            mlp_hidden_mult=args.get("mlp_hidden_mult", 2.0),
            n_gripper_bins=args["n_gripper_bins"],
            n_rot_clusters=args["n_rot_clusters"],
            z_lo=z_lo, z_hi=z_hi,
            img_size=args.get("img_size", 504),
            past_n=args.get("mv_past_n", 8),
            past_enc_dim=args.get("mv_past_enc_dim", 64),
            temporal_mlp_hidden=args.get("temporal_mlp_hidden", 512),
        ).to(device).eval()
    elif kind == "dino_volume_scene_hires_ablate":
        from .model_hires import DinoVolumeSceneHiresAblate
        model = DinoVolumeSceneHiresAblate(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=args["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"],
            img_size=args.get("img_size", 504),
            past_n=args.get("past_n", 0),
            past_z_delta=args.get("past_z_delta", False),
            past_use_uv=args.get("past_use_uv", False),
            past_uv_delta=args.get("past_uv_delta", False),
            past_use_xyz=args.get("past_use_xyz", False),
            past_xyz_delta=args.get("past_xyz_delta", False),
            mlp_layers=args.get("mlp_layers", 1),
            mlp_residual=args.get("mlp_residual", False),
            mlp_hidden_mult=args.get("mlp_hidden_mult", 2.0),
            backbone_layer_fusion=args.get("backbone_layer_fusion", False),
        ).to(device).eval()
    else:
        model = DinoVolumeScene(
            n_window=args["n_window"], n_height_bins=args["n_height_bins"],
            pred_size=args["pred_size"], feat_dim=args["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"],
        ).to(device).eval()
    model.load_state_dict(ck["model"])
    return {"model": model, "cfg": cfg, "args": args, "step": ck["step"],
            "kind": kind}


@torch.no_grad()
def decode_baseline_motion_tracks(
        model, rgb: torch.Tensor,
        K_in: torch.Tensor, T_w2c: torch.Tensor, cfg: dict) -> dict:
    """Decode :class:`DinoCLSUVZBaseline` output to the same shape as
    :func:`decode_baseline`.

    The model regresses (u, v, world-z) per waypoint in the SCENE-CAM image
    frame. To get world-XYZ for the controller we ray-plane-intersect each
    UV at the predicted Z. Requires the live scene-cam intrinsic + extrinsic
    (both at the same image size the model was trained on).
    """
    from .model_baseline_motion_tracks import uvz_to_world_xyz
    out = model(rgb)
    uvz = out["uvz_pred"]                                       # (1, T, 3)
    xyz = uvz_to_world_xyz(uvz, K_in.unsqueeze(0), T_w2c.unsqueeze(0))
    xyz = xyz[0].cpu().numpy()                                  # (T, 3)
    grip = out["grip_logits"][0]
    rot = out["rot_logits"][0]
    g_bin = grip.argmax(dim=-1).cpu().numpy()
    g_lo, g_hi = cfg["grip_range"]
    grip_val = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / cfg["n_gripper_bins"]
    rot_bin = rot.argmax(dim=-1).cpu().numpy()
    quat = cfg["rot_centroids"][rot_bin]
    return {
        "xyz_world": xyz.astype(np.float32),
        "grip_bin": g_bin.astype(np.int64),
        "grip": grip_val.astype(np.float32),
        "rot_bin": rot_bin.astype(np.int64),
        "quat": quat,
        "raw_out": out,
    }


@torch.no_grad()
def decode_baseline(model, rgb: torch.Tensor, cfg: dict) -> dict:
    """Decode ``DinoCLSXYZBaseline`` output to the same shape as ``decode``.

    Baseline directly regresses world-frame EE position (no volume, no
    unproject, no workspace mask). Grip + rot heads decode the same way
    as the volume model. Returns the same fields the rest of deploy.py
    consumes (``xyz_world``, ``quat``, ``grip``), minus volume-specific
    fields (``pix_uv``, ``z_bin``, ``height_m``).
    """
    out = model(rgb)
    xyz = out["xyz_pred"][0].cpu().numpy()                     # (T, 3) world
    grip = out["grip_logits"][0]
    rot = out["rot_logits"][0]

    g_bin = grip.argmax(dim=-1).cpu().numpy()
    g_lo, g_hi = cfg["grip_range"]
    grip_val = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / cfg["n_gripper_bins"]

    rot_bin = rot.argmax(dim=-1).cpu().numpy()
    centroids = cfg["rot_centroids"]                          # (K, 4) np.float32
    quat = centroids[rot_bin]

    return {
        "xyz_world": xyz.astype(np.float32),
        "grip_bin": g_bin.astype(np.int64),
        "grip": grip_val.astype(np.float32),
        "rot_bin": rot_bin.astype(np.int64),
        "quat": quat,
        "raw_out": out,
    }


@torch.no_grad()
def decode_volumetric(model, rgb: torch.Tensor, start_pix: torch.Tensor,
                       cfg: dict, K_in: np.ndarray,
                       T_cam_in_world: np.ndarray) -> dict:
    """Decode for ``DinoVolumeSceneVolumetric``.

    The model's forward needs ``(K_in, T_w2c)`` so it can unproject the
    image-feature grid into world voxels. The argmax over the
    ``(T, Z, P, P)`` volume picks one voxel per timestep, and that voxel's
    pre-computed world position IS the predicted EE XYZ — no separate
    ray-cast step. ``height_m`` is just the z component.

    Same return schema as ``decode`` so deploy_panels.log_volume_panels
    works without changes.
    """
    device = rgb.device
    K_t = torch.from_numpy(K_in.astype(np.float32)).unsqueeze(0).to(device)
    T_w2c = np.linalg.inv(T_cam_in_world).astype(np.float32)
    T_t = torch.from_numpy(T_w2c).unsqueeze(0).to(device)
    out = model(rgb, start_pix, K_t, T_t)
    vol = out["volume_logits"][0]                             # (T, Z, P, P)
    voxel_pos = out["voxel_positions"][0]                     # (Z, P, P, 3)
    grip = out["grip_logits"][0]
    rot = out["rot_logits"][0]
    T, Z, P, _ = vol.shape

    flat = vol.reshape(T, -1).argmax(dim=-1).cpu().numpy()
    z_bin = flat // (P * P)
    yx_flat = flat % (P * P)
    y_idx = yx_flat // P
    x_idx = yx_flat % P

    # Image pixel of the predicted voxel (cell center, in img_size coords).
    cell = cfg["img_size"] / P
    y_pix = (y_idx + 0.5) * cell
    x_pix = (x_idx + 0.5) * cell
    pix_uv = np.stack([x_pix, y_pix], axis=1).astype(np.float32)

    # World XYZ — direct lookup, NO ray-cast.
    vp_np = voxel_pos.cpu().numpy()
    xyz_world = vp_np[z_bin, y_idx, x_idx]                    # (T, 3)
    height_m = xyz_world[:, 2].astype(np.float32)

    g_bin = grip.argmax(dim=-1).cpu().numpy()
    g_lo, g_hi = cfg["grip_range"]
    grip_val = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / cfg["n_gripper_bins"]

    rot_bin = rot.argmax(dim=-1).cpu().numpy()
    centroids = cfg["rot_centroids"]
    quat = centroids[rot_bin]

    return {
        "pix_uv": pix_uv,
        "z_bin": z_bin.astype(np.int64),
        "height_m": height_m,
        "grip_bin": g_bin.astype(np.int64),
        "grip": grip_val.astype(np.float32),
        "rot_bin": rot_bin.astype(np.int64),
        "quat": quat,
        "xyz_world": xyz_world.astype(np.float32),
        "raw_out": out,
    }


@torch.no_grad()
def decode_volumetric_multiview(
    model, rgb_list: list[torch.Tensor], start_pix: torch.Tensor, cfg: dict,
    K_list: list[np.ndarray], T_cam_in_world_list: list[np.ndarray],
    border_mask_px: int = 5,
    past_grip: torch.Tensor | None = None,
    past_z: torch.Tensor | None = None,
) -> dict:
    """Multi-view sibling of ``decode_volumetric``.

    Each list element corresponds to one entry in ``model.views`` (same
    order). The model returns a combined ``volume_logits`` of shape
    ``(T, N*Z, P, P)`` and ``voxel_positions`` of shape ``(N*Z, P, P, 3)``;
    we argmax flat and look up the world position directly — same as the
    single-view path.

    ``border_mask_px`` zeroes a frame of that width around the pred-grid
    edge of EVERY view's slab before argmax. Per-pixel features at the
    borders get little training supervision and tend to accumulate
    spurious probability mass at inference (OOD object placements
    sometimes latch onto the image edge). Set to 0 to disable.
    """
    device = rgb_list[0].device
    K_t = [torch.from_numpy(K.astype(np.float32)).unsqueeze(0).to(device)
            for K in K_list]
    T_t = [torch.from_numpy(np.linalg.inv(T).astype(np.float32))
                .unsqueeze(0).to(device)
            for T in T_cam_in_world_list]
    # Border mask is applied INSIDE the model so the per-voxel grip/rot
    # argmax-gather and the downstream xyz argmax see the same masked
    # volume. The model returns the already-masked ``volume_logits``.
    out = model(rgb_list, start_pix, K_t, T_t,
                 border_mask_px=int(border_mask_px),
                 past_grip=past_grip, past_z=past_z)
    vol = out["volume_logits"][0]                                # (T, NZ, P, P)
    voxel_pos = out["voxel_positions"][0]                         # (NZ, P, P, 3)
    grip = out["grip_logits"][0]
    rot = out["rot_logits"][0]
    T, NZ, P, _ = vol.shape

    flat = vol.reshape(T, -1).argmax(dim=-1).cpu().numpy()
    nz_idx = flat // (P * P)
    yx = flat % (P * P)
    y_idx = yx // P; x_idx = yx % P

    # Pixel coords are reported in the PRIMARY view (views[0]) frame for
    # any voxel whose nz_idx falls into the primary slab; for voxels lifted
    # from a non-primary view the pixel meaning is that view's grid — viz
    # downstream just needs *some* pixel so log_volume_panels' overlay
    # works. We use the primary slab's cell center for voxels in slab 0
    # and nan for the rest (panel code should already skip nan).
    cell = cfg["img_size"] / P
    y_pix = (y_idx + 0.5) * cell
    x_pix = (x_idx + 0.5) * cell
    pix_uv = np.stack([x_pix, y_pix], axis=1).astype(np.float32)
    # Mark non-primary-slab predictions with NaN so callers know not to
    # overlay them on the primary image.
    n_views = max(1, len(out.get("views", [None])))
    in_primary = (nz_idx < (NZ // n_views))
    pix_uv[~in_primary] = np.nan

    vp_np = voxel_pos.cpu().numpy()
    xyz_world = vp_np[nz_idx, y_idx, x_idx]                       # (T, 3)
    height_m = xyz_world[:, 2].astype(np.float32)

    g_bin = grip.argmax(dim=-1).cpu().numpy()
    g_lo, g_hi = cfg["grip_range"]
    grip_val = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / cfg["n_gripper_bins"]

    rot_bin = rot.argmax(dim=-1).cpu().numpy()
    centroids = cfg["rot_centroids"]
    quat = centroids[rot_bin]

    return {
        "pix_uv": pix_uv,
        "z_bin": nz_idx.astype(np.int64),
        "height_m": height_m,
        "grip_bin": g_bin.astype(np.int64),
        "grip": grip_val.astype(np.float32),
        "rot_bin": rot_bin.astype(np.int64),
        "quat": quat,
        "xyz_world": xyz_world.astype(np.float32),
        "raw_out": out,
    }


@torch.no_grad()
def decode_v3(model, rgb_list: list[torch.Tensor], cfg: dict,
                K_list: list[np.ndarray], T_cam_in_world_list: list[np.ndarray],
                past_grip: torch.Tensor | None = None,
                past_z: torch.Tensor | None = None,
                border_mask_px: int = 5) -> dict:
    """Decode for the v3 factorized model.

    Same return schema as ``decode_volumetric_multiview`` so the deploy
    panel code can consume it unchanged.
    """
    device = rgb_list[0].device
    K_t = [torch.from_numpy(K.astype(np.float32)).unsqueeze(0).to(device)
            for K in K_list]
    T_t = [torch.from_numpy(np.linalg.inv(T).astype(np.float32))
                .unsqueeze(0).to(device)
            for T in T_cam_in_world_list]
    out = model(rgb_list, K_in=K_t, T_w2c=T_t,
                 past_grip=past_grip, past_z=past_z,
                 border_mask_px=int(border_mask_px))

    xyz_world = out["chosen_xyz_world"][0].cpu().numpy()                 # (T, 3)
    grip = out["grip_logits"][0]
    rot = out["rot_logits"][0]
    T = xyz_world.shape[0]

    height_m = xyz_world[:, 2].astype(np.float32)

    g_bin = grip.argmax(dim=-1).cpu().numpy()
    g_lo, g_hi = cfg["grip_range"]
    grip_val = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / cfg["n_gripper_bins"]

    rot_bin = rot.argmax(dim=-1).cpu().numpy()
    centroids = cfg["rot_centroids"]
    quat = centroids[rot_bin]

    # Primary-view pixel coords (cell-centered).
    P = int(out["pred_size"])
    cell = cfg["img_size"] / P
    chosen_view = out["chosen_view"][0].cpu().numpy()
    chosen_y = out["chosen_y"][0].cpu().numpy()
    chosen_x = out["chosen_x"][0].cpu().numpy()
    y_pix = (chosen_y + 0.5) * cell
    x_pix = (chosen_x + 0.5) * cell
    pix_uv = np.stack([x_pix, y_pix], axis=1).astype(np.float32)
    in_primary = (chosen_view == 0)
    pix_uv[~in_primary] = np.nan

    return {
        "pix_uv": pix_uv,
        "z_bin": out["chosen_z"][0].cpu().numpy().astype(np.int64),
        "height_m": height_m,
        "grip_bin": g_bin.astype(np.int64),
        "grip": grip_val.astype(np.float32),
        "rot_bin": rot_bin.astype(np.int64),
        "quat": quat,
        "xyz_world": xyz_world.astype(np.float32),
        "raw_out": out,
    }


def unproject_pixel_to_world(
    pix_uv: np.ndarray, height_m: np.ndarray,
    K_in: np.ndarray, T_cam_in_world: np.ndarray,
) -> np.ndarray:
    """Per-T (u, v, world-Z) → world (x, y, z) — same formula as the
    pre-refactor ``unproject_voxel_to_world`` in deploy_yam_2view.py.

    Ray from cam through pixel, intersected with the world-Z plane at
    ``height_m``. Pixel coords should be **cell centers** (use
    ``(bin + 0.5) * cell_size`` upstream).
    """
    K_inv = np.linalg.inv(K_in)
    out = np.zeros((len(pix_uv), 3), dtype=np.float32)
    for t in range(len(pix_uv)):
        ray_cam = K_inv @ np.array([pix_uv[t, 0], pix_uv[t, 1], 1.0])
        h = float(height_m[t])
        denom = T_cam_in_world[2, :3] @ ray_cam
        if abs(denom) < 1e-6:
            out[t] = np.array([np.nan, np.nan, h])
            continue
        d_cam = (h - T_cam_in_world[2, 3]) / denom
        p_cam = ray_cam * d_cam
        out[t] = (T_cam_in_world @ np.append(p_cam, 1.0))[:3].astype(np.float32)
    return out


def compute_workspace_mask(
    K_in: np.ndarray,
    T_cam_in_world: np.ndarray,
    T_lfr: np.ndarray,
    bin_centers_z: np.ndarray,
    pred_size: int,
    image_size: int,
    margin_robot: float = 0.05,
    near: float = 0.05,
) -> torch.Tensor:
    """(Z, P, P) bool — same logic as ``compute_workspace_mask`` in
    deploy_yam_2view.py:919. True = voxel is in front of the camera AND
    in front of both arm bases (so the EE can physically be there).
    Cameron's "ignore predictions behind the robot base" rule lives here.
    """
    Z = len(bin_centers_z)
    cell = image_size / pred_size
    px = (np.arange(pred_size) + 0.5) * cell
    py = (np.arange(pred_size) + 0.5) * cell
    PX, PY = np.meshgrid(px, py, indexing="xy")
    K_inv = np.linalg.inv(K_in)
    pix_h = np.stack([PX, PY, np.ones_like(PX)], axis=-1)            # (P, P, 3)
    ray_dir_cam = pix_h @ K_inv.T                                    # (P, P, 3)
    T_row2 = T_cam_in_world[2, :3]
    T23 = T_cam_in_world[2, 3]
    denom = ray_dir_cam @ T_row2                                     # (P, P)
    safe = np.where(np.abs(denom) > 1e-6, denom, np.sign(denom + 1e-9) * 1e-6)
    d_at_bin = (np.asarray(bin_centers_z)[:, None, None] - T23) / safe[None]  # (Z, P, P)
    T_world2cam = np.linalg.inv(T_cam_in_world)
    cam_z_left = float(T_world2cam[2, 3])
    cam_z_right = float(T_world2cam[2, :3] @ T_lfr[:3, 3] + T_world2cam[2, 3])
    cam_z_robot = max(cam_z_left, cam_z_right) + margin_robot
    valid = (d_at_bin > near) & (d_at_bin < cam_z_robot) & (np.abs(denom)[None] > 1e-6)
    return torch.from_numpy(valid).bool()


@torch.no_grad()
def decode(model: DinoVolumeScene, rgb: torch.Tensor, start_pix: torch.Tensor,
            cfg: dict,
            K_in: np.ndarray = None,
            T_cam_in_world: np.ndarray = None,
            T_lfr: np.ndarray = None,
            border_mask_px: int = 4,
            past_features: dict = None) -> dict:
    """rgb (1, 3, S, S), start_pix (1, 2). Returns dict of (T, …) numpy arrays.

    When ``K_in`` + ``T_cam_in_world`` + ``T_lfr`` are all provided:
    1. **Workspace mask** is applied (``vol.masked_fill(~ws, -1e9)``)
       before the argmax, blocking voxels behind the robot base or
       behind the camera. Matches the pre-refactor ``decode_all_steps``.
    2. ``xyz_world`` is included (ray-cast at predicted-height plane,
       cell-center pixel coords).

    ``border_mask_px`` zeroes a frame of that width around the pred-grid
    edge before argmax. Per-pixel features at the borders get little
    training supervision (GTs rarely land there) and tend to accumulate
    spurious probability mass at inference. Set to 0 to disable.
    """
    if past_features is not None:
        out = model(rgb, start_pix, **past_features)
    else:
        out = model(rgb, start_pix)
    vol = out["volume_logits"][0]                   # (T, Z, P, P)
    grip = out["grip_logits"][0]                     # (T, n_grip)
    rot = out["rot_logits"][0]                       # (T, n_rot)
    T, Z, P, _ = vol.shape

    # ── workspace mask (only when caller provides geometry) ────────
    z_lo, z_hi = cfg["z_range"]
    bin_centers_z = z_lo + (np.arange(cfg["n_z"]) + 0.5) * (z_hi - z_lo) / cfg["n_z"]
    if K_in is not None and T_cam_in_world is not None and T_lfr is not None:
        ws = compute_workspace_mask(
            K_in, T_cam_in_world, T_lfr, bin_centers_z,
            pred_size=cfg["pred_size"], image_size=cfg["img_size"],
        ).to(vol.device)
        vol_masked = vol.masked_fill(~ws.unsqueeze(0), -1e9)
        n_valid = int(ws.sum()); n_total = int(ws.numel())
        if n_valid < 1000:
            print(f"  ws_mask warn: only {n_valid}/{n_total} valid voxels "
                  f"(extrinsic may be wrong)")
    else:
        vol_masked = vol

    # ── border mask (always applied; cheap) ───────────────────────
    if border_mask_px > 0 and border_mask_px * 2 < P:
        border_keep = torch.zeros(P, P, dtype=torch.bool, device=vol.device)
        border_keep[border_mask_px:P - border_mask_px,
                    border_mask_px:P - border_mask_px] = True
        vol_masked = vol_masked.masked_fill(
            ~border_keep.view(1, 1, P, P), -1e9)

    flat = vol_masked.reshape(T, -1).argmax(dim=-1).cpu().numpy()
    z_bin = flat // (P * P)
    yx_flat = flat % (P * P)
    # CELL-CENTER pixel coords (matches OLD unproject_voxel_to_world).
    cell = cfg["img_size"] / P
    y_pix = ((yx_flat // P) + 0.5) * cell
    x_pix = ((yx_flat % P) + 0.5) * cell
    pix_uv = np.stack([x_pix, y_pix], axis=1)        # (T, 2) xy

    height_m = bin_centers_z[z_bin]                  # (T,) directly index centers

    g_bin = grip.argmax(dim=-1).cpu().numpy()
    g_lo, g_hi = cfg["grip_range"]
    grip_val = g_lo + (g_bin + 0.5) * (g_hi - g_lo) / cfg["n_gripper_bins"]

    rot_bin = rot.argmax(dim=-1).cpu().numpy()
    centroids = cfg["rot_centroids"]                 # (K, 4) np.float32
    quat = centroids[rot_bin]                         # (T, 4)

    result = {
        "pix_uv": pix_uv.astype(np.float32),
        "z_bin": z_bin.astype(np.int64),
        "height_m": height_m.astype(np.float32),
        "grip_bin": g_bin.astype(np.int64),
        "grip": grip_val.astype(np.float32),
        "rot_bin": rot_bin.astype(np.int64),
        "quat": quat,
        "raw_out": out,                              # for viz reuse
    }
    if K_in is not None and T_cam_in_world is not None:
        result["xyz_world"] = unproject_pixel_to_world(
            pix_uv, height_m, K_in, T_cam_in_world)
    return result