tests/test_object_interaction.py

"""Synthetic, data-free tests for fpgm.objects.interaction.

No URDF and no recorded episode data: a small stub stands in for
:class:`~fpgm.robot.urdf.RobotModel`, returning pre-scripted per-frame link
poses. ``InteractionModel.solve`` only ever calls ``link_poses(...)``, once per
frame in order, so a call-counter is enough to script an exact scenario.
"""

from __future__ import annotations

import numpy as np

from fpgm.objects.interaction import InteractionConfig, InteractionModel, summarize
from fpgm.objects.types import InteractionState, ObjectAlignment


def _pose(pos: list[float] | np.ndarray, rot: np.ndarray | None = None) -> np.ndarray:
    """Build a (4, 4) SE3 matrix from a position and an optional (3, 3) rotation."""
    mat = np.eye(4, dtype=np.float64)
    if rot is not None:
        mat[:3, :3] = rot
    mat[:3, 3] = np.asarray(pos, dtype=np.float64)
    return mat


def _rotz(radians: float) -> np.ndarray:
    c, s = np.cos(radians), np.sin(radians)
    return np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]])


def _grip_links(grip_pose: np.ndarray, half_gap: float = 0.01) -> dict[str, np.ndarray]:
    """Finger + base-link poses whose finger midpoint is exactly ``grip_pose``'s origin."""
    rot = grip_pose[:3, :3]
    pos = grip_pose[:3, 3]
    offset = rot @ np.array([0.0, half_gap, 0.0])
    left = np.eye(4)
    left[:3, :3] = rot
    left[:3, 3] = pos + offset
    right = np.eye(4)
    right[:3, :3] = rot
    right[:3, 3] = pos - offset
    return {
        "left_inner_finger": left,
        "right_inner_finger": right,
        "robotiq_85_base_link": grip_pose.copy(),
    }


class _ScriptedRobot:
    """Stub for RobotModel: returns one pre-scripted link_poses dict per call.

    solve() calls link_poses() exactly once per frame, in ascending frame order,
    so a plain call-counter index reproduces an exact scripted scenario without
    needing to interpret joint angles at all.
    """

    def __init__(self, frames: list[dict[str, np.ndarray]]) -> None:
        self._frames = frames
        self._idx = 0

    def link_poses(self, joint_positions: np.ndarray, gripper: float) -> dict[str, np.ndarray]:
        assert self._idx < len(self._frames), "solve() called link_poses more times than scripted"
        pose = self._frames[self._idx]
        self._idx += 1
        return pose


def _alignment(transform: np.ndarray) -> ObjectAlignment:
    return ObjectAlignment(
        transform=transform,
        scale=1.0,
        points_world=np.zeros((1, 3)),
    )


def _run(
    frames: list[dict[str, np.ndarray]],
    gripper: list[float],
    obj_transform: np.ndarray,
    cfg: InteractionConfig | None = None,
):
    n = len(frames)
    robot = _ScriptedRobot(frames)
    model = InteractionModel(robot, cfg or InteractionConfig())
    joint_positions = np.zeros((n, 7))
    timestamps = np.arange(n, dtype=np.float64) / 15.0
    return model.solve(_alignment(obj_transform), joint_positions, np.array(gripper), timestamps)


class TestGraspFollowsRotation:
    def test_object_rotates_and_orbits_with_gripper(self):
        grip0 = _pose([0.3, 0.0, 0.1])
        obj0 = _pose([0.32, 0.0, 0.1])  # 0.02 m offset from the grasp point, in x

        rot90 = _rotz(np.pi / 2)
        grip2_pos = rot90 @ np.array([0.3, 0.0, 0.1])
        grip2 = _pose(grip2_pos, rot90)

        frames = [_grip_links(grip0), _grip_links(grip0), _grip_links(grip2)]
        traj = _run(frames, gripper=[0.0, 0.3, 0.3], obj_transform=obj0)

        assert traj.states[0] is InteractionState.FREE
        assert traj.states[1] is InteractionState.GRASPED
        assert traj.states[2] is InteractionState.GRASPED

        # The object was rigidly attached with a fixed offset in the grip frame;
        # after a 90 degree rotation about world Z, that offset -- and hence the
        # object's position -- should itself have rotated by 90 degrees (orbited).
        expected_pos = grip2_pos + rot90 @ np.array([0.02, 0.0, 0.0])
        assert np.allclose(traj.transforms[2][:3, 3], expected_pos, atol=1e-9)
        # This is the user's explicit requirement: assert on the rotation matrix
        # itself, not just on position.
        assert np.allclose(traj.transforms[2][:3, :3], rot90, atol=1e-9)


class TestGraspFollowsLift:
    def test_object_rises_by_exactly_the_lift_distance(self):
        grip0 = _pose([0.3, 0.0, 0.1])
        obj0 = grip0.copy()  # object exactly at the grasp point

        grip1 = _pose([0.3, 0.0, 0.2])  # +0.1 m in z

        frames = [_grip_links(grip0), _grip_links(grip0), _grip_links(grip1)]
        traj = _run(frames, gripper=[0.0, 0.3, 0.3], obj_transform=obj0)

        assert traj.states[1] is InteractionState.GRASPED
        assert traj.states[2] is InteractionState.GRASPED
        delta = traj.transforms[2][:3, 3] - traj.transforms[1][:3, 3]
        assert np.allclose(delta, [0.0, 0.0, 0.1], atol=1e-9)


class TestPushGain:
    def test_object_moves_by_exactly_advance_times_gain(self):
        obj0 = _pose([0.5, 0.0, 0.1])
        grip0 = _pose([0.42, 0.0, 0.1])  # dist 0.08, outside contact_distance_m
        grip1 = _pose([0.47, 0.0, 0.1])  # dist 0.03, within contact; advance = 0.05

        frames = [_grip_links(grip0), _grip_links(grip1)]
        cfg = InteractionConfig(push_gain=4.0)
        traj = _run(frames, gripper=[0.0, 0.0], obj_transform=obj0, cfg=cfg)

        assert traj.states[0] is InteractionState.FREE
        assert traj.states[1] is InteractionState.PUSHED
        delta = traj.transforms[1][:3, 3] - traj.transforms[0][:3, 3]
        assert np.allclose(delta, [0.05 * 4.0, 0.0, 0.0], atol=1e-9)
        # Rotation must be untouched by a push.
        assert np.allclose(traj.transforms[1][:3, :3], np.eye(3), atol=1e-9)


class TestPushDoesNotFollowRetreat:
    def test_retreating_gripper_leaves_object_stationary(self):
        obj0 = _pose([0.5, 0.0, 0.1])
        grip0 = _pose([0.47, 0.0, 0.1])  # dist 0.03, already in contact
        grip1 = _pose([0.465, 0.0, 0.1])  # retreats by 0.005, dist 0.035, still in contact

        frames = [_grip_links(grip0), _grip_links(grip1)]
        traj = _run(frames, gripper=[0.0, 0.0], obj_transform=obj0)

        assert traj.states[1] is InteractionState.FREE
        assert np.allclose(traj.transforms[1], traj.transforms[0], atol=1e-12)


class TestHysteresis:
    def test_oscillating_gripper_does_not_toggle_state(self):
        grip = _pose([0.3, 0.0, 0.1])
        obj0 = grip.copy()

        frames = [_grip_links(grip) for _ in range(5)]
        # Crosses grasp_gripper_threshold (0.15) on the way up, then oscillates
        # between values above and below it but always above
        # release_gripper_threshold (0.10) -- a bare single-threshold state
        # machine would flicker free/grasped on frames 2 and 4.
        gripper = [0.05, 0.20, 0.12, 0.20, 0.12]
        traj = _run(frames, gripper=gripper, obj_transform=obj0)

        assert traj.states[0] is InteractionState.FREE
        assert traj.states[1:] == [InteractionState.GRASPED] * 4


class TestNoGraspWhenFar:
    def test_far_object_stays_free_and_stationary(self):
        obj0 = _pose([5.0, 0.0, 0.1])  # far from the gripper
        grip = _pose([0.3, 0.0, 0.1])

        frames = [_grip_links(grip), _grip_links(grip), _grip_links(grip)]
        traj = _run(frames, gripper=[0.0, 0.3, 0.3], obj_transform=obj0)

        assert all(s is InteractionState.FREE for s in traj.states)
        for t in range(len(traj)):
            assert np.allclose(traj.transforms[t], obj0, atol=1e-12)


class TestSummarize:
    def test_produces_one_row_per_state_span(self):
        obj0 = _pose([0.5, 0.0, 0.1])
        grip0 = _pose([0.42, 0.0, 0.1])
        grip1 = _pose([0.47, 0.0, 0.1])

        frames = [_grip_links(grip0), _grip_links(grip1)]
        traj = _run(frames, gripper=[0.0, 0.0], obj_transform=obj0)

        text = summarize(traj)
        assert "free" in text
        assert "pushed" in text
        assert len(text.splitlines()) == 1 + 2  # header + one row per span