src/fpgm/objects/scene.py

"""Stage 5: render robot + object mesh + point cloud into one image, world frame.

Everything here lives in the world (robot-base) frame established by stages 1-4:
the robot is posed by forward kinematics, the object by
:class:`~fpgm.objects.types.ObjectTrajectory`'s per-frame transform, and the
observed point cloud is carried along unchanged from stage 3. Putting all three
in one image is the entire point of reconstructing a mesh instead of drawing a
2D sprite: only a real mesh, rendered through a real depth buffer, can show
which of the robot and the object is actually in front at a given pixel.

:class:`SceneRenderer` does exactly that -- it reuses
:class:`~fpgm.robot.render.RobotRenderer` wholesale (same persistent
``pyrender.Scene``, same OSMesa offscreen context) and adds the object as one
more node in that same scene, so pyrender's own z-buffer resolves robot/object
occlusion instead of any hand-rolled compositing order.

The rest of this module (:func:`project_point_cloud`, :func:`draw_point_cloud`,
:func:`free_camera`, :func:`state_color`, :func:`draw_scene_hud`) is pure numpy
+ OpenCV and needs no GL context at all.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import trimesh
from PIL import Image

from fpgm.geometry.camera import Camera
from fpgm.geometry.transforms import invert_se3
from fpgm.objects.types import InteractionState, ObjectMesh
from fpgm.robot.render import RobotRenderer
from fpgm.types import CameraIntrinsics
from fpgm.utils.logging import get_logger
from fpgm.viz.overlays import draw_hud

logger = get_logger(__name__)

# fpgm.robot.render (imported above) calls _configure_headless_gl() -- setting
# PYOPENGL_PLATFORM=osmesa and stubbing pyrender.viewer -- as a side effect of
# its own module import, *before* it imports pyrender itself. Importing it
# first guarantees that has already happened by the time this line runs, so
# this module never needs (and must never duplicate) that setup itself.
import pyrender  # noqa: E402

#: Robot-base (``panda_link0``) frame convention assumed throughout this
#: module: +Z is up. This matches the physical Franka mount and is the only
#: place that assumption is baked in -- see :func:`free_camera`.
_WORLD_UP = np.array([0.0, 0.0, 1.0])


@dataclass
class SceneRenderResult:
    """One rendered frame of robot + object, from a single pyrender scene.

    ``robot_mask`` and ``object_mask`` are disjoint by construction -- see
    :meth:`SceneRenderer.render` for how they are computed.
    """

    color: np.ndarray  # (H, W, 3) uint8, RGB
    depth: np.ndarray  # (H, W) float32, metres; 0 where nothing was hit
    robot_mask: np.ndarray  # (H, W) bool
    object_mask: np.ndarray  # (H, W) bool

    @property
    def mask(self) -> np.ndarray:
        """``robot_mask | object_mask``.

        Lets a :class:`SceneRenderResult` duck-type as a
        :class:`~fpgm.robot.render.RenderResult` for
        :func:`fpgm.robot.overlay.composite`, which only ever reads ``.color``
        and ``.mask``.
        """
        return self.robot_mask | self.object_mask


class SceneRenderer:
    """Offscreen renderer for the robot AND its manipulated object, in one scene.

    All robot geometry/camera/lighting setup is delegated to
    :class:`~fpgm.robot.render.RobotRenderer` -- this class does not
    reimplement any of it, per that module's own admonition that scene
    construction (not per-frame rendering) dominates runtime on OSMesa's
    software rasteriser. The object gets exactly one extra node in
    ``RobotRenderer``'s persistent ``pyrender.Scene``, re-posed per frame
    exactly like a robot link, never rebuilt.

    **Mask strategy** (chosen deliberately; the alternative -- per-node colour
    tagging read back from the colour buffer -- was rejected because this
    OSMesa build does not reliably expose an object-id/alpha channel, see
    :class:`~fpgm.robot.render.RobotRenderer`'s own note on ``RenderFlags.RGBA``):

    1. Render robot + object together (one pass) for the real composited
       image and combined depth -- this is the pass where occlusion between
       them is resolved correctly, by pyrender's own z-buffer.
    2. Temporarily detach the robot's link nodes and render the object alone
       (second pass, same camera/light state) to get the object's own,
       unoccluded depth at every pixel it would cover if nothing were in front
       of it.
    3. A pixel counts as "object" wherever the object-alone depth agrees with
       the combined depth (within :attr:`_MASK_DEPTH_ATOL`) -- i.e. the object
       really is the nearest surface there. Every other combined-mask pixel is
       "robot" -- with only two entities ever in the scene, that is exact,
       including at silhouette boundaries where the arm partially occludes the
       object.

    This costs one extra render call per frame (cheap relative to scene
    construction) and needs no changes to :class:`RobotRenderer` itself.
    """

    #: Depth agreement tolerance (metres) between the object-alone pass and the
    #: combined pass, for classifying a pixel as "object visible here".
    _MASK_DEPTH_ATOL = 1e-4

    def __init__(
        self,
        link_meshes: dict[str, list[tuple[trimesh.Trimesh, np.ndarray]]],
        object_mesh: ObjectMesh,
        width: int,
        height: int,
        *,
        object_color_rgb: tuple[int, int, int] = (220, 120, 60),
        **robot_kwargs: object,
    ) -> None:
        """Build the persistent robot+object scene graph.

        Args:
            link_meshes: Robot visual geometry -- see
                :class:`~fpgm.robot.render.RobotRenderer`.
            object_mesh: The manipulated object's mesh, in its own canonical
                frame (e.g. from :mod:`fpgm.objects.proxy` or TRELLIS). Re-posed
                per frame via :meth:`render`.
            width: Render width in pixels.
            height: Render height in pixels.
            object_color_rgb: Flat RGB fill used when `object_mesh` carries no
                per-vertex colour, so a proxy mesh is visually distinguishable
                from the robot rather than falling back to pyrender's default
                grey.
            **robot_kwargs: Forwarded to :class:`~fpgm.robot.render.RobotRenderer`
                (``ambient``, ``bg_color``, ``znear``, ``zfar``,
                ``light_intensity``).

        Raises:
            RenderError: If the OSMesa offscreen GL context cannot be created
                (raised by the underlying ``RobotRenderer``).
        """
        self._robot = RobotRenderer(link_meshes, width, height, **robot_kwargs)
        self.width = self._robot.width
        self.height = self._robot.height

        pmesh = self._build_object_pyrender_mesh(object_mesh, object_color_rgb)
        self._object_node = pyrender.Node(mesh=pmesh, matrix=np.eye(4))
        self._robot.scene.add_node(self._object_node)

        # Flat list of every robot link's pyrender.Node -- toggled off/on for
        # the object-alone mask pass in render(); see the class docstring.
        # Reaches into RobotRenderer's "private" _link_nodes/_renderer rather
        # than duplicating its scene-graph bookkeeping, which is exactly the
        # integration point RobotRenderer's own docstring points at (build the
        # scene once, only update poses per frame).
        self._link_nodes = [
            node for entries in self._robot._link_nodes.values() for node, _ in entries
        ]

    @staticmethod
    def _build_object_pyrender_mesh(
        object_mesh: ObjectMesh, object_color_rgb: tuple[int, int, int]
    ) -> pyrender.Mesh:
        # process=False: a textured mesh's `uv` is per-vertex, so trimesh's
        # default vertex welding would desync a UV row from the vertex it
        # belongs to (same reason scripts/run_object_pipeline.py's mesh save
        # /load round-trip disables it).
        if object_mesh.uv is not None and object_mesh.texture is not None:
            visual = trimesh.visual.TextureVisuals(
                uv=np.asarray(object_mesh.uv, dtype=np.float32),
                image=Image.fromarray(np.asarray(object_mesh.texture, dtype=np.uint8)),
            )
            tri = trimesh.Trimesh(
                vertices=np.asarray(object_mesh.vertices, dtype=np.float64),
                faces=np.asarray(object_mesh.faces, dtype=np.int64),
                visual=visual,
                process=False,
            )
            # Deliberately no vertex_colors assignment on this branch: setting
            # `tri.visual.vertex_colors` replaces `visual` with a fresh
            # ColorVisuals wholesale (trimesh visuals are one-kind-at-a-time),
            # which is exactly how the texture used to get silently thrown
            # away here even after it was loaded correctly.
        else:
            tri = trimesh.Trimesh(
                vertices=np.asarray(object_mesh.vertices, dtype=np.float64),
                faces=np.asarray(object_mesh.faces, dtype=np.int64),
                process=False,
            )
            if object_mesh.vertex_colors is not None:
                tri.visual.vertex_colors = object_mesh.vertex_colors
            else:
                fill = np.array([*object_color_rgb, 255], dtype=np.uint8)
                tri.visual.vertex_colors = np.tile(fill, (len(tri.vertices), 1))
        # smooth=False: consistent with RobotRenderer's own meshes (see its
        # docstring) and harmless here even without per-face colours.
        return pyrender.Mesh.from_trimesh(tri, smooth=False)

    def render(
        self,
        link_poses: dict[str, np.ndarray],
        object_pose: np.ndarray,
        camera: Camera,
    ) -> SceneRenderResult:
        """Render one frame: robot posed by `link_poses`, object posed by `object_pose`.

        Args:
            link_poses: See :meth:`~fpgm.robot.render.RobotRenderer.render`.
            object_pose: ``(4, 4)`` mesh-canonical -> world transform for this
                frame, e.g. one row of
                :attr:`~fpgm.objects.types.ObjectTrajectory.transforms`.
            camera: Camera to render from; must match this renderer's resolution.

        Returns:
            A :class:`SceneRenderResult` with disjoint ``robot_mask``/``object_mask``.

        Raises:
            ValueError: If `camera`'s resolution does not match this renderer's.
            KeyError: If `link_poses` is missing a link this renderer has geometry for.
        """
        self._robot.scene.set_pose(self._object_node, np.asarray(object_pose, dtype=np.float64))
        combined = self._robot.render(link_poses, camera)

        for node in self._link_nodes:
            self._robot.scene.remove_node(node)
        try:
            _, obj_depth = self._robot._renderer.render(self._robot.scene)
        finally:
            for node in self._link_nodes:
                self._robot.scene.add_node(node)

        obj_depth = np.asarray(obj_depth, dtype=np.float32)
        combined_mask = combined.mask
        object_mask = (
            combined_mask
            & (obj_depth > 0.0)
            & (np.abs(obj_depth - combined.depth) <= self._MASK_DEPTH_ATOL)
        )
        robot_mask = combined_mask & ~object_mask
        return SceneRenderResult(
            color=combined.color,
            depth=combined.depth,
            robot_mask=robot_mask,
            object_mask=object_mask,
        )

    def close(self) -> None:
        """Release the underlying GL context. Safe to call more than once."""
        self._robot.close()

    def __enter__(self) -> SceneRenderer:
        return self

    def __exit__(self, *exc_info: object) -> None:
        self.close()


def project_point_cloud(
    points_world: np.ndarray,
    colors: np.ndarray,
    camera: Camera,
    image_shape: tuple[int, int],
) -> tuple[np.ndarray, np.ndarray]:
    """Project a world-frame point cloud to pixels, ordered nearest-drawn-last.

    Points behind the camera or outside `image_shape` are dropped -- there is
    nothing meaningful to draw for them. The rest are sorted by depth,
    farthest first, so a caller drawing them in the returned order paints
    nearer points last (i.e. on top), which is a correct occlusion order for
    dots with no size/depth-buffer of their own.

    Args:
        points_world: ``(N, 3)`` float, world-frame points.
        colors: ``(N, 3)`` uint8, one colour per point, same order as
            `points_world`.
        camera: Camera to project through.
        image_shape: ``(H, W)`` of the target image.

    Returns:
        ``(uv, colors_sorted)``: ``(M, 2)`` float pixel coordinates and
        ``(M, 3)`` uint8 colours, ``M <= N``, ordered far-to-near.

    Raises:
        ValueError: If `points_world` and `colors` disagree on point count.
    """
    points_world = np.asarray(points_world, dtype=np.float64)
    colors = np.asarray(colors)
    if points_world.shape[0] != colors.shape[0]:
        raise ValueError(
            f"project_point_cloud: {points_world.shape[0]} points but "
            f"{colors.shape[0]} colours"
        )

    uv, depth = camera.project(points_world)
    height, width = image_shape[0], image_shape[1]
    in_front = depth > 0.0
    in_bounds = (uv[:, 0] >= 0) & (uv[:, 0] < width) & (uv[:, 1] >= 0) & (uv[:, 1] < height)
    keep = in_front & in_bounds

    uv, depth, colors = uv[keep], depth[keep], colors[keep]
    order = np.argsort(depth)[::-1]  # descending depth: farthest first, nearest last
    return uv[order], colors[order]


def draw_point_cloud(
    frame_bgr: np.ndarray, uv: np.ndarray, colors: np.ndarray, radius: int = 1
) -> np.ndarray:
    """Draw projected point-cloud dots onto a frame. Pure cv2, no GL involved.

    Args:
        frame_bgr: ``(H, W, 3)`` uint8 BGR frame.
        uv: ``(N, 2)`` pixel coordinates, typically from
            :func:`project_point_cloud` (already ordered so nearer points draw
            last / on top).
        colors: ``(N, 3)`` uint8 colours, interpreted as RGB (matching
            :attr:`~fpgm.objects.types.ObjectAlignment.point_colors` and
            pyrender's RGB convention) and converted to BGR for drawing.
        radius: Dot radius in pixels.

    Returns:
        A copy of `frame_bgr` with the dots drawn; the input is not mutated.
    """
    import cv2

    out = frame_bgr.copy()
    uv = np.asarray(uv)
    colors = np.asarray(colors)
    for i in range(uv.shape[0]):
        center = tuple(np.round(uv[i]).astype(int))
        color_bgr = (int(colors[i, 2]), int(colors[i, 1]), int(colors[i, 0]))
        cv2.circle(out, center, radius, color_bgr, -1, lineType=cv2.LINE_AA)
    return out


def free_camera(
    look_at: np.ndarray,
    distance: float,
    azimuth_deg: float,
    elevation_deg: float,
    intrinsics: CameraIntrinsics,
) -> Camera:
    """Build a synthetic camera orbiting `look_at` at a fixed distance.

    The recorded (episode) camera cannot show whether the reconstructed object
    is correctly placed *in depth* -- a wrong depth still looks plausible from
    the one viewpoint that generated it. A second, independent viewpoint of the
    same world is the only honest way to see that 3D relationship, which is
    what this is for.

    World up is assumed to be ``+Z`` (the ``panda_link0`` convention this whole
    pipeline works in -- see the module docstring).

    Args:
        look_at: ``(3,)`` world-frame point the camera points at.
        distance: Orbit radius, metres.
        azimuth_deg: Rotation around the world +Z axis, degrees. 0 puts the
            camera on the world +X side of `look_at`.
        elevation_deg: Angle above the horizontal plane through `look_at`,
            degrees. 0 is level, 90 is looking straight down.
        intrinsics: Pinhole intrinsics for the returned camera.

    Returns:
        A :class:`~fpgm.geometry.camera.Camera` positioned on the orbit and
        looking exactly at `look_at`.
    """
    look_at = np.asarray(look_at, dtype=np.float64).reshape(3)
    az, el = math.radians(azimuth_deg), math.radians(elevation_deg)
    offset = distance * np.array(
        [math.cos(el) * math.cos(az), math.cos(el) * math.sin(az), math.sin(el)]
    )
    cam_pos = look_at + offset

    forward = look_at - cam_pos
    forward_norm = np.linalg.norm(forward)
    if forward_norm < 1e-9:
        raise ValueError("free_camera: distance must be > 0")
    forward = forward / forward_norm

    right = np.cross(forward, _WORLD_UP)
    right_norm = np.linalg.norm(right)
    if right_norm < 1e-6:
        # Looking straight up/down world +Z: world_up gives no constraint on
        # "right", so fall back to a fixed reference axis instead.
        fallback = np.array([1.0, 0.0, 0.0])
        right = np.cross(forward, fallback)
        right_norm = np.linalg.norm(right)
        if right_norm < 1e-6:
            fallback = np.array([0.0, 1.0, 0.0])
            right = np.cross(forward, fallback)
            right_norm = np.linalg.norm(right)
    right = right / right_norm
    down = np.cross(forward, right)  # already unit: forward, right orthonormal

    cam_to_world = np.eye(4, dtype=np.float64)
    cam_to_world[:3, 0] = right
    cam_to_world[:3, 1] = down
    cam_to_world[:3, 2] = forward
    cam_to_world[:3, 3] = cam_pos

    world_to_cam = invert_se3(cam_to_world)
    return Camera(intrinsics, world_to_cam)


#: BGR colours (cv2 convention), one per InteractionState, chosen for maximum
#: separation: grey (nothing happening), amber (in contact), crimson (attached).
_STATE_COLORS_BGR: dict[InteractionState, tuple[int, int, int]] = {
    InteractionState.FREE: (150, 150, 150),
    InteractionState.PUSHED: (0, 165, 255),
    InteractionState.GRASPED: (40, 30, 220),
}


def state_color(state: InteractionState) -> tuple[int, int, int]:
    """Stable BGR colour for an :class:`~fpgm.objects.types.InteractionState`."""
    return _STATE_COLORS_BGR[state]


def draw_scene_hud(
    frame_bgr: np.ndarray,
    frame_idx: int,
    timestamp_s: float,
    state: InteractionState,
    object_speed_mps: float,
    push_gain: float,
    origin: tuple[int, int] = (12, 12),
) -> np.ndarray:
    """Draw the stage-5 HUD: frame/time, interaction state, object speed, push_gain.

    `push_gain` is drawn as a loud, explicit warning line whenever it is not
    ``1.0`` -- pushed motion is deliberately amplified for visibility (see
    :class:`fpgm.objects.interaction.InteractionConfig`), and that amplified
    motion must never be mistaken for a physical measurement.

    Args:
        frame_bgr: ``(H, W, 3)`` uint8 BGR frame to draw over.
        frame_idx: Frame index, for the HUD text.
        timestamp_s: Frame timestamp, seconds.
        state: This frame's :class:`~fpgm.objects.types.InteractionState`.
        object_speed_mps: Object speed at this frame, metres/second.
        push_gain: The trajectory's
            :attr:`~fpgm.objects.types.ObjectTrajectory.push_gain`.
        origin: Top-left corner of the HUD panel, pixels.

    Returns:
        A copy of `frame_bgr` with the HUD drawn on.
    """
    lines = [
        f"frame={frame_idx}  t={timestamp_s:.2f}s",
        f"state={state.value}",
        f"object speed = {object_speed_mps:.3f} m/s",
    ]
    if not math.isclose(push_gain, 1.0):
        lines.append(f"!! push_gain={push_gain:.1f}x -- VISUALISATION ONLY, not physical !!")
    return draw_hud(frame_bgr, lines, origin=origin)