"""
test_kalman_odom.py — Unit tests for KalmanOdomFilter and StereoVO.

Runs without a ROS2 environment (no rclpy imports).
"""

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_visual_odom.kalman_odom_filter import KalmanOdomFilter
from saltybot_visual_odom.stereo_vo import StereoVO, _rot_to_yaw


# ── KalmanOdomFilter tests ────────────────────────────────────────────────────

class TestKalmanOdomFilter:

    def test_initial_pose_zero(self):
        kf = KalmanOdomFilter()
        x, y, th = kf.pose
        assert x == pytest.approx(0.0)
        assert y == pytest.approx(0.0)
        assert th == pytest.approx(0.0)

    def test_predict_straight(self):
        kf = KalmanOdomFilter()
        # Set velocity to 1 m/s forward, then predict 1s
        kf.update_wheel(1.0, 0.0)
        kf.predict(1.0)
        x, y, th = kf.pose
        assert x == pytest.approx(1.0, abs=0.2)
        assert abs(y) < 0.1
        assert th == pytest.approx(0.0, abs=0.1)

    def test_predict_pure_rotation(self):
        kf = KalmanOdomFilter()
        kf.update_wheel(0.0, math.pi / 2)
        kf.predict(1.0)
        _, _, th = kf.pose
        assert th == pytest.approx(math.pi / 2, abs=0.1)

    def test_angle_wrap(self):
        kf = KalmanOdomFilter()
        kf.update_wheel(0.0, math.pi)
        kf.predict(1.1)   # should wrap past π
        _, _, th = kf.pose
        assert -math.pi <= th <= math.pi

    def test_wheel_slip_inflates_covariance(self):
        kf_normal = KalmanOdomFilter()
        kf_slip   = KalmanOdomFilter(slip_noise_multiplier=10.0)
        for _ in range(5):
            kf_normal.update_wheel(1.0, 0.0, slipping=False)
            kf_slip.update_wheel(1.0, 0.0, slipping=True)
            kf_normal.predict(0.1)
            kf_slip.predict(0.1)
        cov_normal = kf_normal.covariance_3x3[0, 0]
        cov_slip   = kf_slip.covariance_3x3[0, 0]
        # Slip covariance should be larger
        assert cov_slip > cov_normal

    def test_rtabmap_correction(self):
        kf = KalmanOdomFilter()
        # Drive straight for 5s to accumulate drift
        for _ in range(50):
            kf.update_wheel(1.0, 0.0)
            kf.predict(0.1)
        x_before, y_before, th_before = kf.pose
        # RTAB-Map says we're actually 1m behind
        kf.update_rtabmap(x_before - 1.0, 0.0, 0.0)
        x_after, _, _ = kf.pose
        # Pose should have moved toward correction
        assert x_after < x_before

    def test_visual_odom_fusion(self):
        kf = KalmanOdomFilter()
        kf.update_wheel(1.0, 0.0, slipping=False)
        kf.update_visual(0.95, 0.0, valid=True)   # VO slightly different
        kf.predict(0.1)
        v, _ = kf.velocity
        # Fused velocity should be between 0.95 and 1.0
        assert 0.9 <= v <= 1.05


# ── StereoVO tests ────────────────────────────────────────────────────────────

class TestStereoVO:

    def _make_K(self) -> np.ndarray:
        return np.array([
            [600.0,   0.0, 320.0],
            [  0.0, 600.0, 240.0],
            [  0.0,   0.0,   1.0],
        ])

    def test_no_K_returns_invalid(self):
        vo = StereoVO()
        prev = np.random.rand(10, 2).astype(np.float32) * 100
        curr = prev + np.random.rand(10, 2).astype(np.float32) * 2
        depth = np.ones((480, 640), dtype=np.float32) * 1.5
        est = vo.update(prev, curr, depth)
        assert not est.valid

    def test_static_scene_near_zero_motion(self):
        vo = StereoVO(K=self._make_K())
        # No real motion: prev_pts == curr_pts → zero translation
        pts = np.random.rand(80, 2).astype(np.float32)
        pts[:, 0] *= 640
        pts[:, 1] *= 480
        depth = np.ones((480, 640), dtype=np.float32) * 2.0
        est = vo.update(pts, pts.copy(), depth)
        # With identical points, E is degenerate → may be invalid
        # Just ensure it doesn't raise
        assert isinstance(est.valid, bool)

    def test_loop_closure_correction(self):
        vo = StereoVO(K=self._make_K())
        # Manually set accumulated pose
        vo._pose[0, 3] = 5.0
        vo._pose[1, 3] = 0.0
        # RTAB-Map says we're at (4.5, 0)
        vo.correct_from_loop_closure(4.5, 0.0, 0.0, alpha=1.0)
        assert vo._pose[0, 3] == pytest.approx(4.5, abs=0.01)

    def test_reset_clears_state(self):
        vo = StereoVO(K=self._make_K())
        vo._pose[0, 3] = 10.0
        vo.reset()
        est = vo._invalid_est()
        assert est.x == pytest.approx(0.0)


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