"""Dataset stubs for ACT training.

Two datasets are provided:

* `RandomDemoDataset`: on-the-fly random tensors with the right shapes, useful
  for smoke-testing the pipeline without having any real data on disk.

* `EpisodicHDF5Dataset`: loader for the per-episode HDF5 layout produced by
  `tools/zips_to_hdf5.py`. Each `episode_*.hdf5` contains:
      /qpos                  (T, state_dim)        float32
      /action                (T, action_dim)       float32
      /images/<key>          (T, H, W) or (T,H,W,3) uint8     for each k in image_keys
      /wrench/<f>            (T, 6)                float32   for f in 0..4
      /timestamps            (T,)                  float64
  Grayscale images are auto-replicated to 3 channels for the ResNet backbone.
"""

from __future__ import annotations

import glob
import os
from dataclasses import dataclass, field

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

try:
    import h5py
except ImportError:  # h5py is only required when actually loading from disk
    h5py = None


@dataclass
class DatasetSpec:
    """Mirrors the obs schema the policy expects, so dataset and ACTConfig stay in sync."""

    state_dim: int                 # per-frame qpos dim (e.g. 22)
    action_dim: int
    image_keys: list[str]
    image_shape: tuple[int, int, int]  # (C, H, W) — what the policy receives (after replication)
    tactile_dim: int = 0           # per-frame tactile dim (e.g. 30 = 5 wrenches × 6)
    # Wrench keys to concatenate into the tactile vector (e.g. ["wrench_1"] or
    # ["wrench_0","wrench_1","wrench_2","wrench_3","wrench_4"]). When tactile_dim > 0
    # the concatenation must produce exactly tactile_dim values.
    tactile_keys: list[str] = field(default_factory=list)
    chunk_size: int = 100
    # Observation history: when > 1, the dataset returns the concatenation of the
    # last N frames (oldest first), front-padded with the first frame for t < N-1.
    # The policy then sees state of dim (state_dim * state_history), and similar
    # for tactile.
    state_history: int = 1
    tactile_history: int = 1
    # If True, the dataset returns ``action - qpos`` (per-step delta) as the
    # supervised target instead of the absolute action. Stats are then computed
    # on the deltas (much smaller std). Requires state_dim == action_dim.
    delta_actions: bool = False
    # Optional dim-subset masks. When set, the dataset slices qpos[indices] /
    # action[indices] per frame; state_dim/action_dim above refer to the
    # post-slice per-frame dim (i.e. len(indices)). Indices are applied AFTER
    # history concat is built (each frame is sliced, then frames concatenated).
    state_indices: list[int] | None = None
    action_indices: list[int] | None = None
    # Per-finger image keys for the in-graph tactile-fusion branch (e.g.
    # ["raw_0", ..., "raw_4"]). When set, the dataset returns:
    #   "tactile_frames": (n_fingers, T_pair, H, W) uint8 grayscale
    #   "tactile_bg":     (n_fingers, H, W)         uint8 grayscale
    # and does NOT return a "tactile" tensor — the in-graph encoder produces it.
    # ``tactile_image_pair`` controls the number of consecutive frames stacked
    # per sample (the encoder expects T=2 for vision modalities).
    tactile_image_keys: list[str] = field(default_factory=list)
    tactile_image_pair: int = 2
    # HDF5 group holding per-finger image frames. Default ``"images"`` reads
    # ``/images/<key>`` at native resolution and compression — slow for the
    # in-graph branch. After running ``tools/precompute_tactile_resized.py``
    # set this to ``"tactile_imgs"`` to read the pre-resized uncompressed
    # version (~10x faster I/O).
    tactile_image_group: str = "images"
    # HDF5 group holding per-episode per-finger bg images (uint8 grayscale).
    # ``"tf_bg"`` is the native-resolution variant from precompute_tactile_bg.py;
    # ``"tactile_bg_224"`` is the pre-resized variant from precompute_tactile_resized.py.
    tactile_bg_group: str = "tf_bg"


# --------------------------- Synthetic dataset (smoke test) -------------------------------


class RandomDemoDataset(Dataset):
    """Returns random tensors with the right shapes — for shape/sanity testing only."""

    def __init__(self, spec: DatasetSpec, length: int = 64):
        self.spec = spec
        self.length = length

    def __len__(self) -> int:
        return self.length

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        s = self.spec
        sample: dict[str, torch.Tensor] = {}
        for key in s.image_keys:
            sample[key] = torch.rand(*s.image_shape)
        if s.state_dim > 0:
            sample["state"] = torch.randn(s.state_dim)
        if s.tactile_dim > 0:
            sample["tactile"] = torch.randn(s.tactile_dim)
        sample["action"] = torch.randn(s.chunk_size, s.action_dim)
        sample["action_is_pad"] = torch.zeros(s.chunk_size, dtype=torch.bool)
        return sample


# --------------------------- HDF5 episode dataset -----------------------------------------


def _to_3chw(img: np.ndarray, target_hw: tuple[int, int]) -> torch.Tensor:
    """Convert an HDF5 image frame to a uint8 (3, H, W) torch tensor.

    Accepts (H, W) grayscale or (H, W, 3) RGB. Resizes via a simple slice/pad if
    `target_hw` differs from the source — kept minimal here; preprocess with the
    desired final resolution at conversion time for best results.
    """
    if img.ndim == 2:
        img = np.stack([img, img, img], axis=0)         # (3, H, W)
    elif img.ndim == 3 and img.shape[-1] == 3:
        img = np.transpose(img, (2, 0, 1))              # (3, H, W)
    elif img.ndim == 3 and img.shape[0] == 3:
        pass                                            # already (3, H, W)
    else:
        raise ValueError(f"unexpected image shape {img.shape}")

    H_t, W_t = target_hw
    H_s, W_s = img.shape[1], img.shape[2]
    if (H_s, W_s) != (H_t, W_t):
        # Center-crop or zero-pad to target. For a real run, resize at conversion time.
        out = np.zeros((3, H_t, W_t), dtype=img.dtype)
        h = min(H_s, H_t)
        w = min(W_s, W_t)
        ys, ys_t = (H_s - h) // 2, (H_t - h) // 2
        xs, xs_t = (W_s - w) // 2, (W_t - w) // 2
        out[:, ys_t:ys_t + h, xs_t:xs_t + w] = img[:, ys:ys + h, xs:xs + w]
        img = out
    return torch.from_numpy(np.ascontiguousarray(img))


class EpisodicHDF5Dataset(Dataset):
    """Random-access dataset over a directory of per-episode HDF5 files.

    Each __getitem__ samples one (episode, t) pair: returns the obs at time t
    plus the ground-truth action chunk a[t : t + chunk_size], padded with zeros and
    a boolean `action_is_pad` mask if the episode ends inside the chunk.
    """

    IMAGE_PATH = "/images/{key}"
    # Tactile features live under "/<group>/<leaf>". Default group is "wrench" so
    # plain leaf names ("1") and "wrench_1" both resolve to "/wrench/1". Any
    # "<group>_<leaf>" key (e.g. "sitr_0") routes to "/<group>/<leaf>".
    DEFAULT_TACTILE_GROUP = "wrench"
    QPOS_PATH = "/qpos"
    ACTION_PATH = "/action"

    def __init__(
        self,
        root: str | None,
        spec: DatasetSpec,
        episode_glob: str = "episode_*.hdf5",
        files: list[str] | None = None,
    ):
        """Either pass ``root`` (+ optional ``episode_glob``) to discover files, or
        pass an explicit ``files`` list (useful for train/val splits).
        """
        if h5py is None:
            raise ImportError("h5py is not installed. `pip install h5py` to use EpisodicHDF5Dataset.")
        self.spec = spec

        if files is not None:
            if not files:
                raise FileNotFoundError("EpisodicHDF5Dataset: explicit file list is empty.")
            self.files = list(files)
        else:
            if root is None:
                raise ValueError("EpisodicHDF5Dataset: must pass either root or files.")
            discovered = sorted(glob.glob(os.path.join(root, episode_glob)))
            if not discovered:
                raise FileNotFoundError(f"No HDF5 episodes found under {root!r} matching {episode_glob!r}.")
            self.files = discovered

        # Cache episode lengths.
        self.episode_lengths: list[int] = []
        for path in self.files:
            with h5py.File(path, "r") as f:
                self.episode_lengths.append(f[self.ACTION_PATH].shape[0])
        self.cum_lengths = np.cumsum(self.episode_lengths)

        if spec.tactile_dim > 0 and not spec.tactile_keys and not spec.tactile_image_keys:
            raise ValueError(
                "DatasetSpec.tactile_dim > 0 but both tactile_keys and tactile_image_keys are empty."
            )

    def __len__(self) -> int:
        return int(self.cum_lengths[-1])

    def _locate(self, idx: int) -> tuple[int, int]:
        ep = int(np.searchsorted(self.cum_lengths, idx, side="right"))
        prev = 0 if ep == 0 else int(self.cum_lengths[ep - 1])
        return ep, idx - prev

    @classmethod
    def _tactile_path(cls, name: str) -> str:
        """Map a user-facing tactile key to its HDF5 path.

        ``"wrench_1"`` → ``"/wrench/1"`` (legacy)
        ``"sitr_0"``   → ``"/sitr/0"``
        ``"1"``        → ``"/wrench/1"`` (default group)
        """
        if "_" in name:
            group, leaf = name.split("_", 1)
            return f"/{group}/{leaf}"
        return f"/{cls.DEFAULT_TACTILE_GROUP}/{name}"

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        s = self.spec
        ep_idx, t = self._locate(idx)
        path = self.files[ep_idx]
        ep_len = self.episode_lengths[ep_idx]

        sample: dict[str, torch.Tensor] = {}
        target_hw = (s.image_shape[1], s.image_shape[2])
        with h5py.File(path, "r") as f:
            for key in s.image_keys:
                img = f[self.IMAGE_PATH.format(key=key)][t]    # (H, W) or (H, W, 3) uint8
                sample[key] = _to_3chw(img, target_hw)

            if s.state_dim > 0:
                hist = max(1, s.state_history)
                # Indices [t-H+1, ..., t], clamped to >= 0 (front-pad with frame 0).
                ts = [max(0, t - hist + 1 + k) for k in range(hist)]
                # h5py wants sorted unique indices for fancy indexing; build the
                # window by manual reads (fast for small H).
                seq = np.stack([f[self.QPOS_PATH][ti].astype(np.float32) for ti in ts])
                if s.state_indices is not None:
                    seq = seq[:, s.state_indices]
                sample["state"] = torch.from_numpy(seq.reshape(-1))   # (H*state_dim,)

            if s.tactile_image_keys:
                # In-graph branch consumes raw image pairs + per-episode bg per
                # finger. Stack a (n_fingers, T_pair, H, W) uint8 tensor and a
                # (n_fingers, H, W) uint8 bg tensor; the model does the resize +
                # normalization on the GPU (or skips resize if frames are
                # already at the encoder input resolution).
                n_pair = max(1, s.tactile_image_pair)
                ts = [max(0, t - n_pair + 1 + k) for k in range(n_pair)]
                frames_per_finger = []
                bg_per_finger = []
                for key in s.tactile_image_keys:
                    img_ds = f[f"/{s.tactile_image_group}/{key}"]
                    pair = np.stack([img_ds[ti] for ti in ts], axis=0)  # (T_pair, H, W) uint8
                    frames_per_finger.append(pair)
                    bg_ds = f[f"/{s.tactile_bg_group}/{key}"]
                    bg_per_finger.append(np.asarray(bg_ds, dtype=np.uint8))
                tactile_frames = np.stack(frames_per_finger, axis=0)    # (n_f, T_pair, H, W)
                tactile_bg = np.stack(bg_per_finger, axis=0)            # (n_f, H, W)
                sample["tactile_frames"] = torch.from_numpy(tactile_frames)
                sample["tactile_bg"] = torch.from_numpy(tactile_bg)
            elif s.tactile_keys and s.tactile_dim > 0:
                hist = max(1, s.tactile_history)
                ts = [max(0, t - hist + 1 + k) for k in range(hist)]
                # For each timestep, concat all tactile keys; then concat across timesteps.
                per_t = []
                for ti in ts:
                    parts = [f[self._tactile_path(k)][ti].astype(np.float32) for k in s.tactile_keys]
                    per_t.append(np.concatenate(parts, axis=-1))
                tac = np.concatenate(per_t, axis=-1)               # (H * sum_dim,)
                expected = s.tactile_dim * hist
                if tac.shape[0] != expected:
                    raise ValueError(
                        f"tactile_dim*history={expected} but concat of {s.tactile_keys}×{hist} gave {tac.shape[0]}"
                    )
                sample["tactile"] = torch.from_numpy(tac)

            end = min(t + s.chunk_size, ep_len)
            action = f[self.ACTION_PATH][t:end].astype(np.float32)
            if s.action_indices is not None:
                action = action[:, s.action_indices]
            if s.delta_actions:
                # delta[k] = action[t+k] - qpos[t+k] = qpos[t+k+1] - qpos[t+k]
                # (action at the last episode step is the repeated final pose, so the
                # resulting delta there is zero — a sensible "do nothing" target.)
                q_chunk = f[self.QPOS_PATH][t:end].astype(np.float32)
                # For delta_actions to be well-defined we must subtract qpos sliced
                # to the same indices as action.
                if s.action_indices is not None:
                    q_chunk = q_chunk[:, s.action_indices]
                if q_chunk.shape[1] != action.shape[1]:
                    raise ValueError(
                        f"delta_actions requires action_dim == state_dim per frame; "
                        f"got action {action.shape[1]} vs qpos {q_chunk.shape[1]}."
                    )
                action = action - q_chunk

        pad = s.chunk_size - action.shape[0]
        if pad > 0:
            action = np.concatenate([action, np.zeros((pad, s.action_dim), dtype=np.float32)], axis=0)
        is_pad = np.zeros(s.chunk_size, dtype=bool)
        if pad > 0:
            is_pad[-pad:] = True

        sample["action"] = torch.from_numpy(action)
        sample["action_is_pad"] = torch.from_numpy(is_pad)
        return sample