src/fpgm/data/pointworld.py

"""PointWorld-DROID dataset access: selective HF download + flows.h5 reading.

PointWorld ships as ``tar.zst`` split into numbered ``.part-*`` files and must
be restored with the dataset's own ``recover_dataset_from_parts.sh
--include-prefix ...`` script -- there is no way to untar a lone
``.part-0000`` directly, so this module always fetches that recovery script
alongside whatever data prefix it needs and shells out to it rather than
reimplementing zstd+tar restoration.
"""

from __future__ import annotations

import io
import json
import os
import re
import stat
import subprocess
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import h5py
import numpy as np
from huggingface_hub import snapshot_download
from PIL import Image

from fpgm.types import DataError, SceneFlowClip
from fpgm.utils.io import ensure_dir
from fpgm.utils.logging import get_logger

#: Clip groups are keyed ``"{start}:{end}"``; every other root key is episode-level metadata.
_CLIP_KEY_RE = re.compile(r"\d+:\d+")

logger = get_logger(__name__)

HF_REPO_ID = "nvidia/PointWorld-DROID"
RECOVER_SCRIPT_NAME = "recover_dataset_from_parts.sh"

#: 1.23 TB and not selectively restorable -- refused outright, not just "not
#: implemented", so a future caller can't accidentally wire this up either.
FORBIDDEN_PREFIXES = ("droid/depth_320x180",)


def _check_prefix_allowed(prefix: str) -> None:
    for forbidden in FORBIDDEN_PREFIXES:
        if prefix.startswith(forbidden):
            raise ValueError(
                f"refusing to download {prefix!r}: {forbidden} is 1.23 TB and not "
                "cherry-pickable. This is a hard guardrail, not a suggestion."
            )


@dataclass
class ShardInfo:
    """One entry of ``_shards_manifest.json``'s ``"shards"`` list."""

    shard_id: str
    member_count: int
    members: list[str]
    remote_prefix: str
    raw_member_bytes: int

    @property
    def episode_uuids(self) -> list[str]:
        """uuid for every ``*_flows.h5`` member, derived from its filename."""
        uuids = []
        for member in self.members:
            name = Path(member).name
            if name.endswith("_flows.h5"):
                uuids.append(name[: -len("_flows.h5")])
        return uuids


class ShardManifest:
    """Parsed ``_shards_manifest.json``: which shard holds which episode."""

    def __init__(self, shards: list[ShardInfo]):
        self._shards = {s.shard_id: s for s in shards}
        self._episode_to_shard = {uuid: s.shard_id for s in shards for uuid in s.episode_uuids}

    @classmethod
    def from_json(cls, path: Path) -> ShardManifest:
        raw = json.loads(Path(path).read_text())
        shards = [
            ShardInfo(
                shard_id=s["shard_id"],
                member_count=s["member_count"],
                members=s["members"],
                remote_prefix=s["remote_prefix"],
                raw_member_bytes=s["raw_member_bytes"],
            )
            for s in raw["shards"]
        ]
        return cls(shards)

    def shard(self, shard_id: str) -> ShardInfo:
        try:
            return self._shards[shard_id]
        except KeyError:
            raise DataError(f"unknown shard id {shard_id!r}") from None

    def shard_for_episode(self, uuid: str) -> str:
        try:
            return self._episode_to_shard[uuid]
        except KeyError:
            raise DataError(f"episode {uuid!r} not found in any shard of the manifest") from None

    def episodes_in_shard(self, shard_id: str) -> list[str]:
        return self.shard(shard_id).episode_uuids


class PointWorldStore:
    """Selective downloader + accessor for the PointWorld-DROID HF dataset."""

    def __init__(
        self,
        local_dir: Path,
        hf_repo: str = HF_REPO_ID,
        hf_token: str | None = None,
    ):
        self.local_dir = ensure_dir(Path(local_dir))
        self.hf_repo = hf_repo
        self.hf_token = hf_token or os.environ.get("HF_TOKEN")
        self._manifest: ShardManifest | None = None

    # -- download ----------------------------------------------------------

    def _snapshot(self, allow_patterns: list[str]) -> Path:
        return Path(
            snapshot_download(
                repo_id=self.hf_repo,
                repo_type="dataset",
                local_dir=self.local_dir,
                allow_patterns=allow_patterns,
                token=self.hf_token,
            )
        )

    def _recover_script(self) -> Path:
        self._snapshot([RECOVER_SCRIPT_NAME])
        script = self.local_dir / RECOVER_SCRIPT_NAME
        script.chmod(script.stat().st_mode | stat.S_IEXEC)
        return script

    def _recover(self, include_prefix: str) -> None:
        """Unpack the downloaded ``package.tar.zst.part-*`` for ``include_prefix``.

        ``--packages`` and ``--out`` are both mandatory in
        ``recover_dataset_from_parts.sh`` (it prints usage and exits non-zero
        without them), and passing only ``--include-prefix`` made every
        ``download_shard``/``download_cameras`` call fail at the extraction step.
        Both point at ``local_dir``: the packaged parts live under it, and the
        script's own layout restores members back into the same tree.

        Every path is resolved to an absolute one first. ``local_dir`` arrives
        relative in the normal case (``download_pointworld.py`` defaults it to
        ``data/pointworld``), and ``cwd=self.local_dir`` then re-anchors that same
        relative string against the new working directory -- so the interpreter
        looked for ``data/pointworld/data/pointworld/recover_dataset_from_parts.sh``
        and raised ``FileNotFoundError`` after the multi-GB download had already
        succeeded.
        """
        root = self.local_dir.resolve()
        script = self._recover_script().resolve()
        logger.info("recovering %s via %s", include_prefix, script.name)
        subprocess.run(
            [
                str(script),
                "--packages", str(root),
                "--out", str(root),
                "--include-prefix", include_prefix,
            ],
            cwd=root,
            check=True,
        )

    def download_cameras(self) -> Path:
        """Fetch and restore all ~42,935 ``*_cameras.json`` (~19.5 MB packed)."""
        prefix = "droid/cameras"
        _check_prefix_allowed(prefix)
        self._snapshot([f"{prefix}/package.tar.zst.part-*"])
        self._recover(prefix)
        return self.local_dir / prefix

    def download_manifest(self) -> Path:
        """Fetch ``_shards_manifest.json`` (~3.9 MB, not packed -- no restore needed)."""
        rel = "droid/flows-fs-optimized/_shards_manifest.json"
        self._snapshot([rel])
        return self.local_dir / rel

    def download_shard(self, shard_id: str) -> Path:
        """Fetch and restore one flows shard (~3.5 GB packed, ~62 episodes)."""
        prefix = f"droid/flows-fs-optimized/{shard_id}"
        _check_prefix_allowed(prefix)
        self._snapshot([f"{prefix}/package.tar.zst.part-*"])
        self._recover(prefix)
        return self.local_dir / prefix

    # -- manifest ------------------------------------------------------------

    @property
    def manifest(self) -> ShardManifest:
        """Loads (downloading first if necessary) and caches the shard manifest."""
        if self._manifest is None:
            manifest_path = self.local_dir / "droid/flows-fs-optimized/_shards_manifest.json"
            if not manifest_path.exists():
                manifest_path = self.download_manifest()
            self._manifest = ShardManifest.from_json(manifest_path)
        return self._manifest

    def shard_for_episode(self, uuid: str) -> str:
        return self.manifest.shard_for_episode(uuid)

    def episodes_in_shard(self, shard_id: str) -> list[str]:
        return self.manifest.episodes_in_shard(shard_id)

    # -- flows / cameras access -----------------------------------------------

    def flows_path(self, uuid: str, shard_id: str | None = None) -> Path:
        """Locate ``<uuid>_flows.h5`` in the restored tree.

        ``recover_dataset_from_parts.sh`` extracts every shard's members into a
        single flat ``droid/flows-fs-optimized/`` directory -- the ``shard-NNNNNN``
        level exists only in the *packaged* download, not in the restored layout.
        The flat location is therefore tried first, with the sharded path kept as a
        fallback for callers who lay the tree out per shard themselves.
        """
        root = self.local_dir / "droid/flows-fs-optimized"
        flat = root / f"{uuid}_flows.h5"
        if flat.exists():
            return flat
        shard_id = shard_id or self.shard_for_episode(uuid)
        return root / shard_id / f"{uuid}_flows.h5"

    def cameras_path(self, uuid: str) -> Path:
        return self.local_dir / "droid/cameras" / f"{uuid}_cameras.json"

    def load_camera_calibration(self, uuid: str) -> dict[str, Any]:
        """Parsed ``<uuid>_cameras.json`` (per-serial ``optimized_extrinsics`` etc)."""
        path = self.cameras_path(uuid)
        if not path.exists():
            raise DataError(f"cameras json not found for {uuid!r}: {path}")
        return json.loads(path.read_text())

    def open_flows(self, uuid: str, shard_id: str | None = None) -> FlowsReader:
        """A context-managed :class:`FlowsReader` over ``<uuid>_flows.h5``."""
        return FlowsReader(self.flows_path(uuid, shard_id), episode_uuid=uuid)


class FlowsReader:
    """Reads one episode's ``*_flows.h5`` and yields :class:`SceneFlowClip` objects.

    Must be used as a context manager (``with store.open_flows(uuid) as
    reader:``) so the underlying HDF5 file handle is always closed.
    """

    def __init__(self, path: Path, episode_uuid: str | None = None):
        if not Path(path).exists():
            raise DataError(f"flows.h5 not found: {path}")
        self.path = Path(path)
        self.episode_uuid = episode_uuid
        self._file: h5py.File | None = None

    def __enter__(self) -> FlowsReader:
        self._file = h5py.File(self.path, "r")
        return self

    def __exit__(self, *exc_info: object) -> None:
        if self._file is not None:
            self._file.close()
            self._file = None

    @property
    def file(self) -> h5py.File:
        if self._file is None:
            raise DataError("FlowsReader must be used as a context manager (`with ... as reader:`)")
        return self._file

    def clip_keys(self) -> list[str]:
        """``"{start}:{end}"`` clip group keys, ordered by start frame.

        The file's root also holds episode-level summaries that are *not* clips
        (``ee_pos_motion_magnitudes``, ``ee_rot_motion_magnitudes``,
        ``has_gripper_movement``), so keys are filtered on the ``start:end`` shape
        rather than taken wholesale.
        """
        keys = [k for k in self.file.keys() if _CLIP_KEY_RE.fullmatch(k)]
        return sorted(keys, key=lambda k: int(k.split(":")[0]))

    def camera_serials(self, clip_key: str) -> list[str]:
        """Camera serials present for one clip (group names ``camera_<serial>_ext``)."""
        group = self.file[clip_key]
        return [
            name[len("camera_") : -len("_ext")]
            for name in group.keys()
            if name.startswith("camera_") and name.endswith("_ext")
        ]

    def read_clip(self, clip_key: str, camera_serial: str) -> SceneFlowClip:
        """Decode one clip/camera group into a :class:`~fpgm.types.SceneFlowClip`."""
        clip_group = self.file[clip_key]
        cam_group = clip_group[f"camera_{camera_serial}_ext"]
        start_str, end_str = clip_key.split(":")

        def arr(name: str) -> np.ndarray:
            return np.asarray(cam_group[name])

        initial_rgb = None
        if "initial_rgb" in cam_group:
            jpeg_bytes = bytes(cam_group["initial_rgb"][0])
            initial_rgb = np.array(Image.open(io.BytesIO(jpeg_bytes)).convert("RGB"))

        # Dense depth exists only for the clip's first frame -- one (H, W) array,
        # not a (T, H, W) stream -- so it is read straight through with no T axis.
        initial_depth = None
        if "initial_depth" in cam_group:
            initial_depth = arr("initial_depth").astype(np.uint16)

        scene_normals = None
        if "scene_normals" in cam_group:
            scene_normals = arr("scene_normals").astype(np.int8)

        # Proprioception is shared by every camera of a clip, so it is stored one
        # level up, on the clip group -- not inside the per-camera group.
        gripper_pose = (
            np.asarray(clip_group["gripper_pose"]).astype(np.float32)
            if "gripper_pose" in clip_group
            else None
        )
        gripper_open = (
            np.asarray(clip_group["gripper_open"]).astype(bool)
            if "gripper_open" in clip_group
            else None
        )
        # Also clip-level, shared by every camera -- kept mainly so the clip's
        # video/trajectory alignment can be *measured* (verify_annotation_stride)
        # rather than assumed. See fpgm.types.ClipTiming's docstring.
        joint_positions = (
            np.asarray(clip_group["joint_positions"]).astype(np.float32)
            if "joint_positions" in clip_group
            else None
        )

        return SceneFlowClip(
            key=clip_key,
            start=int(start_str),
            end=int(end_str),
            camera_serial=camera_serial,
            # scene_flows is stored float16 on disk; upcast for downstream numeric stability.
            scene_flows=arr("scene_flows").astype(np.float32),
            scene_visibility=arr("scene_visibility").astype(bool),
            scene_depth_valid=arr("scene_depth_valid_mask").astype(bool),
            scene_colors=arr("scene_colors").astype(np.uint8),
            intrinsic=arr("intrinsic").astype(np.float64),
            extrinsic=arr("extrinsic").astype(np.float64),
            initial_rgb=initial_rgb,
            gripper_pose=gripper_pose,
            gripper_open=gripper_open,
            joint_positions=joint_positions,
            initial_depth=initial_depth,
            scene_normals=scene_normals,
        )

    def read_all(self, camera_serial: str | None = None) -> Iterator[SceneFlowClip]:
        """Yield every clip, optionally restricted to one camera serial."""
        for clip_key in self.clip_keys():
            serials = [camera_serial] if camera_serial else self.camera_serials(clip_key)
            for serial in serials:
                yield self.read_clip(clip_key, serial)


#: Strides tried by :func:`verify_annotation_stride`. 2 is what has been measured
#: on every clip seen so far; 1 and the small neighbours are kept as candidates so
#: the function still tells the truth if a future shard turns out different.
_DEFAULT_CANDIDATE_STRIDES: tuple[int, ...] = (1, 2, 3, 4)


def verify_annotation_stride(
    clip_joint_positions: np.ndarray,
    trajectory_joint_positions: np.ndarray,
    clip_start: int,
    candidate_strides: tuple[int, ...] = _DEFAULT_CANDIDATE_STRIDES,
    atol: float = 1e-4,
) -> int:
    """Measure the clip-frame -> trajectory-row stride from real joint data.

    A PointWorld clip's own frame axis is not guaranteed to be one-to-one with
    ``trajectory.h5``'s row axis -- on every clip checked so far, clip frame ``t``
    is trajectory row ``2 * (clip_start + t)``, not ``clip_start + t`` (see
    :class:`fpgm.types.ClipTiming`'s docstring). Hardcoding that ``2`` here would
    repeat exactly the mistake that produced the original bug: a plausible-looking
    constant nobody checks against the data again. So instead, this tries each
    ``candidate_strides`` value, reads the corresponding rows out of
    ``trajectory_joint_positions``, and returns whichever stride's rows match
    ``clip_joint_positions`` almost exactly -- these are two independent readings
    of the same physical joint encoders at the same instants, so a correct stride
    matches to float32 precision and a wrong one does not match at all.

    Args:
        clip_joint_positions: ``(T, 7)`` joint positions from the clip's own h5
            group (``SceneFlowClip.joint_positions``).
        trajectory_joint_positions: ``(L, 7)`` (or more columns; only the leading 7
            are compared) joint positions from the episode's ``trajectory.h5``
            (``observation/robot_state/joint_positions``).
        clip_start: The clip's declared start index (the ``start`` in its
            ``"{start}:{end}"`` key).
        candidate_strides: Strides to test, in the order tried; the first exact
            match wins ties.
        atol: Maximum-abs-error tolerance for a candidate to count as a match.

    Returns:
        The stride (from ``candidate_strides``) whose implied trajectory rows best
        match ``clip_joint_positions``.

    Raises:
        DataError: If no candidate stride keeps every implied row inside
            ``trajectory_joint_positions``, or if the best-matching candidate's
            error still exceeds ``atol`` -- i.e. nothing measured actually fits,
            which should stop the pipeline rather than silently proceed on a guess.
    """
    clip_jp = np.asarray(clip_joint_positions, dtype=np.float64)
    traj_jp = np.asarray(trajectory_joint_positions, dtype=np.float64)
    n = clip_jp.shape[0]
    if n == 0:
        raise DataError("clip_joint_positions is empty; cannot measure an annotation stride")

    n_cols = min(clip_jp.shape[1], traj_jp.shape[1])
    errors: dict[int, float] = {}
    for stride in candidate_strides:
        rows = clip_start * stride + stride * np.arange(n)
        if rows[-1] >= traj_jp.shape[0]:
            continue  # this stride would read past the end of trajectory.h5
        errors[stride] = float(
            np.max(np.abs(clip_jp[:, :n_cols] - traj_jp[rows][:, :n_cols]))
        )

    if not errors:
        raise DataError(
            f"no candidate stride in {candidate_strides} keeps clip_start={clip_start}, "
            f"n={n} frames inside a {traj_jp.shape[0]}-row trajectory.h5"
        )

    best_stride = min(errors, key=errors.get)
    if errors[best_stride] > atol:
        raise DataError(
            f"no candidate stride in {candidate_strides} matches trajectory_joint_positions "
            f"within atol={atol} for clip_start={clip_start}; best was stride={best_stride} "
            f"with max abs error {errors[best_stride]:.6g} (all errors: {errors})"
        )
    return best_stride