src/fpgm/data/episode.py

"""Joins :class:`~fpgm.data.droid_raw.DroidRawClient` (raw mp4s/metadata) with
:class:`~fpgm.data.pointworld.PointWorldStore` (scene-flow annotations) into
one :class:`fpgm.types.Episode`.

Camera extrinsic precedence, most to least authoritative:

1. PointWorld's ``optimized_extrinsics`` from ``*_cameras.json``, if present.
2. The clip's own ``extrinsic`` in ``*_flows.h5`` (static per clip; this is
   what the flow annotations were actually rendered against, so absent (1) it
   should already agree with it).
3. DROID's raw metadata 6-DOF pose, converted via
   :func:`fpgm.data.droid_raw.cam2world_vector_to_world2cam` -- the least
   refined of the three, used only as a last resort (e.g. a camera role with
   no scene-flow annotations at all).

``ClipTiming.trajectory_fps`` is DROID's nominal 15 Hz control rate. The
public 1.0.1 raw release does not embed this in ``trajectory.h5``, so the
15.0 Hz fallback is the *expected* path, not a last resort -- it is logged at
INFO so a run's provenance is visible without being alarming.
``ClipTiming.mp4_fps`` always comes from the container itself, per camera,
because it is not guaranteed to match the trajectory rate.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import h5py
import numpy as np

from fpgm.data.droid_raw import (
    DroidRawClient,
    RawEpisodeAssets,
    cam2world_vector_to_world2cam,
    read_mp4_fps,
)
from fpgm.data.pointworld import PointWorldStore
from fpgm.types import CameraAssets, ClipTiming, DataError, Episode, SceneFlowClip
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: DROID's nominal robot control rate (Hz); see the module docstring.
DEFAULT_TRAJECTORY_FPS = 15.0

#: Internal-only key stashed on `Episode.metadata` so `EpisodeBuilder.clip_timing`
#: can recover the trajectory_fps decided at build time without re-deriving it.
_TRAJECTORY_FPS_METADATA_KEY = "_trajectory_fps"

#: Internal-only key stashed on `Episode.metadata` carrying the local
#: `trajectory.h5` path, so callers (e.g. `VelocityPipeline`) can measure
#: `ClipTiming.annotation_stride` against `observation/robot_state/joint_positions`
#: without threading `RawEpisodeAssets` through the whole pipeline.
_TRAJECTORY_H5_PATH_METADATA_KEY = "_trajectory_h5_path"


def _read_trajectory_fps(trajectory_h5_path: Path | None) -> float:
    """Best-effort read of the control rate from ``trajectory.h5``; else the default."""
    if trajectory_h5_path is None or not trajectory_h5_path.exists():
        logger.info(
            "no trajectory.h5 available; using DROID's nominal %.1f Hz control rate",
            DEFAULT_TRAJECTORY_FPS,
        )
        return DEFAULT_TRAJECTORY_FPS
    with h5py.File(trajectory_h5_path, "r") as f:
        for key in ("fps", "control_hz", "trajectory_fps"):
            if key in f.attrs:
                return float(f.attrs[key])
    logger.info(
        "trajectory.h5 has no fps attribute; using DROID's nominal %.1f Hz control rate",
        DEFAULT_TRAJECTORY_FPS,
    )
    return DEFAULT_TRAJECTORY_FPS


class EpisodeBuilder:
    """Builds a fully-joined :class:`~fpgm.types.Episode` from raw + annotated sources."""

    def __init__(self, droid_client: DroidRawClient, pointworld_store: PointWorldStore):
        self.droid_client = droid_client
        self.pointworld_store = pointworld_store

    def build(
        self,
        uuid: str,
        *,
        camera_roles: tuple[str, ...] = ("ext1", "ext2"),
        shard_id: str | None = None,
    ) -> Episode:
        """Fetch (if not already cached) and join every source for ``uuid``."""
        raw = self.droid_client.download_episode(uuid, roles=camera_roles)

        with self.pointworld_store.open_flows(uuid, shard_id=shard_id) as reader:
            clips = list(reader.read_all())
        if not clips:
            logger.warning("episode %s has no scene-flow clips in its flows.h5", uuid)

        cameras = self._build_cameras(raw, clips)

        trajectory_fps = _read_trajectory_fps(raw.trajectory_h5_path)
        metadata = dict(raw.metadata)
        metadata[_TRAJECTORY_FPS_METADATA_KEY] = trajectory_fps
        if raw.trajectory_h5_path is not None:
            metadata[_TRAJECTORY_H5_PATH_METADATA_KEY] = str(raw.trajectory_h5_path)

        return Episode(
            episode_id=raw.episode_id,
            task_instruction=raw.metadata.get("current_task", ""),
            trajectory_length=int(raw.metadata.get("trajectory_length", 0)),
            cameras=cameras,
            metadata=metadata,
            clips=clips,
        )

    def clip_timing(self, episode: Episode, clip: SceneFlowClip) -> ClipTiming:
        """Build :class:`~fpgm.types.ClipTiming` for one clip of a built `episode`.

        Reads `mp4_fps` from that clip's camera's mp4 container (not assumed
        equal to `trajectory_fps` -- see the module docstring). `annotation_stride`
        is *measured* (see `fpgm.data.pointworld.verify_annotation_stride`) against
        `trajectory.h5`'s own joint positions whenever both that file and the
        clip's `joint_positions` are available; otherwise this falls back to
        `ClipTiming`'s documented default and says so, rather than pretending the
        fallback is itself a measurement.
        """
        camera = episode.cameras.get(clip.camera_serial)
        if camera is None or camera.mp4_path is None:
            raise DataError(
                f"no downloaded mp4 for camera serial {clip.camera_serial!r}; "
                "cannot resolve mp4_fps for ClipTiming"
            )
        trajectory_fps = float(
            episode.metadata.get(_TRAJECTORY_FPS_METADATA_KEY, DEFAULT_TRAJECTORY_FPS)
        )
        annotation_stride = self._measure_annotation_stride(episode, clip)
        return ClipTiming(
            clip_start_frame=clip.start,
            clip_end_frame=clip.end,
            trajectory_fps=trajectory_fps,
            mp4_fps=read_mp4_fps(Path(camera.mp4_path)),
            annotation_stride=annotation_stride,
        )

    @staticmethod
    def _measure_annotation_stride(episode: Episode, clip: SceneFlowClip) -> int:
        """Best-effort call into `verify_annotation_stride`; falls back, loudly, if
        the data it needs (clip.joint_positions, a local trajectory.h5) is missing.
        """
        from fpgm.data.pointworld import verify_annotation_stride

        traj_path = episode.metadata.get(_TRAJECTORY_H5_PATH_METADATA_KEY)
        if clip.joint_positions is None or not traj_path or not Path(traj_path).exists():
            logger.warning(
                "cannot measure annotation_stride for clip %s (missing joint_positions "
                "or trajectory.h5); falling back to the unverified default %d",
                clip.key,
                ClipTiming.__dataclass_fields__["annotation_stride"].default,
            )
            return int(ClipTiming.__dataclass_fields__["annotation_stride"].default)

        with h5py.File(traj_path, "r") as f:
            trajectory_joint_positions = np.asarray(
                f["observation/robot_state/joint_positions"]
            )
        return verify_annotation_stride(
            clip.joint_positions, trajectory_joint_positions, clip.start
        )

    def _build_cameras(
        self, raw: RawEpisodeAssets, clips: list[SceneFlowClip]
    ) -> dict[str, CameraAssets]:
        # Intrinsics only exist in PointWorld's per-clip annotations, so a camera
        # with no clip at all cannot get a CameraAssets entry -- there's nothing
        # to combine with its (raw-only) extrinsic.
        clip_by_serial = {c.camera_serial: c for c in clips}

        camera_calibration: dict[str, Any] = {}
        try:
            camera_calibration = self.pointworld_store.load_camera_calibration(raw.episode_id.uuid)
        except DataError:
            logger.debug(
                "no *_cameras.json for %s; falling back to flows.h5/raw extrinsics",
                raw.episode_id,
            )

        cameras: dict[str, CameraAssets] = {}
        for role, serial in raw.camera_serial_by_role.items():
            clip = clip_by_serial.get(serial)
            if clip is None:
                logger.warning(
                    "camera %s (role=%s) has no scene-flow clip; no intrinsics available, skipping",
                    serial,
                    role,
                )
                continue
            extrinsic = self._resolve_extrinsic(serial, raw, clip, camera_calibration)
            cameras[serial] = CameraAssets(
                serial=serial,
                role=role,
                intrinsic=clip.intrinsic,
                extrinsic=extrinsic,
                mp4_path=str(raw.mp4_paths[role]) if role in raw.mp4_paths else None,
            )
        return cameras

    @staticmethod
    def _resolve_extrinsic(
        serial: str,
        raw: RawEpisodeAssets,
        clip: SceneFlowClip,
        camera_calibration: dict[str, Any],
    ) -> np.ndarray:
        """See the extrinsic precedence documented in this module's docstring."""
        camera_entry = camera_calibration.get(serial) if camera_calibration else None
        optimized = camera_entry.get("optimized_extrinsics") if camera_entry else None
        if optimized is not None:
            arr = np.asarray(optimized, dtype=np.float64)
            if arr.shape == (4, 4):
                return arr
            if arr.size == 6:
                # Same 6-DOF cam-to-world convention as raw DROID metadata; not
                # independently verified for *this* field, so this branch is a
                # defensive fallback rather than the expected shape.
                return cam2world_vector_to_world2cam(arr.reshape(6))
            logger.warning(
                "unexpected optimized_extrinsics shape %s for %s; ignoring", arr.shape, serial
            )
        if clip.extrinsic is not None:
            return clip.extrinsic
        return raw.camera_extrinsics[serial]