"""Differential-drive and mecanum kinematics tests."""

import math

import pytest

from herding.world.diffdrive import (
    heading_speed_to_wheels, kinematics_step,
    mecanum_inverse, mecanum_kinematics_step,
    velocity_to_mecanum_wheels, velocity_to_wheels,
)


WHEEL_R = 0.038
WHEEL_B = 0.28
MAX_OMEGA = 70.0
MAX_LINEAR = WHEEL_R * MAX_OMEGA
DT = 0.016


def test_kinematics_zero_input_is_identity():
    x, y, h = kinematics_step(1.0, 2.0, 0.5, 0.0, 0.0, WHEEL_R, WHEEL_B, DT)
    assert (x, y, h) == (1.0, 2.0, 0.5)


def test_kinematics_forward_motion():
    # Equal wheel speeds → pure translation along the heading.
    x, y, h = kinematics_step(0.0, 0.0, 0.0, 10.0, 10.0, WHEEL_R, WHEEL_B, DT)
    assert h == 0.0
    assert math.isclose(x, 10.0 * WHEEL_R * DT)
    assert y == 0.0


def test_kinematics_pure_rotation():
    # Opposite wheel speeds → pure rotation, position unchanged.
    x, y, h = kinematics_step(0.0, 0.0, 0.0, -5.0, 5.0, WHEEL_R, WHEEL_B, DT)
    assert (x, y) == (0.0, 0.0)
    assert h > 0.0


def test_kinematics_heading_wrapped_to_pi():
    _, _, h = kinematics_step(0.0, 0.0, math.pi - 0.01, 100.0, -100.0,
                              WHEEL_R, WHEEL_B, DT)
    assert -math.pi <= h <= math.pi


def test_velocity_to_wheels_zero_velocity():
    left, right = velocity_to_wheels(0.0, 0.0, 0.0,
                                     MAX_LINEAR, WHEEL_R, MAX_OMEGA)
    assert (left, right) == (0.0, 0.0)


def test_velocity_to_wheels_aligned_forward():
    # Target straight ahead → equal positive wheel speeds.
    left, right = velocity_to_wheels(1.0, 0.0, 0.0,
                                     MAX_LINEAR, WHEEL_R, MAX_OMEGA, k_turn=4.0)
    assert math.isclose(left, right, abs_tol=1e-6)
    assert left > 0.0


def test_velocity_to_wheels_perpendicular_target_spins():
    # Target 90° from heading → forward speed ≈ 0, wheels equal-and-opposite.
    left, right = velocity_to_wheels(0.0, 1.0, 0.0,
                                     MAX_LINEAR, WHEEL_R, MAX_OMEGA, k_turn=4.0)
    assert left + right == pytest.approx(0.0, abs=1e-6)
    assert right > 0.0  # turning CCW (left of heading is +y for h=0)


def test_velocity_to_wheels_clamped_to_max_omega():
    # Far overshoot — both wheel commands clamped at ±MAX_OMEGA.
    left, right = velocity_to_wheels(-1.0, 0.0, 0.0,
                                     MAX_LINEAR, WHEEL_R, MAX_OMEGA, k_turn=20.0)
    assert -MAX_OMEGA <= left <= MAX_OMEGA
    assert -MAX_OMEGA <= right <= MAX_OMEGA


def test_heading_speed_to_wheels_aligned():
    left, right = heading_speed_to_wheels(0.0, 10.0, 0.0, MAX_OMEGA)
    assert math.isclose(left, right, abs_tol=1e-6)
    assert left > 0.0


def test_heading_speed_to_wheels_reverse_target_forwards_zero():
    left, right = heading_speed_to_wheels(math.pi, 10.0, 0.0, MAX_OMEGA)
    # cos(π) clamped at 0 → no forward; pure rotation.
    assert left + right == pytest.approx(0.0, abs=1e-6)


# ---------------------------------------------------------------------------
# Mecanum kinematics tests
# ---------------------------------------------------------------------------

LX = 0.14   # half wheel_base_x
LY = 0.14   # half wheel_base_y


def test_mecanum_kinematics_zero_is_identity():
    x, y, h = mecanum_kinematics_step(
        1.0, 2.0, 0.5, 0.0, 0.0, 0.0, 0.0, WHEEL_R, LX, LY, DT,
    )
    assert (x, y, h) == (1.0, 2.0, 0.5)


def test_mecanum_kinematics_pure_forward():
    # All 4 wheels equal → pure forward (vx_body > 0, vy_body = 0).
    w = 10.0
    x, y, h = mecanum_kinematics_step(
        0.0, 0.0, 0.0, w, w, w, w, WHEEL_R, LX, LY, DT,
    )
    assert h == pytest.approx(0.0, abs=1e-9)
    assert y == pytest.approx(0.0, abs=1e-9)
    assert math.isclose(x, w * WHEEL_R * DT, rel_tol=1e-6)


def test_mecanum_kinematics_pure_strafe():
    # Strafe right (positive vy_body) with zero forward:
    # vx_body = (w_fl+w_fr+w_rl+w_rr)*r/4 = 0 → sum of wheels = 0
    # vy_body = (-w_fl+w_fr+w_rl-w_rr)*r/4 > 0
    # Use w_fl=-10, w_fr=10, w_rl=10, w_rr=-10.
    w_fl, w_fr, w_rl, w_rr = -10.0, 10.0, 10.0, -10.0
    x, y, h = mecanum_kinematics_step(
        0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT,
    )
    assert h == pytest.approx(0.0, abs=1e-9)
    assert x == pytest.approx(0.0, abs=1e-9)
    expected_vy = (-w_fl + w_fr + w_rl - w_rr) * WHEEL_R / 4.0
    assert math.isclose(y, expected_vy * DT, rel_tol=1e-6)


def test_mecanum_kinematics_pure_rotation():
    # Pure rotation: vx_body=0, vy_body=0, omega>0.
    # w_fl=-10, w_fr=10, w_rl=-10, w_rr=10 → all sums cancel except omega.
    w_fl, w_fr, w_rl, w_rr = -10.0, 10.0, -10.0, 10.0
    x, y, h = mecanum_kinematics_step(
        0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT,
    )
    assert x == pytest.approx(0.0, abs=1e-9)
    assert y == pytest.approx(0.0, abs=1e-9)
    assert h > 0.0


def test_mecanum_inverse_roundtrip():
    # Inverse → forward: pick desired body velocities, compute wheels,
    # then verify forward kinematics recovers the same velocities.
    vx_b = 0.5
    vy_b = 0.3
    omega = 0.2
    w_fl, w_fr, w_rl, w_rr = mecanum_inverse(
        vx_b, vy_b, omega, WHEEL_R, LX, LY, MAX_OMEGA,
    )
    vx_check = (w_fl + w_fr + w_rl + w_rr) * WHEEL_R / 4.0
    vy_check = (-w_fl + w_fr + w_rl - w_rr) * WHEEL_R / 4.0
    omega_check = (-w_fl + w_fr - w_rl + w_rr) * WHEEL_R / (4.0 * (LX + LY))
    assert math.isclose(vx_b, vx_check, rel_tol=1e-6)
    assert math.isclose(vy_b, vy_check, rel_tol=1e-6)
    assert math.isclose(omega, omega_check, rel_tol=1e-6)


def test_mecanum_inverse_clamped():
    # Request an extreme velocity — all wheels should be clamped.
    w_fl, w_fr, w_rl, w_rr = mecanum_inverse(
        100.0, 100.0, 50.0, WHEEL_R, LX, LY, MAX_OMEGA,
    )
    assert max(abs(w_fl), abs(w_fr), abs(w_rl), abs(w_rr)) <= MAX_OMEGA


def test_velocity_to_mecanum_wheels_zero():
    result = velocity_to_mecanum_wheels(
        0.0, 0.0, 0.0, 0.0, MAX_LINEAR, WHEEL_R, LX, LY, MAX_OMEGA,
        wheel_base=WHEEL_B,
    )
    assert result == (0.0, 0.0, 0.0, 0.0)


def test_velocity_to_mecanum_wheels_forward():
    w_fl, w_fr, w_rl, w_rr = velocity_to_mecanum_wheels(
        1.0, 0.0, 0.0, 0.0, MAX_LINEAR, WHEEL_R, LX, LY, MAX_OMEGA,
        wheel_base=WHEEL_B,
    )
    # All 4 wheels should be positive and roughly equal.
    assert all(w > 0.0 for w in (w_fl, w_fr, w_rl, w_rr))
    assert math.isclose(w_fl, w_rr, rel_tol=1e-6)
    assert math.isclose(w_fr, w_rl, rel_tol=1e-6)


def test_velocity_to_mecanum_wheels_clamped():
    # Extreme input — all wheels within max.
    ws = velocity_to_mecanum_wheels(
        1.0, 1.0, 1.0, 0.0, MAX_LINEAR, WHEEL_R, LX, LY, MAX_OMEGA,
        wheel_base=WHEEL_B,
    )
    assert all(abs(w) <= MAX_OMEGA for w in ws)