import numpy as np
import pinocchio as pin
import pytest

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


def test_joint_limit_mid_regularization_points_toward_range_center():
    kinematics = RM75Kinematics()
    solver = RM75IkSolver(kinematics)
    limits = kinematics.limits
    midpoint = 0.5 * (limits.lower + limits.upper)
    joint_range = limits.upper - limits.lower
    options = IkOptions(
        posture_weight=0.0,
        joint_limit_mid_weight=2e-5,
        joint_motion_weights=(1.0, 1.0, 1.0, 1.0, 0.3, 0.3, 0.2),
    )

    center_hessian, center_gradient = solver._regularization_terms(
        midpoint, midpoint, options, damping=0.1
    )
    np.testing.assert_allclose(center_gradient, 0.0, atol=1e-15)

    above = midpoint + 0.4 * joint_range
    below = midpoint - 0.4 * joint_range
    _, above_gradient = solver._regularization_terms(
        above, above, options, damping=0.1
    )
    _, below_gradient = solver._regularization_terms(
        below, below, options, damping=0.1
    )
    assert np.all(above_gradient > 0.0)
    np.testing.assert_allclose(below_gradient, -above_gradient, rtol=1e-12)
    assert center_hessian[6, 6] < center_hessian[0, 0]


def test_per_joint_step_limits_and_hard_position_limits_are_combined():
    kinematics = RM75Kinematics()
    solver = RM75IkSolver(kinematics)
    limits = kinematics.limits
    midpoint = 0.5 * (limits.lower + limits.upper)
    configured = np.array([0.05, 0.05, 0.05, 0.05, 0.08, 0.08, 0.10])
    options = IkOptions(joint_step_limits_rad=tuple(configured))

    lower, upper = solver._step_bounds(midpoint, options)
    np.testing.assert_allclose(lower, -configured)
    np.testing.assert_allclose(upper, configured)

    near_upper = midpoint.copy()
    near_upper[6] = limits.upper[6] - 0.01
    lower, upper = solver._step_bounds(near_upper, options)
    assert upper[6] == pytest.approx(0.01)
    assert lower[6] == pytest.approx(-0.10)


@pytest.mark.parametrize(
    "kwargs",
    [
        {"joint_limit_mid_weight": -1.0},
        {"joint_motion_weights": (1.0,) * 6 + (0.0,)},
        {"joint_step_limits_rad": (0.05,) * 6},
        {"joint_step_limits_rad": (0.05,) * 6 + (float("nan"),)},
    ],
)
def test_invalid_joint_regularization_options_are_rejected(kwargs):
    with pytest.raises(ValueError):
        IkOptions(**kwargs)