tests/test_geometry.py

"""Synthetic, data-free tests for fpgm.geometry."""

from __future__ import annotations

import numpy as np
import pytest

from fpgm.config import ConventionConfig
from fpgm.geometry.camera import Camera
from fpgm.geometry.convention import _bilinear_sample, detect_scene_flow_convention
from fpgm.geometry.transforms import (
    assert_valid_se3,
    invert_se3,
    matrix_to_rpy,
    pose6_to_matrix,
    pose_error,
    rpy_to_matrix,
    transform_points,
)
from fpgm.types import (
    AmbiguousConventionError,
    CameraIntrinsics,
    FrameConvention,
    GeometryError,
)


def _make_camera(position=(0.3, -0.2, 0.5), rotvec=(0.1, 0.2, -0.15)) -> Camera:
    intrinsics = CameraIntrinsics(fx=500.0, fy=480.0, cx=320.0, cy=240.0, width=640, height=480)
    world_to_cam = pose6_to_matrix(np.array(position), np.array(rotvec))
    return Camera(intrinsics, world_to_cam)


class TestSE3:
    def test_invert_se3_roundtrip(self):
        transform = pose6_to_matrix(
            np.array([1.2, -0.4, 3.1]), np.array([0.3, -0.7, 0.2])
        )
        inv = invert_se3(transform)
        assert np.allclose(transform @ inv, np.eye(4), atol=1e-10)
        assert np.allclose(inv @ transform, np.eye(4), atol=1e-10)

    def test_transform_points_matches_matrix_multiply(self):
        transform = pose6_to_matrix(np.array([1.0, 2.0, 3.0]), np.array([0.1, 0.0, 0.0]))
        pts = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [2.0, 3.0, 4.0]])
        out = transform_points(transform, pts)
        homog = np.concatenate([pts, np.ones((3, 1))], axis=1)
        expected = (transform @ homog.T).T[:, :3]
        assert np.allclose(out, expected, atol=1e-10)

    def test_assert_valid_se3_accepts_identity(self):
        assert_valid_se3(np.eye(4))

    def test_assert_valid_se3_rejects_non_orthonormal(self):
        bad = np.eye(4)
        bad[:3, :3] *= 2.0  # scaling breaks orthonormality
        with pytest.raises(GeometryError):
            assert_valid_se3(bad)

    def test_assert_valid_se3_rejects_bad_bottom_row(self):
        bad = np.eye(4)
        bad[3, :] = [0.0, 0.0, 0.1, 1.0]
        with pytest.raises(GeometryError):
            assert_valid_se3(bad)

    def test_camera_init_rejects_invalid_extrinsic(self):
        intrinsics = CameraIntrinsics(fx=1.0, fy=1.0, cx=0.0, cy=0.0, width=10, height=10)
        bad = np.eye(4)
        bad[:3, :3] *= 3.0
        with pytest.raises(GeometryError):
            Camera(intrinsics, bad)


class TestPoseError:
    def test_zero_at_identical_poses(self):
        pose = pose6_to_matrix(np.array([0.3, -0.1, 0.7]), np.array([0.1, 0.2, -0.3]))
        err = pose_error(pose, pose)
        assert np.allclose(err, np.zeros(6), atol=1e-12)

    def test_position_component_matches_direct_difference(self):
        current = pose6_to_matrix(np.array([0.0, 0.0, 0.0]), np.zeros(3))
        target = pose6_to_matrix(np.array([0.1, -0.2, 0.3]), np.zeros(3))
        err = pose_error(current, target)
        assert np.allclose(err[:3], [0.1, -0.2, 0.3], atol=1e-10)
        assert np.allclose(err[3:], np.zeros(3), atol=1e-10)

    def test_rotation_component_recovers_relative_rotvec(self):
        current = np.eye(4)
        rotvec = np.array([0.0, 0.0, np.pi / 2])
        target = pose6_to_matrix(np.zeros(3), rotvec)
        err = pose_error(current, target)
        assert np.allclose(err[3:], rotvec, atol=1e-10)


class TestRpyConvention:
    def test_rpy_to_matrix_reproduces_panda_joint_origins(self):
        # panda_joint2: rpy(-pi/2, 0, 0) -> [[1,0,0],[0,0,1],[0,-1,0]] (measured
        # against yourdfpy's own parse of the real URDF).
        rot = rpy_to_matrix(np.array([-np.pi / 2, 0.0, 0.0]))
        expected = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]])
        assert np.allclose(rot, expected, atol=1e-9)

        # panda_joint3: rpy(pi/2, 0, 0) -> the transpose-ish opposite.
        rot3 = rpy_to_matrix(np.array([np.pi / 2, 0.0, 0.0]))
        expected3 = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
        assert np.allclose(rot3, expected3, atol=1e-9)

    def test_identity_rpy_is_identity(self):
        assert np.allclose(rpy_to_matrix(np.zeros(3)), np.eye(3), atol=1e-12)

    def test_matrix_to_rpy_roundtrip(self):
        rng = np.random.default_rng(0)
        for _ in range(20):
            rpy = rng.uniform(-np.pi / 2 + 0.01, np.pi / 2 - 0.01, size=3)
            rot = rpy_to_matrix(rpy)
            recovered = matrix_to_rpy(rot)
            assert np.allclose(rpy_to_matrix(recovered), rot, atol=1e-9)


class TestCameraProjection:
    def test_project_unproject_roundtrip(self):
        cam = _make_camera()
        rng = np.random.default_rng(0)
        cam_pts = rng.uniform(low=[-0.5, -0.5, 1.0], high=[0.5, 0.5, 4.0], size=(50, 3))
        world_pts = cam.cam_to_world(cam_pts)

        uv, depth = cam.project(world_pts)
        assert np.allclose(depth, cam_pts[:, 2], atol=1e-8)

        recovered = cam.unproject(uv, depth)
        assert np.allclose(recovered, world_pts, atol=1e-6)

    def test_project_cam_matches_project_via_extrinsic(self):
        cam = _make_camera()
        rng = np.random.default_rng(1)
        cam_pts = rng.uniform(low=[-0.5, -0.5, 1.0], high=[0.5, 0.5, 4.0], size=(20, 3))
        world_pts = cam.cam_to_world(cam_pts)

        uv_world, depth_world = cam.project(world_pts)
        uv_cam, depth_cam = cam.project_cam(cam_pts)
        assert np.allclose(uv_world, uv_cam, atol=1e-8)
        assert np.allclose(depth_world, depth_cam, atol=1e-8)

    def test_project_behind_camera_does_not_raise(self):
        cam = _make_camera(position=(0, 0, 0), rotvec=(0, 0, 0))
        pts_behind = np.array([[0.0, 0.0, -1.0], [0.1, -0.1, -2.0]])
        uv, depth = cam.project(pts_behind)
        assert np.all(depth <= 0)
        assert np.all(np.isfinite(uv))  # must not raise / must not produce garbage NaNs from a raise

    def test_rescaled_projection_is_proportional(self):
        cam = _make_camera()
        cam2 = cam.rescaled(1280, 960)  # exactly 2x
        rng = np.random.default_rng(2)
        cam_pts = rng.uniform(low=[-0.5, -0.5, 1.0], high=[0.5, 0.5, 4.0], size=(10, 3))
        world_pts = cam.cam_to_world(cam_pts)

        uv1, depth1 = cam.project(world_pts)
        uv2, depth2 = cam2.project(world_pts)

        assert np.allclose(uv2, uv1 * 2.0, atol=1e-6)
        assert np.allclose(depth1, depth2, atol=1e-8)  # depth is metric, resolution-independent

    def test_rescaled_does_not_mutate_original(self):
        cam = _make_camera()
        original_width = cam.K.width
        _ = cam.rescaled(100, 100)
        assert cam.K.width == original_width


class TestConventionDetection:
    def _synthetic_scene(self, rotvec):
        """Build a camera + synthetic image where the ground truth is CAMERA-frame points."""
        intrinsics = CameraIntrinsics(fx=100.0, fy=100.0, cx=50.0, cy=50.0, width=100, height=100)
        world_to_cam = pose6_to_matrix(np.array([0.0, 0.0, 0.0]), np.array(rotvec))
        camera = Camera(intrinsics, world_to_cam)

        rng = np.random.default_rng(42)
        grid = np.linspace(-0.3, 0.3, 5)
        xs, ys = np.meshgrid(grid, grid)
        zs = np.full(xs.size, 2.0)
        cam_pts = np.stack([xs.ravel(), ys.ravel(), zs], axis=1)  # ground truth: CAMERA frame

        initial_rgb = rng.integers(0, 256, size=(100, 100, 3)).astype(np.uint8)
        uv_true, _ = camera.project_cam(cam_pts)
        scene_colors = _bilinear_sample(initial_rgb, uv_true).astype(np.uint8)
        return camera, cam_pts, scene_colors, initial_rgb

    def test_detects_camera_convention_by_construction(self):
        # A 180-degree rotation about X flips the sign of Z for the (wrong) WORLD
        # hypothesis, pushing every point behind the camera -- a decisive signal.
        camera, cam_pts, scene_colors, initial_rgb = self._synthetic_scene(
            rotvec=[np.pi, 0.0, 0.0]
        )
        cfg = ConventionConfig()
        result = detect_scene_flow_convention(camera, cam_pts, scene_colors, initial_rgb, cfg)
        assert result.convention == FrameConvention.CAMERA
        assert result.margin > cfg.min_margin
        assert result.frac_in_bounds > cfg.min_frac_in_bounds

    def test_raises_ambiguous_when_hypotheses_tie(self):
        # An identity extrinsic makes project() and project_cam() identical, so both
        # hypotheses score exactly the same regardless of the scene -- the detector
        # must refuse to guess rather than pick one arbitrarily.
        intrinsics = CameraIntrinsics(fx=100.0, fy=100.0, cx=50.0, cy=50.0, width=100, height=100)
        camera = Camera(intrinsics, np.eye(4))

        rng = np.random.default_rng(7)
        grid = np.linspace(-0.3, 0.3, 5)
        xs, ys = np.meshgrid(grid, grid)
        zs = np.full(xs.size, 2.0)
        pts = np.stack([xs.ravel(), ys.ravel(), zs], axis=1)
        initial_rgb = rng.integers(0, 256, size=(100, 100, 3)).astype(np.uint8)
        uv_true, _ = camera.project_cam(pts)
        scene_colors = _bilinear_sample(initial_rgb, uv_true).astype(np.uint8)

        cfg = ConventionConfig()
        with pytest.raises(AmbiguousConventionError):
            detect_scene_flow_convention(camera, pts, scene_colors, initial_rgb, cfg)