"""AR-view (feetech UMI) scene+wrist dataset for smithbot_v4 training.

Mirrors the interface of ``yam_local_train_mirror/lib/dataset.py:YamScene``
so ``train.py`` needs only the import swap. Reads ``record_dataset.py``
output (a flat ``state/*.npz`` + ``scene/*.jpg`` + ``wrist/*.jpg`` tree)
plus ``episodes.json`` for demonstration segmentation.

Schema references:
- ``/data/cameron/claude_feetech_controller/ar_view/DATASET_FORMAT.md``
- ``/data/cameron/claude_feetech_controller/ar_view/TRAINING_DATA_README.md``

Per-sample output matches YamScene exactly (train.py + loss + viz eat
the same tensors) — see YamScene docstring for target semantics.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
import torch
from torch.utils.data import Dataset

# Re-use everything that isn't robot-specific from yams's dataset module.
from .dataset import (
    IMAGENET_MEAN, IMAGENET_STD,
    ROT_KMEANS_K,
    _imagenet_norm_chw,
    _R_to_quat, _quat_canonical,
    _kmeans_quats, _nearest_centroid_idx,
    _project_world_to_uv,
)

# Where the frozen wrist_cam → umi transform lives (fleet + puget mirror).
_WRIST_UMI_CAL_PATHS = (
    "/data/cameron/claude_feetech_controller/ar_view/wrist_umi_calibration.json",
    str(Path.home() / "ar_view/wrist_umi_calibration.json"),
)

# Constant T_board→tcp offset derived from cad's UMI MuJoCo model via
# render_core.Layer.T_board_w (see preprocess/derive_tcp_offset_v2.py).
# Training target = TCP fingertip position, NOT the ArUco board center.
_TCP_OFFSET_PATHS = (
    "/data/cameron/repos/smithbot_v4/preprocess/tcp_offset.json",
    str(Path.home() / "smithbot_v4/preprocess/tcp_offset.json"),
)


def _load_tcp_offset() -> Optional[np.ndarray]:
    """Load the constant 4x4 T_board→tcp. Returns None if unavailable — the
    dataset then falls back to using umi_pose[:3,3] as the target (board center)."""
    for p in _TCP_OFFSET_PATHS:
        try:
            d = json.loads(Path(p).read_text())
        except Exception:
            continue
        T = np.asarray(d["T_board_tcp"], dtype=np.float64)
        if T.shape == (4, 4):
            return T
    return None

# AR-view auto-derive fallback if a task doesn't have a manual gripper range.
_ARVIEW_GRIP_RANGE_PAD = 0.05


def _load_wrist_umi_calibration() -> Optional[np.ndarray]:
    """Load the frozen 4x4 T_wrist_cam→umi (see TRAINING_DATA_README.md §wrist calibration).
    Returns None if the calibration file isn't available — callers should
    then fall back to per-frame wrist_cam_pose from state npz."""
    for p in _WRIST_UMI_CAL_PATHS:
        try:
            d = json.loads(Path(p).read_text())
        except Exception:
            continue
        T = np.asarray(d["T_wrist_cam_umi"], dtype=np.float64)
        if T.shape == (4, 4):
            return T
    return None


def _read_arview_scene_frame(root: Path, frame_idx: int):
    p = root / "scene" / f"{frame_idx:06d}.jpg"
    bgr = cv2.imread(str(p))
    if bgr is None:
        raise RuntimeError(f"scene frame missing: {p}")
    return bgr


def _read_arview_wrist_frame(root: Path, frame_idx: int):
    p = root / "wrist" / f"{frame_idx:06d}.jpg"
    bgr = cv2.imread(str(p))
    return bgr


def _median_pose_over_visible(poses: np.ndarray, visible: np.ndarray) -> np.ndarray:
    """Given (N, 4, 4) camera poses and (N,) visibility bools, return a
    single fixed 4x4. Uses per-element median across the visible subset
    (crude but adequate for a physically-fixed camera whose per-frame PnP
    solve is noisy). Falls back to poses[0] if nothing is visible."""
    mask = np.asarray(visible, dtype=bool)
    if not mask.any():
        return poses[0].astype(np.float64)
    med = np.median(poses[mask], axis=0).astype(np.float64)
    # Re-orthonormalize the rotation block via SVD.
    R = med[:3, :3]
    U, _, Vt = np.linalg.svd(R)
    R_ortho = U @ Vt
    if np.linalg.det(R_ortho) < 0:
        U[:, -1] *= -1
        R_ortho = U @ Vt
    out = np.eye(4, dtype=np.float64)
    out[:3, :3] = R_ortho
    out[:3, 3] = med[:3, 3]
    return out


class ArViewScene(Dataset):
    """AR-view (feetech UMI) training samples. World frame = calib board.

    Reads:
      ``root/state/NNNNNN.npz`` — per-frame poses + gripper (schema in DATASET_FORMAT.md)
      ``root/scene/NNNNNN.jpg`` + ``root/wrist/NNNNNN.jpg`` — 500-px images
      ``root/meta.json`` — intrinsics + FPS + camera conventions
      ``root/scene_wrist/episodes.json`` OR ``root/episodes.json`` — segmentation

    Uses ``umi_pose`` as the end-effector world pose (no FK, unlike YAM).
    Wrist camera pose per frame uses ``wrist_cam_pose`` when
    ``wrist_visible`` is true, else falls back to
    ``umi_pose(f) @ inv(T_wrist_cam→umi)`` from the frozen calibration.
    """

    def __init__(
        self,
        root: Path,
        n_window: int = 32,
        frame_stride: int = 1,          # AR-view is already 4 fps; no further stride
        pred_size: int = 128,
        img_size: int = 504,
        n_z: int = 32,
        n_gripper_bins: int = 32,
        n_rot_clusters: int = ROT_KMEANS_K,
        z_range: Optional[tuple[float, float]] = None,
        grip_range: Optional[tuple[float, float]] = None,
        rot_centroids: Optional[np.ndarray] = None,
        z_pad: float = 0.02,
        past_n: int = 0,
        episodes_json: Optional[Path] = None,
    ):
        self.root = Path(root)
        self.n_window = n_window
        self.frame_stride = frame_stride
        self.pred_size = pred_size
        self.img_size = img_size
        self.n_z = n_z
        self.n_gripper_bins = n_gripper_bins
        self.n_rot_clusters = n_rot_clusters
        self.past_n = int(past_n)

        # ── meta.json ──────────────────────────────────────────────────
        meta = json.loads((self.root / "meta.json").read_text())
        assert meta.get("camera_convention", "").startswith("opencv"), \
            "AR-view meta.json missing opencv camera convention"
        n_frames_total = int(meta["num_frames"])
        scene = meta["scene"]
        wrist = meta.get("wrist", {})
        W_save = int(scene["save_width"])
        # Height matches native aspect ratio; read from a real JPEG.
        bgr0 = _read_arview_scene_frame(self.root, 0)
        H_save, W_verify = bgr0.shape[:2]
        assert W_verify == W_save, \
            f"scene JPEG width {W_verify} != meta.save_width {W_save}"
        K_scene_save = np.asarray(scene["K_save"], dtype=np.float64)
        K_wrist_save = (np.asarray(wrist["K_save"], dtype=np.float64)
                        if wrist and wrist.get("K_save") else None)
        # Rescale intrinsics from save resolution → model input resolution.
        # (Non-square: use the actual save H/W independently for cx/cy scale.)
        def _scale_K(K_src, src_w, src_h, dst):
            K = K_src.copy()
            K[0, 0] *= dst / src_w
            K[0, 2] *= dst / src_w
            K[1, 1] *= dst / src_h
            K[1, 2] *= dst / src_h
            return K
        K_scene_in = _scale_K(K_scene_save, W_save, H_save, img_size)
        K_wrist_in = (_scale_K(K_wrist_save, W_save, H_save, img_size)
                      if K_wrist_save is not None else None)

        # ── Episodes.json ──────────────────────────────────────────────
        # Prefer the explicit path, else the two conventional locations.
        ej_path = None
        candidates = [episodes_json] if episodes_json is not None else []
        candidates += [self.root / "scene_wrist" / "episodes.json",
                        self.root / "episodes.json"]
        for c in candidates:
            if c is not None and Path(c).exists():
                ej_path = Path(c)
                break
        if ej_path is None:
            raise RuntimeError(f"no episodes.json in {self.root} (checked {candidates})")
        ej = json.loads(ej_path.read_text())
        raw_eps = ej.get("episodes", [])
        if not raw_eps:
            raise RuntimeError(f"no episodes listed in {ej_path}")

        # ── Preload all state npzs ONCE (272 frames * ~8 arrays ≈ trivial) ──
        state = []
        for i in range(n_frames_total):
            d = np.load(self.root / "state" / f"{i:06d}.npz")
            state.append({
                "scene_cam_pose": np.asarray(d["scene_cam_pose"], dtype=np.float64),
                "wrist_cam_pose": np.asarray(d["wrist_cam_pose"], dtype=np.float64),
                "umi_pose":       np.asarray(d["umi_pose"],       dtype=np.float64),
                "scene_visible":  bool(d["scene_visible"]),
                "wrist_visible":  bool(d["wrist_visible"]),
                "umi_visible":    bool(d["umi_visible"]),
                "gripper_rad":    float(d["gripper_rad"]) if not np.isnan(d["gripper_rad"]) else np.nan,
            })

        # Frozen wrist_cam→umi (optional but preferred over per-frame stale poses).
        T_wc2umi = _load_wrist_umi_calibration()
        self._T_wrist_cam_to_umi = T_wc2umi   # may be None
        if T_wc2umi is None:
            print("  WARN: wrist_umi_calibration.json not found; falling back to "
                  "per-frame wrist_cam_pose (may be stale on non-visible frames)")

        # Frozen T_board→tcp offset — training targets the fingertip TCP,
        # not the ArUco board center. Derived from cad's UMI MuJoCo model.
        self._T_board_tcp = _load_tcp_offset()
        if self._T_board_tcp is None:
            print("  WARN: tcp_offset.json not found; falling back to board-center "
                  "targets (worse for grasping tasks — fingertip is 21cm from board).")
        else:
            print(f"  using TCP offset (board→tcp translation: "
                  f"{self._T_board_tcp[:3, 3].round(3).tolist()} m, "
                  f"|={np.linalg.norm(self._T_board_tcp[:3, 3]):.3f}m)")

        # ── Build per-episode metadata ─────────────────────────────────
        self.episodes: list[dict] = []
        for raw_ep in raw_eps:
            e_start = int(raw_ep["start"])
            e_end = int(raw_ep["end"])          # inclusive per schema
            n_ep = e_end - e_start + 1
            if n_ep < 2:
                # Need at least one future frame beyond start.
                continue
            # Fix scene camera pose = median of visible per-frame solves within
            # this episode (physically the scene camera is stationary).
            scene_poses_all = np.stack([state[i]["scene_cam_pose"]
                                          for i in range(e_start, e_end + 1)], axis=0)
            scene_vis_all = np.array([state[i]["scene_visible"]
                                       for i in range(e_start, e_end + 1)], dtype=bool)
            T_scene_w = _median_pose_over_visible(scene_poses_all, scene_vis_all)
            T_w2c_scene = np.linalg.inv(T_scene_w)

            # Per-frame arrays for this episode.
            # ee_world/quat targets the TCP (fingertip) when the offset is
            # available, else the board center as a fallback.
            ee_world = np.zeros((n_ep, 3), dtype=np.float64)
            quats = np.zeros((n_ep, 4), dtype=np.float32)
            grips = np.zeros(n_ep, dtype=np.float32)
            umi_pose_all = np.zeros((n_ep, 4, 4), dtype=np.float64)
            grip_valid = np.ones(n_ep, dtype=bool)
            for k, i in enumerate(range(e_start, e_end + 1)):
                T_board = state[i]["umi_pose"]
                umi_pose_all[k] = T_board
                if self._T_board_tcp is not None:
                    T_tcp = T_board @ self._T_board_tcp
                else:
                    T_tcp = T_board
                ee_world[k] = T_tcp[:3, 3]
                quats[k] = _quat_canonical(_R_to_quat(T_tcp[:3, :3])).astype(np.float32)
                g = state[i]["gripper_rad"]
                if np.isnan(g):
                    grip_valid[k] = False
                    grips[k] = 0.0
                else:
                    grips[k] = float(g)

            self.episodes.append({
                "ep_id": raw_ep.get("id", f"ep_{e_start}_{e_end}"),
                "frame0": e_start,       # absolute frame index of local frame 0
                "n": n_ep,
                "K_in": K_scene_in.astype(np.float32),
                "K_wrist_in": K_wrist_in.astype(np.float32) if K_wrist_in is not None else None,
                "T_w2c": T_w2c_scene.astype(np.float32),
                "T_lfr": np.eye(4, dtype=np.float32),   # world already = calib board
                "ee_world": ee_world,
                "quats": quats,
                "grips": grips,
                "grip_valid": grip_valid,
                "umi_pose": umi_pose_all,
            })

        if not self.episodes:
            raise RuntimeError(
                f"no usable episodes in {self.root} (n_window={n_window} "
                f"× frame_stride={frame_stride} requires ep length ≥ {n_window*frame_stride + 1})")

        # ── z_range: auto-derive if not passed ──────────────────────────
        if z_range is None:
            all_z = np.concatenate([ep["ee_world"][:, 2] for ep in self.episodes])
            z_lo = float(all_z.min()) - z_pad
            z_hi = float(all_z.max()) + z_pad
            print(f"  auto z_range: [{z_lo:.4f}, {z_hi:.4f}]m  "
                  f"(data [{all_z.min():.3f}, {all_z.max():.3f}], pad ±{z_pad}m)")
        else:
            z_lo, z_hi = z_range
            print(f"  z_range from caller: [{z_lo:.4f}, {z_hi:.4f}]m")
        self.z_lo, self.z_hi = z_lo, z_hi

        # ── grip_range: auto-derive if not passed ───────────────────────
        if grip_range is None:
            all_g = np.concatenate([ep["grips"][ep["grip_valid"]]
                                     for ep in self.episodes])
            grip_lo = float(all_g.min()) - _ARVIEW_GRIP_RANGE_PAD
            grip_hi = float(all_g.max()) + _ARVIEW_GRIP_RANGE_PAD
            print(f"  auto grip_range: [{grip_lo:.4f}, {grip_hi:.4f}]rad  "
                  f"(data [{all_g.min():.3f}, {all_g.max():.3f}])")
        else:
            grip_lo, grip_hi = grip_range
        self.grip_lo, self.grip_hi = grip_lo, grip_hi

        # ── K-means rotation centroids ──────────────────────────────────
        if rot_centroids is not None:
            assert rot_centroids.shape == (n_rot_clusters, 4)
            self.rot_centroids = rot_centroids.astype(np.float32)
            print(f"  using rot K-means (K={n_rot_clusters}) from caller")
        else:
            all_q = np.concatenate([ep["quats"] for ep in self.episodes], axis=0)
            print(f"  fitting K-means (K={n_rot_clusters}, deterministic seed=0) "
                  f"on {len(all_q)} quaternions…")
            # If we have fewer quaternions than clusters, cap K to n_samples.
            K_eff = min(n_rot_clusters, len(all_q))
            if K_eff < n_rot_clusters:
                print(f"  clamping K from {n_rot_clusters} → {K_eff} "
                      f"(only {len(all_q)} training quats available)")
            self.rot_centroids = _kmeans_quats(all_q, K_eff)
            # Pad up to n_rot_clusters by repeating the first centroid so downstream
            # dimensioning is preserved (unused indices simply never get selected).
            if K_eff < n_rot_clusters:
                pad = np.tile(self.rot_centroids[:1],
                              (n_rot_clusters - K_eff, 1)).astype(np.float32)
                self.rot_centroids = np.concatenate([self.rot_centroids, pad], axis=0)

        # ── Per-frame bin precompute (same as YamScene) ─────────────────
        for ep in self.episodes:
            ep["rot_bin"] = _nearest_centroid_idx(ep["quats"], self.rot_centroids)
            g = (ep["grips"] - self.grip_lo) / (self.grip_hi - self.grip_lo)
            ep["grip_bin"] = np.clip(g * n_gripper_bins, 0,
                                       n_gripper_bins - 1).astype(np.int64)
            z = (ep["ee_world"][:, 2] - self.z_lo) / (self.z_hi - self.z_lo)
            ep["z_bin"] = np.clip(z * n_z, 0, n_z - 1).astype(np.int64)
            ee_h = np.concatenate([ep["ee_world"],
                                     np.ones((ep["n"], 1))], axis=1)
            ee_cam = (ep["T_w2c"] @ ee_h.T).T[:, :3]
            uv = (ep["K_in"] @ ee_cam.T).T
            uv = uv[:, :2] / np.maximum(uv[:, 2:3], 1e-6)
            scale = pred_size / img_size
            yx = np.stack([np.clip(uv[:, 1] * scale, 0, pred_size - 1),
                            np.clip(uv[:, 0] * scale, 0, pred_size - 1)], axis=1)
            ep["yx_bin"] = np.floor(yx).astype(np.int64)
            ep["start_pix_img"] = uv.astype(np.float32)

        # ── Sample index (ep, local_start_frame) ────────────────────────
        # Iterate over every frame that has at least one valid future frame;
        # __getitem__ clips future indices to [0, ep["n"] - 1] so short
        # episodes just "hold last target" for the tail of the window
        # (matches yams's past-features clip pattern at dataset.py:419).
        self.entries: list[tuple[int, int]] = []
        for ei, ep in enumerate(self.episodes):
            for i in range(ep["n"] - 1):
                self.entries.append((ei, i))
        if not self.entries:
            raise RuntimeError("no samples; all episodes are 0- or 1-frame")
        print(f"  built {len(self.entries)} samples across {len(self.episodes)} "
              f"episodes (n_window={n_window}, frame_stride={frame_stride}; "
              f"future frames clipped to episode end)")

    def __len__(self) -> int:
        return len(self.entries)

    def config(self) -> dict:
        return {
            "n_window": self.n_window,
            "frame_stride": self.frame_stride,
            "pred_size": self.pred_size,
            "img_size": self.img_size,
            "n_z": self.n_z,
            "n_gripper_bins": self.n_gripper_bins,
            "n_rot_clusters": self.n_rot_clusters,
            "z_range": (self.z_lo, self.z_hi),
            "grip_range": (self.grip_lo, self.grip_hi),
            "rot_centroids": self.rot_centroids,
        }

    def _wrist_pose_at(self, ep_idx: int, local_frame: int) -> Optional[np.ndarray]:
        """4x4 wrist camera pose in world. Uses the frozen wrist_cam→umi
        transform composed with umi_pose(f) when available (robust to
        wrist_visible=False frames), else the raw per-frame wrist_cam_pose."""
        ep = self.episodes[ep_idx]
        if self._T_wrist_cam_to_umi is not None:
            umi_pose = ep["umi_pose"][local_frame]
            # Empirical verification 2026-07-14: with the TCP offset applied,
            # projecting TCP world → wrist pixel using this composition matches
            # the big-TCP sphere render (which Cameron approved) at (188, 86)
            # within 10-pixel centroid noise. The alternate `umi @ inv(T_wc_umi)`
            # puts TCP at (326, 333) — off-frame — and does not match the render.
            # See preprocess/kp_projection_tcp_viz.py + kp_diff_viz_bigtcp.
            T_world_wrist_cam = umi_pose @ self._T_wrist_cam_to_umi
            return T_world_wrist_cam
        # Fallback: raw per-frame stored wrist_cam_pose (already world-frame).
        abs_frame = ep["frame0"] + local_frame
        # Note: this branch requires the state array — we saved umi_pose
        # into ep but not wrist_cam_pose per-frame. If we hit this we need
        # to reload from disk. Simpler: rely on the frozen calibration path.
        raise RuntimeError(
            "Frozen wrist_umi_calibration missing and per-frame wrist pose "
            "not cached; add T_wrist_cam→umi calibration or extend "
            "ArViewScene.__init__ to persist wrist_cam_pose per episode.")

    def __getitem__(self, idx: int):
        ei, start = self.entries[idx]
        ep = self.episodes[ei]
        abs_start = ep["frame0"] + start
        bgr = _read_arview_scene_frame(self.root, abs_start)
        rgb = _imagenet_norm_chw(bgr, self.img_size)

        s = self.frame_stride
        # Clip future frame indices to episode range; short episodes get
        # "hold last target" padding — the tail of the window replays the
        # final frame's target, teaching the model to sit still at demo end.
        frames = np.clip(
            np.arange(1, self.n_window + 1, dtype=np.int64) * s + start,
            0, ep["n"] - 1,
        )

        out = {
            "rgb": rgb,
            "start_pix": torch.from_numpy(ep["start_pix_img"][start]),
            "target_yx": torch.from_numpy(ep["yx_bin"][frames]),
            "target_z":  torch.from_numpy(ep["z_bin"][frames]),
            "target_grip": torch.from_numpy(ep["grip_bin"][frames]),
            "target_rot": torch.from_numpy(ep["rot_bin"][frames]),
            "target_xyz": torch.from_numpy(
                ep["ee_world"][frames].astype(np.float32)),
            "_bgr_src": bgr,
            "_K_in": torch.from_numpy(ep["K_in"]),
            "_T_w2c": torch.from_numpy(ep["T_w2c"]),
        }

        # Wrist view — same key names as YamScene, so train.py's
        # multi-view handling picks it up automatically.
        if ep["K_wrist_in"] is not None and self._T_wrist_cam_to_umi is not None:
            w_bgr = _read_arview_wrist_frame(self.root, abs_start)
            if w_bgr is not None:
                out["wrist_rgb"] = _imagenet_norm_chw(w_bgr, self.img_size)
                out["_wrist_bgr_src"] = w_bgr
                T_world_wrist_cam = self._wrist_pose_at(ei, start)
                T_w2c_wrist = np.linalg.inv(T_world_wrist_cam)
                out["_K_wrist"] = torch.from_numpy(ep["K_wrist_in"])
                out["_T_w2c_wrist"] = torch.from_numpy(
                    T_w2c_wrist.astype(np.float32))

        # Past-trajectory features (same shape/keys as YamScene).
        if self.past_n > 0:
            past_off = np.arange(self.past_n, 0, -1, dtype=np.int64) * s
            past_frames = np.clip(start - past_off, 0, ep["n"] - 1)
            past_xyz_world = ep["ee_world"][past_frames].astype(np.float32)
            past_uv = _project_world_to_uv(
                past_xyz_world, ep["K_in"], ep["T_w2c"]).astype(np.float32)
            cur_xyz_world = ep["ee_world"][start].astype(np.float32)
            out["past_z"] = torch.from_numpy(past_xyz_world[:, 2].copy())
            out["past_grip"] = torch.from_numpy(
                ep["grips"][past_frames].astype(np.float32))
            out["past_uv"] = torch.from_numpy(past_uv)
            out["past_xyz"] = torch.from_numpy(past_xyz_world)
            out["current_xyz"] = torch.from_numpy(cur_xyz_world)
        return out