tests/test_object_scene.py

"""Tests for fpgm.objects.scene.

Deliberately independent of fpgm.objects.{crop,proxy,align,interaction} (other
workstreams): every non-GL test below builds its own synthetic points/camera
directly, exactly matching the shapes those modules produce. The one GL test
mirrors tests/test_robot_render.py's pattern of skipping cleanly when no OSMesa
offscreen context can be created.
"""

from __future__ import annotations

import numpy as np
import pytest
import trimesh

from fpgm.geometry.camera import Camera
from fpgm.objects.scene import (
    SceneRenderer,
    draw_point_cloud,
    draw_scene_hud,
    free_camera,
    project_point_cloud,
    state_color,
)
from fpgm.objects.types import InteractionState, MeshSource, ObjectMesh
from fpgm.robot.render import RenderError, RobotRenderer
from fpgm.types import CameraIntrinsics


def _camera(width=160, height=120, fx=150.0, fy=150.0, camera_z=-6.0) -> Camera:
    """A camera at world ``(0, 0, camera_z)``, identity rotation, +Z forward."""
    intrinsics = CameraIntrinsics(fx=fx, fy=fy, cx=width / 2.0, cy=height / 2.0,
                                   width=width, height=height)
    world_to_cam = np.eye(4)
    world_to_cam[2, 3] = -camera_z
    return Camera(intrinsics, world_to_cam)


class TestProjectPointCloud:
    def test_uv_matches_camera_project_and_orders_far_to_near(self) -> None:
        camera = _camera()
        # Off-axis in x, distinct world-z depths; colours double as point ids.
        points = np.array([[0.0, 0.0, 5.0], [0.1, 0.0, 3.0], [-0.1, 0.0, 1.0]])
        colors = np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255]], dtype=np.uint8)

        uv, colors_sorted = project_point_cloud(points, colors, camera, (120, 160))

        assert uv.shape == (3, 2)
        assert colors_sorted.shape == (3, 3)

        # Farthest world-z (red) first, nearest (blue) last.
        assert np.array_equal(colors_sorted, np.array([[255, 0, 0], [0, 255, 0], [0, 0, 255]]))

        # Each returned uv exactly matches Camera.project of its own point.
        expected_uv, expected_depth = camera.project(points)
        by_color = {tuple(c): i for i, c in enumerate(colors)}
        for row_uv, row_color in zip(uv, colors_sorted, strict=True):
            i = by_color[tuple(row_color.tolist())]
            assert row_uv == pytest.approx(expected_uv[i])

        # z-ordering: depths of the returned rows are strictly descending.
        depths_in_order = [expected_depth[by_color[tuple(c.tolist())]] for c in colors_sorted]
        assert depths_in_order == sorted(depths_in_order, reverse=True)

    def test_drops_points_behind_camera_and_out_of_bounds(self) -> None:
        camera = _camera(width=160, height=120)
        points = np.array(
            [
                [0.0, 0.0, 5.0],  # kept: in front, in bounds
                [0.0, 0.0, -10.0],  # dropped: behind the camera (camera sits at world z=-6)
                [100.0, 0.0, 1.0],  # dropped: projects far outside the image
            ]
        )
        colors = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]], dtype=np.uint8)

        uv, colors_sorted = project_point_cloud(points, colors, camera, (120, 160))

        assert uv.shape[0] == 1
        assert np.array_equal(colors_sorted[0], [1, 1, 1])

    def test_mismatched_lengths_raises(self) -> None:
        camera = _camera()
        with pytest.raises(ValueError):
            project_point_cloud(
                np.zeros((2, 3)), np.zeros((3, 3), dtype=np.uint8), camera, (120, 160)
            )


class TestDrawPointCloud:
    def test_draws_only_near_the_point_and_preserves_color(self) -> None:
        frame = np.zeros((50, 50, 3), dtype=np.uint8)
        original = frame.copy()
        uv = np.array([[25.0, 25.0]])
        colors_rgb = np.array([[10, 20, 30]], dtype=np.uint8)  # R=10 G=20 B=30

        out = draw_point_cloud(frame, uv, colors_rgb, radius=2)

        # Input not mutated.
        assert np.array_equal(frame, original)
        # Something was drawn at the centre, in BGR order (B=30, G=20, R=10).
        assert tuple(int(v) for v in out[25, 25]) == (30, 20, 10)
        # Far from the drawn point, nothing changed.
        assert np.array_equal(out[0:10, 0:10], original[0:10, 0:10])
        assert np.array_equal(out[40:50, 40:50], original[40:50, 40:50])

    def test_multiple_points_each_draw_their_own_color(self) -> None:
        frame = np.zeros((30, 30, 3), dtype=np.uint8)
        uv = np.array([[5.0, 5.0], [24.0, 24.0]])
        colors_rgb = np.array([[255, 0, 0], [0, 255, 0]], dtype=np.uint8)

        out = draw_point_cloud(frame, uv, colors_rgb, radius=2)

        assert tuple(int(v) for v in out[5, 5]) == (0, 0, 255)  # BGR for RGB red
        assert tuple(int(v) for v in out[24, 24]) == (0, 255, 0)  # BGR for RGB green


class TestFreeCamera:
    def _intrinsics(self, width=160, height=120) -> CameraIntrinsics:
        return CameraIntrinsics(fx=150.0, fy=150.0, cx=width / 2.0, cy=height / 2.0,
                                 width=width, height=height)

    def test_looks_at_target(self) -> None:
        intrinsics = self._intrinsics()
        target = np.array([0.3, -0.1, 0.5])
        camera = free_camera(target, distance=2.0, azimuth_deg=37.0, elevation_deg=15.0,
                              intrinsics=intrinsics)

        uv, depth = camera.project(target.reshape(1, 3))
        assert depth[0] == pytest.approx(2.0, abs=1e-6)
        assert uv[0, 0] == pytest.approx(intrinsics.cx, abs=1e-6)
        assert uv[0, 1] == pytest.approx(intrinsics.cy, abs=1e-6)

    def test_azimuth_180_is_the_opposite_side(self) -> None:
        intrinsics = self._intrinsics()
        target = np.array([0.0, 0.0, 0.0])
        cam_a = free_camera(target, distance=3.0, azimuth_deg=20.0, elevation_deg=0.0,
                             intrinsics=intrinsics)
        cam_b = free_camera(target, distance=3.0, azimuth_deg=200.0, elevation_deg=0.0,
                             intrinsics=intrinsics)

        pos_a = cam_a.cam_to_world(np.zeros((1, 3)))[0]
        pos_b = cam_b.cam_to_world(np.zeros((1, 3)))[0]

        # At elevation 0, azimuth+180 exactly mirrors the camera through the target.
        assert pos_a + pos_b == pytest.approx(2.0 * target, abs=1e-6)
        assert np.linalg.norm(pos_a - target) == pytest.approx(np.linalg.norm(pos_b - target))

    def test_still_looks_at_target_when_looking_straight_down(self) -> None:
        intrinsics = self._intrinsics()
        target = np.array([1.0, 1.0, 1.0])
        camera = free_camera(target, distance=1.5, azimuth_deg=0.0, elevation_deg=90.0,
                              intrinsics=intrinsics)
        uv, depth = camera.project(target.reshape(1, 3))
        assert depth[0] == pytest.approx(1.5, abs=1e-6)
        assert uv[0] == pytest.approx([intrinsics.cx, intrinsics.cy], abs=1e-6)

    def test_zero_distance_raises(self) -> None:
        with pytest.raises(ValueError):
            free_camera(np.zeros(3), distance=0.0, azimuth_deg=0.0, elevation_deg=0.0,
                         intrinsics=self._intrinsics())


class TestStateColor:
    def test_distinct_colors_per_state(self) -> None:
        colors = {state_color(s) for s in InteractionState}
        assert len(colors) == len(list(InteractionState))
        for c in colors:
            assert len(c) == 3
            assert all(0 <= v <= 255 for v in c)


class TestDrawSceneHud:
    def test_shows_push_gain_warning_only_when_amplified(self) -> None:
        # Large enough that neither HUD panel is clipped by the frame edge --
        # a clipped panel would make both variants saturate to the same size.
        frame = np.zeros((400, 700, 3), dtype=np.uint8)

        physical = draw_scene_hud(
            frame, frame_idx=3, timestamp_s=0.2, state=InteractionState.PUSHED,
            object_speed_mps=0.05, push_gain=1.0,
        )
        amplified = draw_scene_hud(
            frame, frame_idx=3, timestamp_s=0.2, state=InteractionState.PUSHED,
            object_speed_mps=0.05, push_gain=4.0,
        )

        # Both draw a HUD panel (frame is modified relative to the blank input).
        assert not np.array_equal(physical, frame)
        assert not np.array_equal(amplified, frame)
        # The amplified HUD has one more line -> a taller panel -> strictly more
        # pixels touched than the non-amplified one.
        changed_physical = int(np.count_nonzero(np.any(physical != frame, axis=-1)))
        changed_amplified = int(np.count_nonzero(np.any(amplified != frame, axis=-1)))
        assert changed_amplified > changed_physical

    def test_does_not_mutate_input(self) -> None:
        frame = np.zeros((80, 200, 3), dtype=np.uint8)
        original = frame.copy()
        draw_scene_hud(frame, 0, 0.0, InteractionState.FREE, 0.0, push_gain=1.0)
        assert np.array_equal(frame, original)


def _build_robot_renderer(link_meshes, width, height) -> RobotRenderer:
    """Construct a RobotRenderer, skipping the test if OSMesa is unavailable."""
    try:
        return RobotRenderer(link_meshes, width, height)
    except RenderError as exc:  # pragma: no cover - environment-dependent
        pytest.skip(f"OSMesa offscreen GL context unavailable: {exc}")


def _build_scene_renderer(link_meshes, object_mesh, width, height) -> SceneRenderer:
    try:
        return SceneRenderer(link_meshes, object_mesh, width, height)
    except RenderError as exc:  # pragma: no cover - environment-dependent
        pytest.skip(f"OSMesa offscreen GL context unavailable: {exc}")


class TestSceneRendererOcclusion:
    """The scene-5 smoke test: robot box + object box overlapping in depth.

    A unit robot "link" box sits at the origin; a smaller object box is placed
    fully in front of it (nearer the camera) and fully within its silhouette,
    so the expected relationship between the three masks is exact and checked
    numerically, not just "something rendered".
    """

    def _scene(self):
        width, height = 160, 120
        intrinsics = CameraIntrinsics(fx=150.0, fy=150.0, cx=80.0, cy=60.0,
                                       width=width, height=height)
        world_to_cam = np.eye(4)
        world_to_cam[2, 3] = 6.0  # camera at world (0, 0, -6), looking down +Z
        camera = Camera(intrinsics, world_to_cam)

        robot_box = trimesh.creation.box(extents=[1.0, 1.0, 1.0])
        link_meshes = {"robot_link": [(robot_box, np.eye(4))]}
        link_poses = {"robot_link": np.eye(4)}

        obj_box = trimesh.creation.box(extents=[0.4, 0.4, 0.4])
        object_mesh = ObjectMesh(
            vertices=np.asarray(obj_box.vertices, dtype=np.float64),
            faces=np.asarray(obj_box.faces, dtype=np.int64),
            source=MeshSource.PROXY_BOX,
        )
        object_pose = np.eye(4)
        object_pose[2, 3] = -1.0  # strictly nearer the camera than the robot box

        return width, height, camera, link_meshes, link_poses, object_mesh, object_pose

    def test_masks_are_disjoint_and_occlusion_is_exact(self) -> None:
        width, height, camera, link_meshes, link_poses, object_mesh, object_pose = self._scene()

        robot_alone = _build_robot_renderer(link_meshes, width, height)
        try:
            reference = robot_alone.render(link_poses, camera)
        finally:
            robot_alone.close()

        scene_renderer = _build_scene_renderer(link_meshes, object_mesh, width, height)
        try:
            result = scene_renderer.render(link_poses, object_pose, camera)
        finally:
            scene_renderer.close()

        assert result.color.shape == (height, width, 3)
        assert result.depth.shape == (height, width)
        assert not (result.robot_mask & result.object_mask).any()
        assert np.array_equal(result.mask, result.robot_mask | result.object_mask)
        assert result.object_mask.any()
        assert result.robot_mask.any()

        # The object sits fully in front of, and fully within the silhouette
        # of, the robot box in this setup -- so punching the object's mask out
        # of the robot-alone reference mask must reproduce the combined
        # scene's robot_mask exactly.
        assert np.array_equal(result.robot_mask, reference.mask & ~result.object_mask)

    def test_close_is_idempotent(self) -> None:
        width, height, camera, link_meshes, link_poses, object_mesh, object_pose = self._scene()
        scene_renderer = _build_scene_renderer(link_meshes, object_mesh, width, height)
        with scene_renderer:
            scene_renderer.render(link_poses, object_pose, camera)
        scene_renderer.close()  # closing an already-closed renderer must not raise