dataset/dataset.py

"""Map-style dataset over the fota_unlabeled tactile parquets.

Each parquet row holds a JPEG-encoded `image` blob plus per-frame metadata
(`obj_name`, pose, force, etc.). The dataset enumerates all parquet files
under `root`, indexes them by file metadata only (no full read at init),
and decodes JPEGs on demand.

For DataLoader use:
  - Each worker keeps an LRU of `cache_files` parquet image-columns in memory
    (PyArrow `ChunkedArray` of binary refs). Random reads within a cached
    file are zero-copy slices; misses pay one full column read (~2 GB).
  - Combine with `ParquetFileShuffleSampler` for shuffle without thrashing:
    shuffles file order each epoch and shuffles indices inside each file,
    so a worker only ever pulls from one cached file at a time.
"""
from __future__ import annotations

import io
import json
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterator, Mapping, Sequence

import numpy as np
import pyarrow.parquet as pq
import torch
from PIL import Image
from torch.utils.data import Dataset, Sampler
from torchvision.transforms import ColorJitter

DEFAULT_DATA_ROOT = Path("/group2/ct/weihanx/tactile_world_model/tactile_vae/data")
DEFAULT_FILE_GLOB = "train-*-of-*.parquet"
DEFAULT_SPLITS_PATH = Path(__file__).with_name("splits.json")
VALID_SPLITS = ("train", "val", "test")

# Metadata columns we expose by default — float32 scalars usable as conditioning.
META_COLUMNS: tuple[str, ...] = (
    "x_mm", "y_mm", "z_mm",
    "quat_x", "quat_y", "quat_z", "quat_w",
    "f_x", "f_y", "f_z",
    "grid_z_max", "grid_z_mean",
)


@dataclass
class ColorJitterConfig:
    """Training-time color-jitter knobs (forwarded to `torchvision.transforms.ColorJitter`).

    Each value is the magnitude of perturbation around the identity. `0.0`
    disables that channel; when all four are zero the config is a no-op.
    Defaults match a mild augmentation reasonable for tactile RGB frames.

    Treat as part of the training config — eval/inference paths should pass
    `color_jitter=None` to the dataset so augmentation is off.
    """
    brightness: float = 0.2
    contrast: float = 0.2
    saturation: float = 0.2
    hue: float = 0.1

    def is_noop(self) -> bool:
        return not any((self.brightness, self.contrast, self.saturation, self.hue))

    def build(self) -> ColorJitter | None:
        if self.is_noop():
            return None
        return ColorJitter(
            brightness=self.brightness,
            contrast=self.contrast,
            saturation=self.saturation,
            hue=self.hue,
        )


def _resolve_color_jitter(
    cfg: ColorJitterConfig | Mapping[str, float] | ColorJitter | None,
) -> ColorJitter | None:
    if cfg is None:
        return None
    if isinstance(cfg, ColorJitter):
        return cfg
    if isinstance(cfg, Mapping):
        cfg = ColorJitterConfig(**cfg)
    if isinstance(cfg, ColorJitterConfig):
        return cfg.build()
    raise TypeError(f"unsupported color_jitter type: {type(cfg).__name__}")


def default_image_transform(
    image_size: int = 128,
    color_jitter: ColorJitterConfig | Mapping[str, float] | ColorJitter | None = None,
) -> Callable[[Image.Image], torch.Tensor]:
    """Resize → (optional) color-jitter → CHW float32 in [0, 1].

    `color_jitter` accepts a `ColorJitterConfig`, a dict of its fields, an
    already-built `torchvision.transforms.ColorJitter`, or `None` to disable.
    Apply jitter on PIL (before tensor conversion) — torchvision's RNG is
    per-worker in PyTorch DataLoaders, so per-sample augmentation is correct
    with `num_workers > 0`.
    """
    jitter = _resolve_color_jitter(color_jitter)

    def _tx(img: Image.Image) -> torch.Tensor:
        if img.size != (image_size, image_size):
            img = img.resize((image_size, image_size), Image.BILINEAR)
        if jitter is not None:
            img = jitter(img)
        arr = np.array(img, dtype=np.uint8, copy=True)  # H, W, 3, writable
        t = torch.from_numpy(arr).permute(2, 0, 1).contiguous().float().div_(255.0) # [0, 1]
        return t
    return _tx


class TactileParquetDataset(Dataset):
    """Lazy parquet dataset that yields decoded RGB tensors (+ optional metadata).

    Args:
        root: directory containing parquet shards.
        file_glob: glob pattern relative to `root`. Defaults to
            `train-*-of-*.parquet`; `.raw.parquet` shards are excluded.
        image_size: output square size if `transform` is the default.
        transform: callable(PIL.Image) -> Tensor. If None, uses
            `default_image_transform(image_size, color_jitter=color_jitter)`.
            Pass an explicit transform to fully override preprocessing;
            `color_jitter` is then ignored.
        color_jitter: training-time RGB augmentation, applied inside the
            default transform. Accepts a `ColorJitterConfig`, a dict of its
            fields, a prebuilt `torchvision.transforms.ColorJitter`, or
            `None` (default — augmentation off). For eval/inference, leave
            as `None`.
        split: one of `"train"`, `"val"`, `"test"`, or `None`. If set, the
            dataset only exposes rows assigned to that split by `splits_path`.
        splits_path: JSON manifest produced by `make_splits.py`. Defaults to
            `tactile_vae/dataset/splits.json`. Only consulted when
            `split is not None`.
        return_meta: if True, `__getitem__` returns `(image, meta_dict)` where
            `meta_dict` has float32 tensors for the columns in `meta_columns`.
        meta_columns: which metadata columns to surface.
        cache_files: how many parquet image-columns to keep in memory per worker.
    """

    def __init__(
        self,
        root: str | Path = DEFAULT_DATA_ROOT,
        file_glob: str = DEFAULT_FILE_GLOB,
        image_size: int = 128,
        transform: Callable[[Image.Image], torch.Tensor] | None = None,
        color_jitter: ColorJitterConfig | Mapping[str, float] | ColorJitter | None = None,
        split: str | None = None,
        splits_path: str | Path | None = None,
        return_meta: bool = False,
        meta_columns: Sequence[str] = META_COLUMNS,
        cache_files: int = 1,
    ) -> None:
        root = Path(root)
        files = sorted(root.glob(file_glob))
        files = [p for p in files if ".raw." not in p.name]
        if not files:
            raise FileNotFoundError(f"no parquet shards under {root} matching {file_glob!r}")

        self.files: list[Path] = files
        self.image_size = image_size
        self.transform = transform or default_image_transform(
            image_size, color_jitter=color_jitter
        )
        self.split = split
        self.return_meta = return_meta
        self.meta_columns: list[str] = list(meta_columns)
        self.cache_files = max(1, int(cache_files))

        # True per-file row counts from parquet metadata (independent of split).
        file_row_counts: list[int] = [pq.ParquetFile(p).metadata.num_rows for p in files]

        # Build `_samples`: (N, 2) int64 array of (file_idx, row_idx), sorted
        # by file_idx then row_idx. Grouping by file keeps contiguous global
        # indices on the same parquet — that's what the file-cache relies on.
        if split is None:
            per_file_rows: list[np.ndarray] = [
                np.arange(n, dtype=np.int64) for n in file_row_counts
            ]
        else:
            if split not in VALID_SPLITS:
                raise ValueError(f"split must be one of {VALID_SPLITS}; got {split!r}")
            spath = Path(splits_path) if splits_path is not None else DEFAULT_SPLITS_PATH
            if not spath.exists():
                raise FileNotFoundError(
                    f"split={split!r} requested but {spath} not found. "
                    "Generate it with `python tactile_vae/dataset/make_splits.py`."
                )
            with spath.open() as f:
                manifest = json.load(f)
            split_map: dict[str, list[int]] = manifest["splits"][split]
            per_file_rows = []
            for p, total_rows in zip(files, file_row_counts):
                rows = split_map.get(p.name, [])
                arr = np.asarray(rows, dtype=np.int64)
                if arr.size and (arr.min() < 0 or arr.max() >= total_rows):
                    raise ValueError(
                        f"split manifest for {p.name} has row index out of range "
                        f"[0, {total_rows}): min={int(arr.min())} max={int(arr.max())}"
                    )
                # Sort within file so cache reads progress monotonically.
                per_file_rows.append(np.sort(arr))

        chunks = [
            np.stack([np.full(rows.shape[0], fi, dtype=np.int64), rows], axis=1)
            for fi, rows in enumerate(per_file_rows)
            if rows.shape[0] > 0
        ]
        self._samples: np.ndarray = (
            np.concatenate(chunks, axis=0)
            if chunks
            else np.zeros((0, 2), dtype=np.int64)
        )
        per_file_counts = [rows.shape[0] for rows in per_file_rows]
        self._per_file_offsets: np.ndarray = np.concatenate(
            [[0], np.cumsum(per_file_counts, dtype=np.int64)]
        )

        # Per-worker caches; populated lazily.
        self._image_cache: "OrderedDict[int, object]" = OrderedDict()
        self._meta_cache: "OrderedDict[int, dict[str, np.ndarray]]" = OrderedDict()

    def __len__(self) -> int:
        return int(self._samples.shape[0])

    @property
    def num_files(self) -> int:
        return len(self.files)

    def file_lengths(self) -> list[int]:
        """Number of samples drawn from each file under the current split."""
        return np.diff(self._per_file_offsets).astype(int).tolist()

    def sample_id(self, idx: int) -> str:
        """Stable ID for sample `idx`: `<file-stem>:<row>` — matches `make_splits.py`."""
        fi, ri = self._locate(idx)
        return f"{self.files[fi].stem}:{ri}"

    def __getstate__(self) -> dict:
        # Strip caches before pickling to workers — pyarrow objects don't
        # pickle cleanly and the cache should be per-worker anyway.
        state = self.__dict__.copy()
        state["_image_cache"] = OrderedDict()
        state["_meta_cache"] = OrderedDict()
        return state

    def _locate(self, idx: int) -> tuple[int, int]:
        n = self.__len__()
        if idx < 0:
            idx += n
        if idx < 0 or idx >= n:
            raise IndexError(idx)
        fi, ri = self._samples[idx]
        return int(fi), int(ri)

    def _get_image_column(self, fi: int):
        col = self._image_cache.get(fi)
        if col is not None:
            self._image_cache.move_to_end(fi)
            return col
        col = pq.read_table(self.files[fi], columns=["image"]).column("image")
        self._image_cache[fi] = col
        while len(self._image_cache) > self.cache_files:
            self._image_cache.popitem(last=False)
        return col

    def _get_meta_columns(self, fi: int) -> dict[str, np.ndarray]:
        if not self.meta_columns:
            return {}
        cached = self._meta_cache.get(fi)
        if cached is not None:
            self._meta_cache.move_to_end(fi)
            return cached
        tbl = pq.read_table(self.files[fi], columns=self.meta_columns)
        cached = {name: tbl.column(name).to_numpy(zero_copy_only=False)
                  for name in self.meta_columns}
        self._meta_cache[fi] = cached
        while len(self._meta_cache) > self.cache_files:
            self._meta_cache.popitem(last=False)
        return cached

    def __getitem__(self, idx: int):
        fi, ri = self._locate(idx)
        col = self._get_image_column(fi)
        raw = col[ri].as_py()
        with Image.open(io.BytesIO(raw)) as im:
            im = im.convert("RGB")
            tensor = self.transform(im)

        if not self.return_meta:
            return tensor

        meta_np = self._get_meta_columns(fi)
        meta = {
            name: torch.as_tensor(float(arr[ri]), dtype=torch.float32)
            for name, arr in meta_np.items()
        }
        return tensor, meta


class ParquetFileShuffleSampler(Sampler[int]):
    """Shuffle that keeps each worker reading from one parquet at a time.

    Shuffles file order per epoch, then yields indices inside each file in a
    shuffled order. With `cache_files=1` on the dataset this means each
    worker only ever pays one file-load cost per file.
    """

    def __init__(
        self,
        dataset: TactileParquetDataset,
        seed: int = 0,
        shuffle_files: bool = True,
        shuffle_within_file: bool = True,
    ) -> None:
        self.dataset = dataset
        self.seed = seed
        self.shuffle_files = shuffle_files
        self.shuffle_within_file = shuffle_within_file
        self.epoch = 0

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

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

    def __iter__(self) -> Iterator[int]:
        rng = np.random.default_rng(self.seed + self.epoch)
        n_files = self.dataset.num_files
        file_order = rng.permutation(n_files) if self.shuffle_files else np.arange(n_files)
        offsets = self.dataset._per_file_offsets
        for fi in file_order:
            start = int(offsets[fi])
            end = int(offsets[fi + 1])
            n = end - start
            if n == 0:
                continue  # this file contributes no samples to the current split
            local = rng.permutation(n) if self.shuffle_within_file else np.arange(n)
            for li in local:
                yield int(start + li)