"""Render 4 overlays for a dataset frame using the saved poses (no PnP round-trip).

Given a state npz that already carries scene_cam_pose / wrist_cam_pose / umi_pose (all
local→world 4x4 in the calib-board GridBoard frame, OpenCV camera convention), we drive
the render_core.Layer camera directly instead of re-detecting boards. Produces:

  A. <task>/f<N>_scene_umi.png       — UMI gripper rendered from scene camera
  B. <task>/f<N>_scene_calib.png     — scene calib board rendered from scene camera
  C. <task>/f<N>_wrist_calib.png     — scene calib board rendered from wrist camera
  D. <task>/f<N>_wrist_umi.png       — UMI gripper rendered from wrist camera, chained
                                       through a pre-calibrated wrist_cam→UMI transform.

The wrist_cam→UMI transform (constant, since the wrist camera is rigidly attached to the
UMI gripper) is recovered once from a chosen "calibration frame" where all three boards
were detected, then cached to `wrist_umi_calibration.json`.

Usage:
    MUJOCO_GL=osmesa python render_dataset_frame.py \
        --dataset-root /data/cameron/puget_ar_datasets/testing_cube_in_bowl \
        --frame 28 --calib-frame 184 \
        --out /data/cameron/ar_renders/testing_cube_in_bowl
"""
from __future__ import annotations
import argparse, json, os, sys
from pathlib import Path
import numpy as np
import cv2

os.environ.setdefault("MUJOCO_GL", "osmesa")

AR_VIEW = Path(__file__).parent.resolve()
sys.path.insert(0, str(AR_VIEW))

from render_core import Layer  # noqa: E402
from aruco_boards import UMI_V2  # noqa: E402
from calib_board import CALIB_BOARD  # noqa: E402


CALIB_PATH = AR_VIEW / "wrist_umi_calibration.json"
SCENE_MODEL = AR_VIEW / "calib_board_model" / "calib_board.xml"
UMI_MODEL   = Path("/data/cameron/cad_recovery/mj_umi_v2/umi_gripper_v2.xml")


def _fovy_deg(K_save: np.ndarray, H_render: int, H_save: int) -> float:
    """K_save is the intrinsics matched to the saved (500-wide) JPEGs.
       Since fovy is resolution-independent (fovy = 2 atan(H/(2 fy))), we can render at any
       size using the same fy scaled to that render height."""
    # fy at the SAVE resolution:
    fy_save = float(K_save[1, 1])
    # fy scales linearly with vertical resolution — but H_render == H_save when we render at
    # the composite's height, so just use fy_save directly.
    if H_render == H_save:
        fy = fy_save
    else:
        fy = fy_save * (H_render / H_save)
    return 2.0 * np.degrees(np.arctan(H_render / (2.0 * fy)))


CYAN = (0, 255, 255)
MAGENTA = (255, 0, 255)


def _blend(real_bgr: np.ndarray, sim_bgr: np.ndarray, weight: float = 0.9,
           outline_color=CYAN) -> np.ndarray:
    """Composite the model onto the real image. Where the sim has drawn ANY pixel we blend
       heavily (default 90% sim) so the overlay is unmistakable; elsewhere the real image
       is preserved. A cyan silhouette outline is drawn along the sim mask so the align
       target is easy to see even when the sim is dark.
    """
    lum = sim_bgr.max(axis=2).astype(np.float32) / 255.0
    mask = (lum > 0.02)
    alpha = np.zeros_like(lum)
    alpha[mask] = weight  # binary mask * weight — no soft falloff, so aligned areas are unambiguous
    a = alpha[..., None]
    out = real_bgr.astype(np.float32) * (1 - a) + sim_bgr.astype(np.float32) * a
    out = np.clip(out, 0, 255).astype(np.uint8)
    # Silhouette outline along the mask boundary
    mask_u8 = (mask.astype(np.uint8)) * 255
    edges = cv2.morphologyEx(mask_u8, cv2.MORPH_GRADIENT,
                             cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)))
    out[edges > 0] = outline_color
    return out


def _hide_floors(layer):
    """Make every plane-type geom in the loaded model invisible (rgba alpha → 0) so the
       ambient floor/checker plane doesn't leak into the render — only real model meshes
       (like the calib board or the UMI gripper) survive on a black background."""
    import mujoco
    for g in range(layer.model.ngeom):
        if int(layer.model.geom_type[g]) == int(mujoco.mjtGeom.mjGEOM_PLANE):
            layer.model.geom_rgba[g] = [0, 0, 0, 0]


def _render_clean(layer, T_camCV_w, fovy, H, W):
    """Layer.render() but with skybox + floor + reflections stripped so the sim background
       is black and only the model's dynamic geoms (the mesh we care about) survive."""
    import mujoco
    from render_core import set_cam
    set_cam(layer.model, layer.extcam, T_camCV_w, fovy)
    mujoco.mj_forward(layer.model, layer.data)
    r = layer.renderers.get((H, W))
    if r is None:
        r = mujoco.Renderer(layer.model, H, W)
        layer.renderers[(H, W)] = r
    r.update_scene(layer.data, camera=layer.extcam)
    # Then strip environment/render flags on the built scene
    scn = r.scene
    for flg in (mujoco.mjtRndFlag.mjRND_SKYBOX,
                mujoco.mjtRndFlag.mjRND_HAZE,
                mujoco.mjtRndFlag.mjRND_REFLECTION,
                mujoco.mjtRndFlag.mjRND_SHADOW,
                mujoco.mjtRndFlag.mjRND_FOG):
        scn.flags[flg] = 0
    return cv2.cvtColor(r.render(), cv2.COLOR_RGB2BGR)


def _label(img: np.ndarray, text: str) -> np.ndarray:
    cv2.rectangle(img, (0, 0), (img.shape[1], 24), (0, 0, 0), -1)
    cv2.putText(img, text, (6, 17), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
    return img


def _project(P_world: np.ndarray, T_cam_world: np.ndarray, K: np.ndarray) -> tuple[int, int, float]:
    """Project a world point into the image. T_cam_world is (camera→world), 4x4, OpenCV convention.
       Returns (u, v, z_cam). z_cam < 0 means behind camera."""
    T_world_cam = np.linalg.inv(T_cam_world)
    P_cam = T_world_cam @ np.array([P_world[0], P_world[1], P_world[2], 1.0])
    z = float(P_cam[2])
    if z <= 1e-6:
        return -1, -1, z
    uv = K @ (P_cam[:3] / z)
    return int(round(uv[0])), int(round(uv[1])), z


# UMI_V2 board is 4 cols × 3 rows of 20 mm markers with 5 mm gaps → 95 × 70 mm bounding box.
# cv2.aruco.GridBoard's frame origin sits at the corner of marker 150 (the first id), so raw
# umi_pose[:3, 3] is at a board corner, not the visible gripper center. UMI_BOARD_CENTER
# is that fixed offset (in umi/board local frame) needed to move the reference to the board
# center — a much more meaningful "gripper anchor" for training.
UMI_BOARD_W = 4 * 0.020 + 3 * 0.005   # 0.095 m
UMI_BOARD_H = 3 * 0.020 + 2 * 0.005   # 0.070 m
UMI_BOARD_CENTER = np.array([UMI_BOARD_W / 2.0, UMI_BOARD_H / 2.0, 0.0])


def _keypoint_world(umi_pose: np.ndarray, ref: str = "center") -> np.ndarray:
    """Return the world-frame reference point for a UMI pose.
       'corner' → cv2 GridBoard origin (buggy behaviour, kept for A/B comparison).
       'center' → center of the physical UMI ArUco board (much closer to the gripper)."""
    if ref == "corner":
        return umi_pose[:3, 3].astype(float)
    local = np.array([*UMI_BOARD_CENTER, 1.0])
    return (umi_pose @ local)[:3]


def _draw_future_traj(base_bgr: np.ndarray, umi_poses: list[np.ndarray], visible: list[bool],
                      T_cam_world: np.ndarray, K: np.ndarray, ref: str = "center") -> np.ndarray:
    """Polyline of UMI origins across `umi_poses`. Fading yellow→red, alpha 0.9→0.3."""
    H, W = base_bgr.shape[:2]
    out = base_bgr.copy()
    pts = []
    for i, (T, vis) in enumerate(zip(umi_poses, visible)):
        u, v, z = _project(_keypoint_world(T, ref), T_cam_world, K)
        if z <= 0 or u < 0 or v < 0 or u >= W or v >= H:
            pts.append(None)
            continue
        pts.append((u, v, vis))
    # segments
    for i in range(len(pts) - 1):
        a, b = pts[i], pts[i + 1]
        if a is None or b is None:
            continue
        t = i / max(1, len(pts) - 1)
        color = (int(80 + 175 * t), int(200 - 60 * t), 255)  # BGR: shift toward orange over time
        cv2.line(out, a[:2], b[:2], color, 2, cv2.LINE_AA)
    # dots
    for i, p in enumerate(pts):
        if p is None:
            continue
        u, v, vis = p
        t = i / max(1, len(pts) - 1)
        r = 4 if i == 0 else 3
        color = (60, 220, 255) if i == 0 else (int(80 + 175 * t), int(200 - 60 * t), 255)
        cv2.circle(out, (u, v), r, color, -1, cv2.LINE_AA)
        # dim ring on stale (umi not visible in that frame — pose reuses previous)
        if not vis:
            cv2.circle(out, (u, v), r + 2, (60, 60, 60), 1, cv2.LINE_AA)
    return out


def _draw_axes_triads(base_bgr: np.ndarray, umi_poses: list[np.ndarray], visible: list[bool],
                      T_cam_world: np.ndarray, K: np.ndarray, axis_len_m: float = 0.02,
                      ref: str = "center") -> np.ndarray:
    """3-axis triad (X=red, Y=green, Z=blue) drawn at every UMI pose. Line alpha fades over time
       so the current frame is bright and the horizon fades out."""
    H, W = base_bgr.shape[:2]
    out = base_bgr.copy()
    N = len(umi_poses)
    for i, (T, vis) in enumerate(zip(umi_poses, visible)):
        origin = _keypoint_world(T, ref)
        R = T[:3, :3]
        u0, v0, z0 = _project(origin, T_cam_world, K)
        if z0 <= 0 or u0 < 0 or v0 < 0 or u0 >= W or v0 >= H:
            continue
        t = i / max(1, N - 1)
        alpha = 0.9 - 0.55 * t
        for axis, color_full in ((0, (60, 60, 255)),   # X - red   (BGR)
                                  (1, (80, 220, 80)),   # Y - green
                                  (2, (255, 180, 60))): # Z - blue-ish (BGR: light blue)
            end_world = origin + R[:, axis] * axis_len_m
            u1, v1, z1 = _project(end_world, T_cam_world, K)
            if z1 <= 0 or u1 < 0 or v1 < 0 or u1 >= W or v1 >= H:
                continue
            color = tuple(int(c * alpha) for c in color_full)
            cv2.line(out, (u0, v0), (u1, v1), color, 1 if i > 0 else 2, cv2.LINE_AA)
        # small dot on origin (dim if the UMI was not actually observed that frame)
        cv2.circle(out, (u0, v0), 2 if i > 0 else 3,
                   (200, 200, 200) if vis else (100, 100, 100), -1, cv2.LINE_AA)
    return out


def _hstack(cells: list[np.ndarray], target_hw: tuple[int, int] = (281, 500),
            gap: int = 4, bg=(0, 0, 0)) -> np.ndarray:
    """Stack `cells` in a single horizontal row, aspect-fit each cell to `target_hw`."""
    H, W = target_hw
    def fit(img):
        ih, iw = img.shape[:2]
        s = min(H / ih, W / iw)
        nh, nw = int(round(ih * s)), int(round(iw * s))
        r = cv2.resize(img, (nw, nh))
        c = np.full((H, W, 3), bg, np.uint8)
        c[(H - nh) // 2:(H - nh) // 2 + nh, (W - nw) // 2:(W - nw) // 2 + nw] = r
        return c
    fitted = [fit(c) for c in cells]
    total_w = len(fitted) * W + (len(fitted) + 1) * gap
    canvas = np.full((H + 2 * gap, total_w, 3), bg, np.uint8)
    x = gap
    for c in fitted:
        canvas[gap:gap + H, x:x + W] = c
        x += W + gap
    return canvas


def _grid_2col(cells: list[np.ndarray], target_hw: tuple[int, int] = (281, 500),
               gap: int = 6, bg=(0, 0, 0)) -> np.ndarray:
    """Stack `cells` into a 2-column grid, aspect-fit each cell into `target_hw` with letterbox."""
    H, W = target_hw
    def fit(img):
        ih, iw = img.shape[:2]
        s = min(H / ih, W / iw)
        nh, nw = int(round(ih * s)), int(round(iw * s))
        r = cv2.resize(img, (nw, nh))
        c = np.full((H, W, 3), bg, np.uint8)
        c[(H - nh) // 2:(H - nh) // 2 + nh, (W - nw) // 2:(W - nw) // 2 + nw] = r
        return c
    fitted = [fit(c) for c in cells]
    rows = (len(fitted) + 1) // 2
    grid_h = rows * H + (rows + 1) * gap
    grid_w = 2 * W + 3 * gap
    canvas = np.full((grid_h, grid_w, 3), bg, np.uint8)
    for i, c in enumerate(fitted):
        r, col = i // 2, i % 2
        y = gap + r * (H + gap)
        x = gap + col * (W + gap)
        canvas[y:y + H, x:x + W] = c
    return canvas


def _load_state(state_dir: Path, frame: int) -> dict:
    p = state_dir / f"{frame:06d}.npz"
    if not p.exists():
        raise FileNotFoundError(p)
    z = np.load(p)
    return {k: z[k] for k in z.files}


def _load_frame(root: Path, subdir: str, frame: int) -> np.ndarray:
    p = root / subdir / f"{frame:06d}.jpg"
    img = cv2.imread(str(p))
    if img is None:
        raise FileNotFoundError(p)
    return img


def _K_save(meta: dict, cam: str) -> np.ndarray:
    return np.asarray(meta[cam]["K_save"], dtype=float)


def _project_to_SO3(M: np.ndarray) -> np.ndarray:
    """Nearest rotation matrix in Frobenius norm (Markley's method): SVD, drop determinant flip."""
    U, _, Vt = np.linalg.svd(M)
    d = np.sign(np.linalg.det(U @ Vt))
    return U @ np.diag([1.0, 1.0, d]) @ Vt


def _pose_median_filter(Ts: list[np.ndarray], k: float = 3.0) -> list[int]:
    """Return indices of poses whose translation is within k*MAD of the median translation.
       Coarse outlier rejection — throws away frames where the underlying PnP obviously drifted."""
    ts = np.array([T[:3, 3] for T in Ts])
    med = np.median(ts, axis=0)
    dev = np.linalg.norm(ts - med, axis=1)
    mad = np.median(dev) + 1e-9
    return [i for i, d in enumerate(dev) if d <= k * mad]


def _load_or_calibrate_wrist_umi(dataset_root: Path, calib_frame: int, force: bool) -> np.ndarray:
    """Return T_wrist_cam→umi (4x4).

    Strategy:
      1. If cached and not forced, return the cached value.
      2. Otherwise sweep every state npz in the dataset, keep only frames where
         scene/wrist/umi are ALL visible (so both PnP solves are trustworthy), compute
         T_wc_umi(f) = inv(umi_pose(f)) @ wrist_cam_pose(f).
      3. MAD-filter outliers on translation, then chordal-average the rotations via SVD
         and take the mean translation. That reduces PnP noise from a single frame by ~√N
         and rejects bad PnP solves (small board area / grazing angle).
      4. If nothing usable is found, fall back to the single `calib_frame`.
    """
    if CALIB_PATH.exists() and not force:
        d = json.loads(CALIB_PATH.read_text())
        return np.asarray(d["T_wrist_cam_umi"], dtype=float)

    state_dir = dataset_root / "state"
    per_frame: list[tuple[int, np.ndarray]] = []
    for p in sorted(state_dir.glob("*.npz")):
        st = np.load(p)
        try:
            if bool(st["scene_visible"]) and bool(st["wrist_visible"]) and bool(st["umi_visible"]):
                T = np.linalg.inv(st["umi_pose"]) @ st["wrist_cam_pose"]
                per_frame.append((int(p.stem), T))
        except KeyError:
            pass

    if not per_frame:
        print(f"[calibration] no frames with all three boards visible — falling back to f{calib_frame}", flush=True)
        st = _load_state(state_dir, calib_frame)
        T = np.linalg.inv(st["umi_pose"]) @ st["wrist_cam_pose"]
        CALIB_PATH.write_text(json.dumps({
            "T_wrist_cam_umi": T.tolist(),
            "recovered_from": {"dataset": str(dataset_root), "frame": int(calib_frame), "n_frames": 1},
        }, indent=2))
        return T

    frames, Ts = zip(*per_frame)
    keep = _pose_median_filter(list(Ts))
    Ts_keep = [Ts[i] for i in keep]
    frames_keep = [frames[i] for i in keep]

    Rsum = np.zeros((3, 3))
    tsum = np.zeros(3)
    for T in Ts_keep:
        Rsum += T[:3, :3]
        tsum += T[:3, 3]
    R_mean = _project_to_SO3(Rsum)
    t_mean = tsum / len(Ts_keep)
    T_mean = np.eye(4)
    T_mean[:3, :3] = R_mean
    T_mean[:3, 3]  = t_mean

    # Diagnostics: translation spread after filtering
    ts = np.array([T[:3, 3] for T in Ts_keep])
    t_std = float(np.linalg.norm(ts - t_mean, axis=1).std())
    ang_dev = []
    for T in Ts_keep:
        Rd = T[:3, :3] @ R_mean.T
        cos = np.clip((np.trace(Rd) - 1) / 2, -1.0, 1.0)
        ang_dev.append(np.degrees(np.arccos(cos)))
    ang_std = float(np.std(ang_dev))

    print(f"[calibration] {len(per_frame)} tri-visible frames → {len(Ts_keep)} after MAD filter", flush=True)
    print(f"[calibration] residual std: translation={t_std*1000:.2f} mm, rotation={ang_std:.2f}°", flush=True)

    CALIB_PATH.write_text(json.dumps({
        "T_wrist_cam_umi": T_mean.tolist(),
        "recovered_from": {
            "dataset": str(dataset_root),
            "method": "multi-frame chordal average, MAD-filtered on translation",
            "n_frames_considered": len(per_frame),
            "n_frames_kept": len(Ts_keep),
            "frames_kept": frames_keep[:20],  # first 20 for reference
            "translation_std_mm": round(t_std * 1000, 3),
            "rotation_std_deg": round(ang_std, 3),
        },
    }, indent=2))
    print(f"[calibration] wrote {CALIB_PATH}", flush=True)
    return T_mean


def _pose_umi_jaw(layer_umi, gripper_rad: float, joint_name: str = "jaw"):
    """Write the servo-derived jaw angle into the UMI model's jaw joint before rendering.
       gripper_rad comes from state/*.npz and already carries the calibration jaw_sign, so
       it can be written to qpos directly (MuJoCo will clip to the joint range on step)."""
    import mujoco
    if not np.isfinite(gripper_rad):
        return
    jid = mujoco.mj_name2id(layer_umi.model, mujoco.mjtObj.mjOBJ_JOINT, joint_name)
    if jid < 0:
        return
    adr = int(layer_umi.model.jnt_qposadr[jid])
    layer_umi.data.qpos[adr] = float(gripper_rad)
    mujoco.mj_forward(layer_umi.model, layer_umi.data)


def _render_one(frame: int, args, meta, T_wc_umi, K_scene, K_wrist,
                fovy_scene, fovy_wrist, Hs, Ws, Hw, Ww,
                layer_scene, layer_umi, ref: str = "center") -> np.ndarray:
    """Render a single frame → horizontal-stack composite (scene | wrist | future).
       Returns the BGR composite ready to write."""
    scene = _load_frame(args.dataset_root, "scene", frame)
    wrist = _load_frame(args.dataset_root, "wrist", frame)
    st = _load_state(args.dataset_root / "state", frame)

    scene_cam_pose = st["scene_cam_pose"]
    wrist_cam_pose = st["wrist_cam_pose"]
    umi_pose       = st["umi_pose"]

    _pose_umi_jaw(layer_umi, float(st.get("gripper_rad", np.nan)))

    # Scene: calib (cyan outline) + UMI (magenta outline)
    sim_scene_calib = _render_clean(layer_scene, layer_scene.T_board_w @ scene_cam_pose, fovy_scene, Hs, Ws)
    sim_scene_umi   = _render_clean(layer_umi,   layer_umi.T_board_w @ np.linalg.inv(umi_pose) @ scene_cam_pose,
                                    fovy_scene, Hs, Ws)
    scene_combo = _blend(scene, sim_scene_calib, outline_color=CYAN)
    scene_combo = _blend(scene_combo, sim_scene_umi, outline_color=MAGENTA)
    scene_combo = _label(scene_combo,
                         f"scene · calib+UMI  f{frame}  "
                         f"sv={int(st['scene_visible'])} uv={int(st['umi_visible'])} "
                         f"grip={float(st.get('gripper_rad', float('nan'))):+.2f}")

    # Wrist: calib (cyan) + UMI (magenta, chained)
    sim_wrist_calib = _render_clean(layer_scene, layer_scene.T_board_w @ wrist_cam_pose, fovy_wrist, Hw, Ww)
    sim_wrist_umi   = _render_clean(layer_umi,   layer_umi.T_board_w @ T_wc_umi, fovy_wrist, Hw, Ww)
    wrist_combo = _blend(wrist, sim_wrist_calib, outline_color=CYAN)
    wrist_combo = _blend(wrist_combo, sim_wrist_umi, outline_color=MAGENTA)
    wrist_combo = _label(wrist_combo,
                         f"wrist · calib+UMI (chained)  f{frame}  wv={int(st['wrist_visible'])}")

    # Future window: next `horizon` UMI poses projected as trajectory + triads on scene view
    horizon = int(args.horizon)
    total = int(meta["num_frames"])
    fut_frames = list(range(frame, min(frame + horizon + 1, total)))
    umi_poses = []
    umi_vis = []
    for f in fut_frames:
        stf = _load_state(args.dataset_root / "state", f)
        umi_poses.append(np.asarray(stf["umi_pose"], dtype=float))
        umi_vis.append(bool(stf["umi_visible"]))
    future_scene = _draw_axes_triads(scene, umi_poses, umi_vis, scene_cam_pose, K_scene, ref=ref)
    future_scene = _draw_future_traj(future_scene, umi_poses, umi_vis, scene_cam_pose, K_scene, ref=ref)
    future_scene = _label(future_scene,
                          f"scene · next {len(fut_frames)-1} keypoints+axes  f{frame}  [ref={ref}]")

    # Wrist view "future" panel — same trajectory + triads projected into the CURRENT wrist
    # camera. Wrist camera moves rigidly with the UMI, so we derive its pose from the current
    # UMI pose + frozen T_wc_umi (avoids the stale wrist_visible=False issue).
    wrist_cam_derived = umi_pose @ T_wc_umi
    future_wrist = _draw_axes_triads(wrist, umi_poses, umi_vis, wrist_cam_derived, K_wrist, ref=ref)
    future_wrist = _draw_future_traj(future_wrist, umi_poses, umi_vis, wrist_cam_derived, K_wrist, ref=ref)
    future_wrist = _label(future_wrist,
                          f"wrist · next {len(fut_frames)-1} keypoints+axes  f{frame}  [ref={ref}]")

    return _hstack([scene_combo, wrist_combo, future_scene, future_wrist], target_hw=(281, 500))


def _parse_frames_arg(spec: str, total: int) -> list[int]:
    if spec == "all":
        return list(range(total))
    if "-" in spec:
        a, b = spec.split("-", 1)
        return list(range(int(a), int(b) + 1))
    return [int(x) for x in spec.split(",") if x.strip()]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset-root", required=True, type=Path,
                    help="dataset root (has meta.json, scene/, wrist/, state/)")
    ap.add_argument("--frames", type=str, default=None,
                    help='which frames to render — "all" | "28" | "0-100" | "28,50,90"')
    ap.add_argument("--frame", type=int, default=None,
                    help="convenience alias for --frames <N>")
    ap.add_argument("--calib-frame", type=int, default=184,
                    help="frame with all three boards visible for wrist_cam→umi recovery")
    ap.add_argument("--out", required=True, type=Path, help="output directory")
    ap.add_argument("--force-calib", action="store_true",
                    help="ignore cached wrist_cam→umi and re-derive")
    ap.add_argument("--horizon", type=int, default=30,
                    help="how many future frames to draw in the trajectory + axes viz")
    ap.add_argument("--name-fmt", type=str, default="f{frame:06d}.jpg",
                    help="output file naming template (jpg default so the scrubber loads fast)")
    ap.add_argument("--jpeg-quality", type=int, default=92)
    ap.add_argument("--ref", choices=["corner", "center"], default="center",
                    help="which UMI point to project as the 'gripper keypoint': "
                         "'corner' = raw cv2 GridBoard origin (marker 150 corner — off the visible gripper), "
                         "'center' = board center (much closer to gripper body). default center.")
    args = ap.parse_args()

    args.out.mkdir(parents=True, exist_ok=True)
    meta = json.loads((args.dataset_root / "meta.json").read_text())
    total = int(meta["num_frames"])

    if args.frames is None and args.frame is None:
        ap.error("supply --frames or --frame")
    frames = _parse_frames_arg(args.frames, total) if args.frames else [int(args.frame)]
    frames = [f for f in frames if 0 <= f < total]
    print(f"[batch] rendering {len(frames)} frame(s): "
          f"{frames if len(frames) <= 8 else str(frames[:4])+' … '+str(frames[-4:])}", flush=True)

    T_wc_umi = _load_or_calibrate_wrist_umi(args.dataset_root, args.calib_frame, args.force_calib)

    K_scene = _K_save(meta, "scene")
    K_wrist = _K_save(meta, "wrist")
    # Probe shapes once — they're constant per dataset
    probe = _load_frame(args.dataset_root, "scene", frames[0])
    Hs, Ws = probe.shape[:2]
    probe = _load_frame(args.dataset_root, "wrist", frames[0])
    Hw, Ww = probe.shape[:2]
    fovy_scene = _fovy_deg(K_scene, Hs, Hs)
    fovy_wrist = _fovy_deg(K_wrist, Hw, Hw)

    print("[layer] loading scene calib_board.xml + measure_board_world (loop-closure) …", flush=True)
    layer_scene = Layer("scene", str(SCENE_MODEL), CALIB_BOARD, "calib", lambda jn: None)
    _hide_floors(layer_scene)
    print(f"[layer] scene reproj={layer_scene.reproj}", flush=True)

    print("[layer] loading umi_gripper_v2.xml + measure_board_world …", flush=True)
    layer_umi = Layer("umi", str(UMI_MODEL), UMI_V2, "ar_back", lambda jn: None)
    _hide_floors(layer_umi)
    print(f"[layer] umi reproj={layer_umi.reproj}", flush=True)

    import time
    t0 = time.time()
    for i, frame in enumerate(frames):
        img = _render_one(frame, args, meta, T_wc_umi, K_scene, K_wrist,
                          fovy_scene, fovy_wrist, Hs, Ws, Hw, Ww,
                          layer_scene, layer_umi, ref=args.ref)
        out_path = args.out / args.name_fmt.format(frame=frame)
        if out_path.suffix.lower() in (".jpg", ".jpeg"):
            cv2.imwrite(str(out_path), img, [int(cv2.IMWRITE_JPEG_QUALITY), args.jpeg_quality])
        else:
            cv2.imwrite(str(out_path), img)
        if (i + 1) % 25 == 0 or i + 1 == len(frames):
            dt = time.time() - t0
            print(f"[batch] {i + 1}/{len(frames)} done  ({dt:.1f}s, {dt/(i+1)*1000:.0f} ms/frame)", flush=True)
    print(f"[done] output dir: {args.out}")


if __name__ == "__main__":
    main()