posed_video

@dataclass(frozen=True)
class PosedFrame:
    """One frame of a ``PosedVideo``: its pose and the depth measured at it.

    A pure data record — decode its RGB pixels through the owning clip with
    :meth:`PosedVideo.image_at`. ``depth`` is the metric map the sensor measured
    at *this* frame, or ``None`` when this frame carries no fresh depth (depth is
    sparser than the RGB stream), so ``depth is not None`` marks exactly the
    frames that measured it.

    Attributes:
        time: The frame's capture instant (the shared pose/depth ``Map`` key).
        pose: The frame's 6-DoF camera pose, with intrinsics.
        depth: The depth measured at this frame, or ``None`` if it measured none.
        index: The frame's position in the clip — the pose's rank in key order,
            which the recorder guarantees is the video frame's index.
    """

    time: datetime
    pose: Pose
    depth: Depth | None
    index: int

    def point_cloud(self) -> np.ndarray:
        """Back-project this frame's depth into ``N×3`` camera-space metres via ``pose``.

        Raises:
            ValueError: If this frame measured no depth.
        """
        if self.depth is None:
            raise ValueError(f"frame {self.time.isoformat()} measured no depth")
        return self.depth.point_cloud(self.pose)

class PosedVideo(Struct):
    """A world-anchored posed RGB-D video: a clip, its poses, and its depth.

    ``poses`` is dense — one per video frame, keyed by that frame's instant;
    ``depth`` is sparse — only the frames a depth sensor reached, keyed on the
    same ``Time`` namespace so it lines up with ``poses``. Use :meth:`frames` to
    walk the clip with each frame's pose and (present-or-absent) depth in step.

    The dense per-frame pose is *also* baked into the ``video`` file itself (a
    timed metadata track), so the ``.mov`` is a self-contained artifact; ``poses``
    is the schema-advertised copy CaptureGraph reads back without demuxing.

    A depth entry is keyed by the exact instant of the pose (video frame) it was
    measured at, so ``depth`` is a subset of ``poses``' keys — that shared key is
    the frame↔depth correspondence.

    Attributes:
        world_map: The ARKit world map the clip is anchored in (relocalize later
            sessions into the same coordinate frame).
        video: The recorded RGB clip.
        poses: Per-frame 6-DoF poses, keyed by frame instant (one per frame).
        depth: Best-effort per-frame metric depth, keyed by frame instant (only
            the frames the depth sensor reached — sparser than ``poses``).
    """

    world_map: ARWorldMap
    video: Video
    poses: Map[Time, Pose]
    depth: Map[Time, Depth]

    def frames(self) -> Iterator[PosedFrame]:
        """Each frame in chronological order, with the depth measured at it.

        Yields one :class:`PosedFrame` per pose entry (per video frame). A
        frame's ``depth`` is the map the sensor measured at that exact instant,
        or ``None`` — so ``frame.depth is not None`` marks precisely the frames
        that carry fresh depth.
        """
        poses = sorted(self.poses.to_dict().items(), key=lambda entry: entry[0])
        depth = self.depth.to_dict()
        for index, (time, pose) in enumerate(poses):
            yield PosedFrame(
                time=time,
                pose=pose,
                depth=depth.get(time),
                index=index,
            )

    def image_at(self, frame: PosedFrame) -> PILImage.Image:
        """Decode ``frame``'s RGB image from the clip (exact, by frame index)."""
        return self.video.frame_at_index(frame.index)