"""
Unit tests for saltybot_uwb_imu_fusion.ekf_math (Issue #573).
No ROS2 or hardware required.
"""

import math
import sys
import os

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from saltybot_uwb_imu_fusion.ekf_math import UwbImuEKF, _wrap_angle


class TestWrapAngle:
    def test_within_range(self):
        assert abs(_wrap_angle(1.0) - 1.0) < 1e-9

    def test_positive_overflow(self):
        assert abs(_wrap_angle(math.pi + 0.1) - (-math.pi + 0.1)) < 1e-9

    def test_negative_overflow(self):
        assert abs(_wrap_angle(-math.pi - 0.1) - (math.pi - 0.1)) < 1e-9


class TestEkfInitialise:
    def test_not_initialised_by_default(self):
        ekf = UwbImuEKF()
        assert not ekf.is_initialised

    def test_initialise_sets_position(self):
        ekf = UwbImuEKF()
        ekf.initialise(3.0, 4.0, heading_rad=0.5)
        assert ekf.is_initialised
        x, y = ekf.position
        assert abs(x - 3.0) < 1e-9
        assert abs(y - 4.0) < 1e-9
        assert abs(ekf.heading - 0.5) < 1e-9


class TestEkfPredict:
    def test_stationary_predict_no_drift(self):
        """Stationary robot with zero IMU input stays near origin."""
        ekf = UwbImuEKF(sigma_vel_drift=0.0)
        ekf.initialise(0.0, 0.0)
        for _ in range(10):
            ekf.predict(ax_body=0.0, ay_body=0.0, omega=0.0, dt=0.01)
        x, y = ekf.position
        assert abs(x) < 0.01
        assert abs(y) < 0.01

    def test_forward_acceleration(self):
        """1 m/s² forward for 0.5 s → ~0.125 m forward displacement."""
        ekf = UwbImuEKF(sigma_imu_accel=0.0, sigma_vel_drift=0.0)
        ekf.initialise(0.0, 0.0, heading_rad=0.0)
        dt = 0.005
        for _ in range(100):  # 0.5 s
            ekf.predict(ax_body=1.0, ay_body=0.0, omega=0.0, dt=dt)
        x, y = ekf.position
        # x ≈ 0.5 * 1 * 0.5² = 0.125 m
        assert 0.10 < x < 0.16, f"x={x}"
        assert abs(y) < 0.01

    def test_yaw_rate(self):
        """π/4 rad/s for 1 s → heading ≈ π/4."""
        ekf = UwbImuEKF(sigma_imu_gyro=0.0)
        ekf.initialise(0.0, 0.0, heading_rad=0.0)
        dt = 0.005
        for _ in range(200):
            ekf.predict(ax_body=0.0, ay_body=0.0, omega=math.pi / 4, dt=dt)
        assert abs(ekf.heading - math.pi / 4) < 0.05

    def test_covariance_grows_with_time(self):
        """Covariance must grow during predict (no updates)."""
        ekf = UwbImuEKF()
        ekf.initialise(0.0, 0.0)
        p0 = float(ekf.covariance_position[0, 0])
        for _ in range(50):
            ekf.predict(ax_body=0.0, ay_body=0.0, omega=0.0, dt=0.01)
        p1 = float(ekf.covariance_position[0, 0])
        assert p1 > p0, f"covariance did not grow: p0={p0}, p1={p1}"


class TestEkfUpdate:
    def test_uwb_update_reduces_covariance(self):
        """UWB update must reduce position covariance."""
        ekf = UwbImuEKF()
        ekf.initialise(0.0, 0.0)
        # Grow uncertainty first
        for _ in range(20):
            ekf.predict(ax_body=0.0, ay_body=0.0, omega=0.0, dt=0.05)
        p_before = float(ekf.covariance_position[0, 0])
        ekf.update_uwb(0.1, 0.1)
        p_after = float(ekf.covariance_position[0, 0])
        assert p_after < p_before

    def test_uwb_corrects_position(self):
        """Large position error corrected by UWB measurement."""
        ekf = UwbImuEKF(sigma_uwb_pos_m=0.1)
        ekf.initialise(5.0, 5.0)
        # Multiple UWB updates at true position (0,0)
        for _ in range(20):
            ekf.update_uwb(0.0, 0.0)
        x, y = ekf.position
        assert abs(x) < 0.5, f"x not corrected: {x}"
        assert abs(y) < 0.5, f"y not corrected: {y}"

    def test_uninitialised_update_initialises(self):
        """Calling update_uwb before initialise should initialise filter."""
        ekf = UwbImuEKF()
        assert not ekf.is_initialised
        ekf.update_uwb(2.0, 3.0)
        assert ekf.is_initialised
        x, y = ekf.position
        assert abs(x - 2.0) < 1e-9
        assert abs(y - 3.0) < 1e-9


class TestDropout:
    def test_velocity_damping_during_dropout(self):
        """Velocity decays during extended UWB dropout."""
        ekf = UwbImuEKF(dropout_vel_damping=0.9)
        ekf.initialise(0.0, 0.0)
        # Give it some velocity
        for _ in range(50):
            ekf.predict(ax_body=2.0, ay_body=0.0, omega=0.0, dt=0.01)
        vx_before, _ = ekf.velocity
        # Let dropout accumulate (> 2 s threshold)
        for _ in range(300):
            ekf.predict(ax_body=0.0, ay_body=0.0, omega=0.0, dt=0.01)
        vx_after, _ = ekf.velocity
        assert abs(vx_after) < abs(vx_before), \
            f"velocity not damped: before={vx_before:.3f}, after={vx_after:.3f}"


class TestCovariance:
    def test_covariance_positive_definite(self):
        """Full covariance matrix must be positive definite at all times."""
        ekf = UwbImuEKF()
        ekf.initialise(1.0, 2.0)
        for _ in range(20):
            ekf.predict(ax_body=0.5, ay_body=0.1, omega=0.05, dt=0.01)
            if _ % 5 == 0:
                ekf.update_uwb(1.0, 2.0)
        eigvals = np.linalg.eigvalsh(ekf.covariance_full)
        assert all(ev > -1e-9 for ev in eigvals), f"Non-PD covariance: {eigvals}"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
