import numpy as np
import pinocchio as pin

from rm75_ik import IkOptions, IkStatus, RM75IkSolver, RM75Kinematics, pose_errors


def test_near_seed_round_trip_converges():
    kinematics = RM75Kinematics()
    solver = RM75IkSolver(kinematics)
    target_q = np.deg2rad([30, -20, 40, 60, -50, 25, 90])
    seed = target_q + np.deg2rad([2, -1, 2, -1, 1, -2, 3])

    result = solver.solve(kinematics.forward(target_q), seed)

    assert result.status is IkStatus.SUCCESS
    assert result.q is not None
    position_error, orientation_error = pose_errors(
        kinematics.forward(result.q), kinematics.forward(target_q)
    )
    assert position_error <= 1e-3
    assert orientation_error <= np.deg2rad(0.1)


def test_invalid_seed_returns_no_solution():
    kinematics = RM75Kinematics()
    result = RM75IkSolver(kinematics).solve(
        kinematics.forward(np.zeros(7)), np.full(7, np.nan)
    )

    assert result.status is IkStatus.INVALID_INPUT
    assert result.q is None


def test_invalid_rotation_returns_no_solution():
    kinematics = RM75Kinematics()
    invalid_target = pin.SE3(2.0 * np.eye(3), np.zeros(3))
    result = RM75IkSolver(kinematics).solve(invalid_target, np.zeros(7))

    assert result.status is IkStatus.INVALID_INPUT
    assert result.q is None


def test_expired_time_budget_returns_no_solution():
    kinematics = RM75Kinematics()
    result = RM75IkSolver(kinematics).solve(
        kinematics.forward(np.deg2rad([30, 20, -40, 60, 50, -25, 90])),
        np.zeros(7),
        IkOptions(time_limit_sec=1e-12),
    )

    assert result.status is IkStatus.TIME_LIMIT
    assert result.q is None