AR-Autopilot/arautopilot/tests/test_heading_ekf.py

"""Tests for 2-state Heading EKF -- Sprint 8."""

from __future__ import annotations

import math

import pytest

from arautopilot.core.heading_ekf import HeadingEKF, _shortest_arc


class TestShortestArc:
    def test_zero_difference(self):
        assert _shortest_arc(90.0, 90.0) == pytest.approx(0.0)

    def test_positive_difference(self):
        assert _shortest_arc(100.0, 90.0) == pytest.approx(10.0)

    def test_negative_difference(self):
        assert _shortest_arc(80.0, 90.0) == pytest.approx(-10.0)

    def test_wrap_around_positive(self):
        # 350 → 10 is +20 degrees (going through north)
        assert _shortest_arc(10.0, 350.0) == pytest.approx(20.0)

    def test_wrap_around_negative(self):
        # 10 → 350 is -20 degrees
        assert _shortest_arc(350.0, 10.0) == pytest.approx(-20.0)

    def test_exactly_180_degrees(self):
        diff = _shortest_arc(270.0, 90.0)
        assert abs(diff) == pytest.approx(180.0)


class TestPredict:
    def test_heading_advances_by_rot_times_dt(self):
        ekf = HeadingEKF(heading_deg=0.0, rot_dps=10.0)
        ekf.predict(dt_s=1.0)
        assert ekf.heading_deg == pytest.approx(10.0)

    def test_heading_wraps_at_360(self):
        ekf = HeadingEKF(heading_deg=355.0, rot_dps=10.0)
        ekf.predict(dt_s=1.0)
        assert ekf.heading_deg == pytest.approx(5.0)

    def test_zero_rot_heading_unchanged(self):
        ekf = HeadingEKF(heading_deg=45.0, rot_dps=0.0)
        ekf.predict(dt_s=1.0)
        assert ekf.heading_deg == pytest.approx(45.0)

    def test_covariance_grows_with_predict(self):
        ekf = HeadingEKF()
        p00_before = ekf._P[0]
        ekf.predict(dt_s=0.1)
        assert ekf._P[0] > p00_before

    def test_predict_symmetry_p01_p10(self):
        ekf = HeadingEKF()
        for _ in range(5):
            ekf.predict(dt_s=0.1)
        assert ekf._P[1] == pytest.approx(ekf._P[2])


class TestUpdateHeading:
    def test_state_moves_toward_measurement(self):
        ekf = HeadingEKF(heading_deg=0.0)
        ekf.update_heading(10.0)
        assert 0.0 < ekf.heading_deg < 10.0

    def test_exact_measurement_moves_fully_when_p_large(self):
        # With very large P and very small noise, Kalman gain → 1
        ekf = HeadingEKF(heading_deg=0.0, _P=[1e6, 0.0, 0.0, 1e6])
        ekf.update_heading(90.0, noise_deg=0.001)
        assert ekf.heading_deg == pytest.approx(90.0, abs=0.1)

    def test_covariance_shrinks_after_update(self):
        ekf = HeadingEKF()
        ekf._P = [100.0, 0.0, 0.0, 100.0]
        p00_before = ekf._P[0]
        ekf.update_heading(10.0, noise_deg=2.0)
        assert ekf._P[0] < p00_before

    def test_wrap_around_innov(self):
        ekf = HeadingEKF(heading_deg=355.0)
        ekf._P = [1e6, 0.0, 0.0, 1e6]
        ekf.update_heading(5.0, noise_deg=0.001)
        # Should go toward 5.0 (via shortest arc +10), not regress
        assert ekf.heading_deg == pytest.approx(5.0, abs=0.5)


class TestUpdateRot:
    def test_state_moves_toward_rot_measurement(self):
        ekf = HeadingEKF(rot_dps=0.0)
        ekf.update_rot(5.0)
        assert 0.0 < ekf.rot_dps < 5.0

    def test_covariance_p11_shrinks_after_rot_update(self):
        ekf = HeadingEKF()
        ekf._P = [100.0, 0.0, 0.0, 100.0]
        p11_before = ekf._P[3]
        ekf.update_rot(2.0, noise_dps=1.0)
        assert ekf._P[3] < p11_before


class TestCovarianceConvergence:
    def test_covariance_converges_with_repeated_updates(self):
        ekf = HeadingEKF(
            process_noise_heading=0.01,
            process_noise_rot=0.1,
            _P=[100.0, 0.0, 0.0, 100.0],
        )
        # 50 predict+update cycles
        for _ in range(50):
            ekf.predict(dt_s=0.1)
            ekf.update_heading(ekf.heading_deg, noise_deg=2.0)
            ekf.update_rot(ekf.rot_dps, noise_dps=1.0)

        # Covariance should have converged to steady-state (not 100 anymore)
        assert ekf._P[0] < 50.0
        assert ekf._P[3] < 50.0

    def test_filter_tracks_constant_heading(self):
        true_heading = 135.0
        ekf = HeadingEKF(heading_deg=0.0)
        for _ in range(100):
            ekf.predict(dt_s=0.1)
            ekf.update_heading(true_heading, noise_deg=2.0)
        assert ekf.heading_deg == pytest.approx(true_heading, abs=2.0)

    def test_filter_tracks_constant_rot(self):
        ekf = HeadingEKF(heading_deg=0.0, rot_dps=0.0)
        true_rot = 3.0
        for _ in range(100):
            ekf.predict(dt_s=0.1)
            ekf.update_rot(true_rot, noise_dps=0.5)
        assert ekf.rot_dps == pytest.approx(true_rot, abs=0.5)

    def test_covariance_property_returns_tuple(self):
        ekf = HeadingEKF()
        cov = ekf.covariance
        assert len(cov) == 4
        assert all(isinstance(v, float) for v in cov)