"""YAM scene-only dataset — reads ``record.py`` output, precomputes
EE pose + gripper + quaternion per frame at init (FK is the slow bit;
do it ONCE, not per-sample), and clusters all quaternions into K-means
centroids that go into the checkpoint so inference can decode back.

Per-sample target convention matches the pre-refactor code:
- predicted frames at ``[start + 1*s, …, start + n_window*s]``
- ``target_yx``  (T, 2) int — voxel pixel in the pred grid
- ``target_z``   (T,)   int — height bin (linear over z_range)
- ``target_grip``(T,)   int — gripper bin (linear over grip_range)
- ``target_rot`` (T,)   int — K-means centroid index for the EE quaternion

Also returns ``start_pix`` (B, 2) — the *current* EEF pixel in the
model-input image, fed to the model as the EEF feature lookup site.
"""
from __future__ import annotations

import pickle
from pathlib import Path
from typing import Optional

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


IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

GRIP_RANGE = (-0.2, 0.8)        # canonical PARA gripper range (radians-ish)
ROT_KMEANS_K = 256              # K-means clusters for quat rotation head
                                # (raised from 64 → 256 on 2026-06-14 to give
                                # 5-8° resolution in populated regions of SO(3)
                                # across the multi-task base; cost is one wide
                                # linear in the rot_head + a slower one-shot
                                # K-means init — negligible at runtime.)


def _imagenet_norm_chw(bgr: np.ndarray, img_size: int) -> torch.Tensor:
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    rgb = cv2.resize(rgb, (img_size, img_size), interpolation=cv2.INTER_AREA)
    x = rgb.astype(np.float32) / 255.0
    x = (x - IMAGENET_MEAN) / IMAGENET_STD
    return torch.from_numpy(x.transpose(2, 0, 1).copy())


_KIN = None


def _fk_ee_in_rbase(joints7: np.ndarray) -> np.ndarray:
    global _KIN
    if _KIN is None:
        import sys
        rf = Path(__file__).resolve().parents[1] / "raiden_fork"
        for p in (rf, rf / "third_party" / "i2rt"):
            if str(p) not in sys.path:
                sys.path.insert(0, str(p))
        from i2rt.robots.kinematics import Kinematics
        from raiden._xml_paths import get_yam_4310_linear_xml_path
        _KIN = Kinematics(get_yam_4310_linear_xml_path(), site_name="grasp_site")
    q = np.zeros(8, dtype=np.float64)
    q[:6] = np.asarray(joints7, dtype=np.float64)[:6]
    return np.asarray(_KIN.fk(q), dtype=np.float64)


def _R_to_quat(R: np.ndarray) -> np.ndarray:
    """3×3 rotation matrix → unit quaternion (w, x, y, z)."""
    t = np.trace(R)
    if t > 0:
        s = np.sqrt(t + 1.0) * 2
        return np.array([0.25 * s,
                         (R[2, 1] - R[1, 2]) / s,
                         (R[0, 2] - R[2, 0]) / s,
                         (R[1, 0] - R[0, 1]) / s])
    if R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
        s = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2
        return np.array([(R[2, 1] - R[1, 2]) / s,
                         0.25 * s,
                         (R[0, 1] + R[1, 0]) / s,
                         (R[0, 2] + R[2, 0]) / s])
    if R[1, 1] > R[2, 2]:
        s = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2
        return np.array([(R[0, 2] - R[2, 0]) / s,
                         (R[0, 1] + R[1, 0]) / s,
                         0.25 * s,
                         (R[1, 2] + R[2, 1]) / s])
    s = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2
    return np.array([(R[1, 0] - R[0, 1]) / s,
                     (R[0, 2] + R[2, 0]) / s,
                     (R[1, 2] + R[2, 1]) / s,
                     0.25 * s])


def _quat_canonical(q: np.ndarray) -> np.ndarray:
    """Disambiguate q vs -q so K-means can converge — force w ≥ 0."""
    return q if q[0] >= 0 else -q


def _kmeans_quats(quats: np.ndarray, K: int, n_iter: int = 100,
                   seed: int = 0) -> np.ndarray:
    """Simple K-means on the unit sphere (cosine sim). Returns (K, 4)."""
    from sklearn.cluster import KMeans
    km = KMeans(n_clusters=K, random_state=seed, n_init=10, max_iter=n_iter)
    km.fit(quats)
    C = km.cluster_centers_
    return (C / (np.linalg.norm(C, axis=1, keepdims=True) + 1e-8)).astype(np.float32)


def _nearest_centroid_idx(quats: np.ndarray, centroids: np.ndarray) -> np.ndarray:
    """For each quat in ``quats`` (N, 4), index of closest centroid (cosine sim)."""
    sims = quats @ centroids.T                                # (N, K)
    return np.argmax(np.abs(sims), axis=1).astype(np.int64)   # |sim| handles antipodal


def _project_world_to_uv(xyz_world: np.ndarray, K: np.ndarray,
                          T_w2c: np.ndarray) -> np.ndarray:
    """Project world-frame ``xyz_world`` (..., 3) to (..., 2) pixel coords
    using the precomputed camera intrinsics ``K`` (already scaled to
    ``img_size``) and world→cam extrinsic. Returns (u, v) in pixel coords
    of whatever resolution ``K`` was scaled for. NaN/clip not applied —
    callers handle out-of-frame as needed."""
    orig_shape = xyz_world.shape
    pts = xyz_world.reshape(-1, 3).astype(np.float64)
    homo = np.concatenate([pts, np.ones((len(pts), 1), dtype=np.float64)], axis=-1)
    cam_pts = (T_w2c.astype(np.float64) @ homo.T).T[:, :3]    # (n, 3) cam frame
    z = cam_pts[:, 2:3] + 1e-8
    norm = cam_pts / z                                          # (n, 3) — last col ≈ 1
    pix = (K.astype(np.float64) @ norm.T).T                     # (n, 3)
    return pix[:, :2].reshape(*orig_shape[:-1], 2)


def _read_scene_frame(ep_dir, frame_idx: int):
    """Read scene-cam frame ``ep_dir/rgb/scene_camera/<idx:010d>.{jpg,png}``.
    Tries .jpg first (new record.py convention), falls back to .png (legacy
    `rd convert` output). Returns BGR ndarray; raises if neither exists."""
    base = ep_dir / "rgb" / "scene_camera" / f"{frame_idx:010d}"
    for ext in (".jpg", ".png"):
        bgr = cv2.imread(str(base) + ext)
        if bgr is not None:
            return bgr
    raise RuntimeError(f"scene_camera frame missing: {base}.{{jpg,png}}")


def _read_wrist_frame(ep_dir, frame_idx: int):
    """Read wrist-cam JPEG. Returns BGR ndarray or None if missing."""
    base = ep_dir / "rgb" / "right_wrist_camera" / f"{frame_idx:010d}"
    for ext in (".jpg", ".png"):
        bgr = cv2.imread(str(base) + ext)
        if bgr is not None:
            return bgr
    return None


def _load_wrist_episode_cfg(ep_dir, img_size: int):
    """Read wrist intrinsics + hand-eye T_ee_wrist_cam from this episode's
    calibration_results.json. Returns (K_wrist_in_imgsize, T_ee_wrist_cam)
    both as float32 numpy arrays, or (None, None) if absent."""
    cal = ep_dir / "calibration_results.json"
    if not cal.exists():
        return None, None
    try:
        import json as _json
        d = _json.loads(cal.read_text())
    except Exception:
        return None, None
    rwc = d.get("right_wrist_camera")
    if not rwc:
        return None, None
    K_src = np.asarray(rwc["intrinsics_save_res"], dtype=np.float64)
    W_src, H_src = rwc.get("save_wh", [1280, 720])
    K_in = K_src.copy()
    K_in[0, 0] *= img_size / W_src; K_in[0, 2] *= img_size / W_src
    K_in[1, 1] *= img_size / H_src; K_in[1, 2] *= img_size / H_src
    he = rwc.get("hand_eye_T_ee_wrist_cam")
    if he is None:
        return None, None
    T_ee_wrist_cam = np.asarray(he, dtype=np.float64)
    return K_in.astype(np.float32), T_ee_wrist_cam.astype(np.float64)


class YamScene(Dataset):
    """Scene-only training samples from a record.py output tree."""

    def __init__(
        self,
        root: Path,
        n_window: int = 16,
        frame_stride: int = 8,
        pred_size: int = 128,
        img_size: int = 405,
        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: tuple[float, float] = GRIP_RANGE,
        rot_centroids: Optional[np.ndarray] = None,
        z_pad: float = 0.01,
        past_n: int = 0,
    ):
        """``z_range=None`` → auto-derive from the dataset's EE z observations
        (min/max with ``z_pad`` m of padding on each side). Pass an explicit
        tuple to override; finetune ckpts must reuse the source ckpt's range
        so the volume head's z-bin semantics stay aligned (mirrors how
        ``rot_centroids`` is reused on finetune)."""
        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
        # z_range is filled in after the per-frame EE pass if it was None.
        self._z_range_arg = z_range
        self.grip_lo, self.grip_hi = grip_range
        # Past-trajectory window (0 = no past features returned in __getitem__).
        self.past_n = int(past_n)

        # ── Precompute per-frame EE pose / quat / grip per episode ──────
        self.episodes: list[dict] = []          # one entry per episode
        for ep in sorted(self.root.iterdir()):
            if not (ep.is_dir() and ep.name.isdigit()):
                continue
            pkls = sorted((ep / "lowdim").glob("*.pkl"))
            if not pkls:
                continue
            fd0 = pickle.load(open(pkls[0], "rb"))
            K_src = np.asarray(fd0["intrinsics"]["scene_camera"], dtype=np.float64)
            T_scene_w = np.asarray(fd0["extrinsics"]["scene_camera"], dtype=np.float64)
            # World-frame convention. Two formats supported:
            #   new (post-2026-06-04):  `_scene_in_rbase=True` → world == right_arm_base,
            #                            no T_lfr; the scene extrinsic is in rbase already.
            #   legacy:                  `_scene_in_lbase=True` + a `T_left_from_right`
            #                            key; world == left_arm_base.
            if fd0.get("_scene_in_rbase", False):
                T_lfr = np.eye(4, dtype=np.float64)
            elif "T_left_from_right" in fd0:
                T_lfr = np.asarray(fd0["T_left_from_right"], dtype=np.float64)
            else:
                raise RuntimeError(
                    f"PKL {pkls[0]} has neither `_scene_in_rbase=True` nor "
                    "`T_left_from_right` key; cannot determine world frame.")
            # Read source image dims from the matching JPEG.
            bgr0 = _read_scene_frame(ep, 0)
            H_src, W_src = bgr0.shape[:2]
            # Scale K to model-input resolution.
            K_in = K_src.copy()
            K_in[0, 0] *= img_size / W_src; K_in[0, 2] *= img_size / W_src
            K_in[1, 1] *= img_size / H_src; K_in[1, 2] *= img_size / H_src
            T_w2c = np.linalg.inv(T_scene_w)

            n = len(pkls)
            ee_world = np.zeros((n, 3), dtype=np.float64)
            quats = np.zeros((n, 4), dtype=np.float32)
            grips = np.zeros(n, dtype=np.float32)
            joints_all = np.zeros((n, 7), dtype=np.float32)
            for i in range(n):
                fd_i = pickle.load(open(pkls[i], "rb"))
                j7 = np.asarray(fd_i["joints"], dtype=np.float64)
                joints_all[i] = j7.astype(np.float32)
                grips[i] = float(j7[-1])                  # gripper = last joint
                T_ee_rbase = _fk_ee_in_rbase(j7)
                T_ee_w = T_lfr @ T_ee_rbase
                ee_world[i] = T_ee_w[:3, 3]
                quats[i] = _quat_canonical(_R_to_quat(T_ee_w[:3, :3])).astype(np.float32)

            # Optional: wrist cam intrinsics + hand-eye, snapshotted by
            # record.py into each episode's calibration_results.json.
            K_wrist_in, T_ee_wrist_cam = _load_wrist_episode_cfg(ep, img_size)
            self.episodes.append({
                "dir": ep,
                "n": n,
                "K_in": K_in.astype(np.float32),
                "T_w2c": T_w2c.astype(np.float32),
                "T_lfr": T_lfr.astype(np.float32),
                "ee_world": ee_world,
                "quats": quats,
                "grips": grips,
                "joints": joints_all,
                "K_wrist_in": K_wrist_in,
                "T_ee_wrist_cam": T_ee_wrist_cam,
            })

        if not self.episodes:
            raise RuntimeError(f"no episodes in {self.root}")

        # ── z_range: auto-derive from EE z observations OR use caller's tuple ──
        # Auto matches how rot_centroids is data-derived: each new training
        # session normalizes to its own data range. Finetunes must pass the
        # source ckpt's range so the volume head's z-bin meaning is preserved.
        if self._z_range_arg 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 = self._z_range_arg
            print(f"  z_range from caller: [{z_lo:.4f}, {z_hi:.4f}]m")
        self.z_lo, self.z_hi = z_lo, z_hi

        # ── K-means rotation centroids ──────────────────────────────────
        # Single source of truth = the checkpoint (dataset_cfg.rot_centroids).
        # If caller passes ``rot_centroids`` (e.g. resuming from a ckpt), use
        # those; otherwise refit deterministically (seed=0). NO external cache
        # file so we never end up with two-copies-that-disagree.
        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 "
                  f"(e.g. resumed ckpt)")
        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…")
            self.rot_centroids = _kmeans_quats(all_q, n_rot_clusters)

        # Bin every frame's quaternion → centroid index, for fast __getitem__.
        for ep in self.episodes:
            ep["rot_bin"] = _nearest_centroid_idx(ep["quats"], self.rot_centroids)
            # Gripper bin: linear in grip_range.
            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)
            # Height bin: linear in z_range.
            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)
            # Pixel target (predicted-grid coords) for every frame.
            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)
            # start_pix in model-input (img_size) coords for the eef-feat lookup.
            ep["start_pix_img"] = uv.astype(np.float32)        # (n, 2) xy at img_size

        # ── Build sample index: (ep_idx, start_frame) ──────────────────
        max_offset = n_window * frame_stride
        self.entries: list[tuple[int, int]] = []
        for ei, ep in enumerate(self.episodes):
            for i in range(ep["n"] - max_offset):
                self.entries.append((ei, i))
        if not self.entries:
            raise RuntimeError(
                f"no samples; need ep length > n_window*frame_stride = {max_offset}")
        print(f"  built {len(self.entries)} samples (n_window={n_window}, "
              f"frame_stride={frame_stride})")

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

    def config(self) -> dict:
        """Per-dataset config payload to save in the checkpoint for inference."""
        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,        # (K, 4) np.float32
        }

    def __getitem__(self, idx: int):
        ei, start = self.entries[idx]
        ep = self.episodes[ei]
        bgr = _read_scene_frame(ep["dir"], start)
        rgb = _imagenet_norm_chw(bgr, self.img_size)
        # Window frames offset 1*s .. n_window*s.
        s = self.frame_stride
        frames = np.arange(1, self.n_window + 1, dtype=np.int64) * s + start

        out = {
            "rgb": rgb,                                              # (3, S, S)
            "start_pix": torch.from_numpy(ep["start_pix_img"][start]),  # (2,) xy in img_size
            "target_yx": torch.from_numpy(ep["yx_bin"][frames]),     # (T, 2)
            "target_z":  torch.from_numpy(ep["z_bin"][frames]),      # (T,)
            "target_grip": torch.from_numpy(ep["grip_bin"][frames]), # (T,)
            "target_rot": torch.from_numpy(ep["rot_bin"][frames]),   # (T,)
            "target_xyz": torch.from_numpy(                          # (T, 3) world EE pos
                ep["ee_world"][frames].astype(np.float32)),
            "_bgr_src": bgr,                                          # for viz (not GPU)
            "_K_in": torch.from_numpy(ep["K_in"]),                    # (3, 3) for baseline viz
            "_T_w2c": torch.from_numpy(ep["T_w2c"]),                  # (4, 4) for baseline viz
        }

        # Wrist data — present if the episode JSON had `right_wrist_camera`.
        # Computes T_w2c for the wrist cam from this frame's joint state:
        #   T_world_wrist_cam = T_world_ee @ T_ee_wrist_cam
        #   T_w2c_wrist       = inv(T_world_wrist_cam)
        if ep["K_wrist_in"] is not None and ep["T_ee_wrist_cam"] is not None:
            w_bgr = _read_wrist_frame(ep["dir"], start)
            if w_bgr is not None:
                out["wrist_rgb"] = _imagenet_norm_chw(w_bgr, self.img_size)
                out["_wrist_bgr_src"] = w_bgr                         # for viz only
                j7 = ep["joints"][start].astype(np.float64)
                T_ee_rbase = _fk_ee_in_rbase(j7)
                T_ee_w = ep["T_lfr"].astype(np.float64) @ T_ee_rbase
                T_world_wrist_cam = T_ee_w @ ep["T_ee_wrist_cam"]
                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 (for ablation experiments that consume
        # past EE state). Frames are [start - past_n*s, ..., start - s],
        # clamped to [0, n-1] when the start is near the episode beginning.
        # All values are returned in their natural units (m, m, px) — the
        # model normalizes per-flag.
        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_n, 3)
            past_uv = _project_world_to_uv(
                past_xyz_world, ep["K_in"], ep["T_w2c"]).astype(np.float32)   # (past_n, 2)
            cur_xyz_world = ep["ee_world"][start].astype(np.float32)         # (3,)
            out["past_z"] = torch.from_numpy(past_xyz_world[:, 2].copy())    # (past_n,) raw meters
            out["past_grip"] = torch.from_numpy(
                ep["grips"][past_frames].astype(np.float32))                  # (past_n,) raw units
            out["past_uv"] = torch.from_numpy(past_uv)                        # (past_n, 2) img_size px
            out["past_xyz"] = torch.from_numpy(past_xyz_world)                # (past_n, 3) world m
            out["current_xyz"] = torch.from_numpy(cur_xyz_world)              # (3,) for delta enc
        return out