"""Unit tests for PoseFusionEKF (no ROS2 required)."""

import math
import pytest
from saltybot_pose_fusion.pose_fusion_ekf import PoseFusionEKF


def make_ekf(**kwargs):
    return PoseFusionEKF(**kwargs)


# ── Initialisation ─────────────────────────────────────────────────────────────

def test_not_initialised_before_uwb():
    ekf = make_ekf()
    assert not ekf.is_initialised


def test_initialised_after_uwb_update():
    ekf = make_ekf()
    ekf.update_uwb_position(1.0, 2.0)
    assert ekf.is_initialised
    x, y = ekf.position
    assert abs(x - 1.0) < 1e-9
    assert abs(y - 2.0) < 1e-9


# ── IMU predict ───────────────────────────────────────────────────────────────

def test_predict_advances_position():
    ekf = make_ekf()
    ekf.update_uwb_position(0.0, 0.0)
    # Give it a forward velocity by injecting a forward acceleration
    ekf.predict_imu(ax_body=1.0, ay_body=0.0, omega=0.0, dt=1.0)
    ekf.predict_imu(ax_body=1.0, ay_body=0.0, omega=0.0, dt=1.0)
    x, y = ekf.position
    assert x > 0.0, f"Expected forward motion, got x={x}"


def test_predict_no_crash_before_init():
    ekf = make_ekf()
    ekf.predict_imu(ax_body=1.0, ay_body=0.0, omega=0.0, dt=0.01)
    assert not ekf.is_initialised


# ── UWB position update ───────────────────────────────────────────────────────

def test_uwb_update_converges():
    ekf = make_ekf(sigma_uwb_pos_m=0.10)
    ekf.update_uwb_position(5.0, 3.0)
    for _ in range(20):
        ekf.predict_imu(ax_body=0.0, ay_body=0.0, omega=0.0, dt=0.05)
        ekf.update_uwb_position(5.0, 3.0)
    x, y = ekf.position
    assert abs(x - 5.0) < 0.5
    assert abs(y - 3.0) < 0.5


def test_uwb_resets_dropout():
    ekf = make_ekf()
    ekf.update_uwb_position(0.0, 0.0)
    ekf.predict_imu(0.0, 0.0, 0.0, dt=5.0)
    assert ekf.uwb_dropout_s > 4.0
    ekf.update_uwb_position(0.0, 0.0)
    assert ekf.uwb_dropout_s == 0.0


# ── UWB heading update ────────────────────────────────────────────────────────

def test_uwb_heading_update():
    ekf = make_ekf(sigma_uwb_head_rad=0.1)
    ekf.update_uwb_position(0.0, 0.0)
    target_yaw = math.radians(45.0)
    for _ in range(10):
        ekf.update_uwb_heading(target_yaw)
    assert abs(ekf.heading - target_yaw) < 0.2


# ── VO velocity update ────────────────────────────────────────────────────────

def test_vo_velocity_update():
    ekf = make_ekf(sigma_vo_vel_m_s=0.05, sigma_vo_omega_r_s=0.03)
    ekf.update_uwb_position(0.0, 0.0)
    ekf.update_vo_velocity(vx_body=1.0, omega=0.0)
    vx, vy = ekf.velocity
    assert vx > 0.5, f"Expected vx>0.5 after VO update, got {vx}"


def test_vo_resets_dropout():
    ekf = make_ekf()
    ekf.update_uwb_position(0.0, 0.0)
    ekf.predict_imu(0.0, 0.0, 0.0, dt=3.0)
    assert ekf.vo_dropout_s > 2.0
    ekf.update_vo_velocity(vx_body=0.0, omega=0.0)
    assert ekf.vo_dropout_s == 0.0


# ── Covariance ────────────────────────────────────────────────────────────────

def test_covariance_decreases_with_updates():
    ekf = make_ekf()
    ekf.update_uwb_position(0.0, 0.0)
    p_init = ekf.position_uncertainty_m()
    for _ in range(30):
        ekf.predict_imu(0.0, 0.0, 0.0, dt=0.05)
        ekf.update_uwb_position(0.0, 0.0)
    assert ekf.position_uncertainty_m() < p_init


def test_covariance_matrix_shape():
    ekf = make_ekf()
    ekf.update_uwb_position(1.0, 1.0)
    assert ekf.covariance_position.shape == (2, 2)
    assert ekf.covariance_full.shape == (6, 6)


# ── Sigma override ────────────────────────────────────────────────────────────

def test_custom_sigma_override():
    """High sigma should produce weaker updates than low sigma."""
    ekf_weak   = make_ekf(sigma_uwb_pos_m=2.0)
    ekf_strong = make_ekf(sigma_uwb_pos_m=0.01)

    for ekf in (ekf_weak, ekf_strong):
        ekf.update_uwb_position(10.0, 0.0)
        ekf.predict_imu(0.0, 0.0, 0.0, dt=0.1)
        ekf.update_uwb_position(0.0, 0.0)

    x_weak,   _ = ekf_weak.position
    x_strong, _ = ekf_strong.position

    # Strong measurement should pull closer to 0
    assert x_strong < x_weak, \
        f"Strong sigma should pull x closer to 0: strong={x_strong:.3f} weak={x_weak:.3f}"
