"""
test_tracking_fusion.py — Unit tests for saltybot_social_tracking pure modules.

Tests cover:
  - KalmanTracker: initialization, predict, update, state accessors
  - source_arbiter: confidence functions, source selection, fusion, bearing

No ROS2 runtime required.
"""

import math
import sys
import os

# Allow running: python -m pytest test/test_tracking_fusion.py
# from the package root without installing the package.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import pytest
import numpy as np

from saltybot_social_tracking.kalman_tracker import KalmanTracker
from saltybot_social_tracking.source_arbiter import (
    uwb_confidence,
    camera_confidence,
    select_source,
    fuse_positions,
    composite_confidence,
    bearing_and_range,
)


# ─────────────────────────────────────────────────────────────────────────────
# KalmanTracker tests
# ─────────────────────────────────────────────────────────────────────────────

class TestKalmanTrackerInit:

    def test_not_initialized_by_default(self):
        kf = KalmanTracker()
        assert not kf.initialized

    def test_initialize_sets_position(self):
        kf = KalmanTracker()
        kf.initialize(3.0, 1.5)
        assert kf.initialized
        x, y = kf.position
        assert abs(x - 3.0) < 1e-9
        assert abs(y - 1.5) < 1e-9

    def test_initialize_sets_zero_velocity(self):
        kf = KalmanTracker()
        kf.initialize(1.0, -2.0)
        vx, vy = kf.velocity
        assert abs(vx) < 1e-9
        assert abs(vy) < 1e-9

    def test_initialize_origin(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        assert kf.initialized
        x, y = kf.position
        assert x == 0.0 and y == 0.0


class TestKalmanTrackerPredict:

    def test_predict_zero_dt_no_change(self):
        kf = KalmanTracker()
        kf.initialize(2.0, 1.0)
        kf.predict(0.0)
        x, y = kf.position
        assert abs(x - 2.0) < 1e-9
        assert abs(y - 1.0) < 1e-9

    def test_predict_negative_dt_no_change(self):
        kf = KalmanTracker()
        kf.initialize(2.0, 1.0)
        kf.predict(-0.1)
        x, y = kf.position
        assert abs(x - 2.0) < 1e-9

    def test_predict_constant_velocity(self):
        """After a position update gives the filter a velocity, predict should extrapolate."""
        kf = KalmanTracker(process_noise=0.001, meas_noise_uwb=0.001)
        kf.initialize(0.0, 0.0)
        # Force filter to track a moving target to build up velocity estimate
        dt = 0.05
        for i in range(40):
            t = i * dt
            kf.predict(dt)
            kf.update(2.0 * t, 0.0, "uwb")   # 2 m/s in x

        # After many updates the velocity estimate should be close to 2 m/s
        vx, vy = kf.velocity
        assert abs(vx - 2.0) < 0.3, f"vx={vx:.3f}"
        assert abs(vy) < 0.2

    def test_predict_grows_uncertainty(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        unc_before = kf.position_uncertainty_m()
        kf.predict(1.0)
        unc_after = kf.position_uncertainty_m()
        assert unc_after > unc_before

    def test_predict_multiple_steps(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        kf.predict(0.1)
        kf.predict(0.1)
        kf.predict(0.1)
        # No assertion on exact value; just verify no exception and state is finite
        x, y = kf.position
        assert math.isfinite(x) and math.isfinite(y)


class TestKalmanTrackerUpdate:

    def test_update_pulls_position_toward_measurement(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        kf.update(5.0, 5.0, "uwb")
        x, y = kf.position
        assert x > 0.0 and y > 0.0
        assert x < 5.0 and y < 5.0   # blended, not jumped

    def test_update_reduces_uncertainty(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        kf.predict(1.0)              # uncertainty grows
        unc_mid = kf.position_uncertainty_m()
        kf.update(0.1, 0.1, "uwb")  # measurement corrects
        unc_after = kf.position_uncertainty_m()
        assert unc_after < unc_mid

    def test_update_converges_to_true_position(self):
        """Many updates from same point should converge to that point."""
        kf = KalmanTracker(meas_noise_uwb=0.01)
        kf.initialize(0.0, 0.0)
        for _ in range(50):
            kf.update(3.0, -1.0, "uwb")
        x, y = kf.position
        assert abs(x - 3.0) < 0.05, f"x={x:.4f}"
        assert abs(y - (-1.0)) < 0.05, f"y={y:.4f}"

    def test_update_camera_source_different_noise(self):
        """Camera and UWB updates should both move state (noise model differs)."""
        kf1 = KalmanTracker(meas_noise_uwb=0.20, meas_noise_cam=0.10)
        kf1.initialize(0.0, 0.0)
        kf1.update(5.0, 0.0, "uwb")
        x_uwb, _ = kf1.position

        kf2 = KalmanTracker(meas_noise_uwb=0.20, meas_noise_cam=0.10)
        kf2.initialize(0.0, 0.0)
        kf2.update(5.0, 0.0, "camera")
        x_cam, _ = kf2.position

        # Camera has lower noise → stronger pull toward measurement
        assert x_cam > x_uwb

    def test_update_unknown_source_defaults_to_camera_noise(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        kf.update(2.0, 0.0, "other")   # unknown source — should not raise
        x, _ = kf.position
        assert x > 0.0

    def test_position_uncertainty_finite(self):
        kf = KalmanTracker()
        kf.initialize(1.0, 1.0)
        kf.predict(0.05)
        kf.update(1.1, 0.9, "uwb")
        assert math.isfinite(kf.position_uncertainty_m())
        assert kf.position_uncertainty_m() >= 0.0

    def test_covariance_copy_is_independent(self):
        kf = KalmanTracker()
        kf.initialize(0.0, 0.0)
        cov = kf.covariance_copy()
        cov[0, 0] = 9999.0   # mutate copy
        assert kf.covariance_copy()[0, 0] != 9999.0


# ─────────────────────────────────────────────────────────────────────────────
# source_arbiter tests
# ─────────────────────────────────────────────────────────────────────────────

class TestUwbConfidence:

    def test_zero_age_returns_quality(self):
        assert abs(uwb_confidence(0.0, 1.5) - 1.0) < 1e-9

    def test_at_timeout_returns_zero(self):
        assert uwb_confidence(1.5, 1.5) == pytest.approx(0.0)

    def test_beyond_timeout_returns_zero(self):
        assert uwb_confidence(2.0, 1.5) == 0.0

    def test_half_timeout_returns_half(self):
        assert uwb_confidence(0.75, 1.5) == pytest.approx(0.5)

    def test_quality_scales_result(self):
        assert uwb_confidence(0.0, 1.5, quality=0.7) == pytest.approx(0.7)

    def test_large_age_returns_zero(self):
        assert uwb_confidence(1e9, 1.5) == 0.0

    def test_zero_timeout_returns_zero(self):
        assert uwb_confidence(0.0, 0.0) == 0.0


class TestCameraConfidence:

    def test_zero_age_full_quality(self):
        assert camera_confidence(0.0, 1.0, quality=1.0) == pytest.approx(1.0)

    def test_at_timeout_zero(self):
        assert camera_confidence(1.0, 1.0) == pytest.approx(0.0)

    def test_beyond_timeout_zero(self):
        assert camera_confidence(2.0, 1.0) == 0.0

    def test_quality_scales(self):
        # age=0, quality=0.8
        assert camera_confidence(0.0, 1.0, quality=0.8) == pytest.approx(0.8)

    def test_halfway(self):
        assert camera_confidence(0.5, 1.0) == pytest.approx(0.5)


class TestSelectSource:

    def test_both_above_threshold_fused(self):
        assert select_source(0.8, 0.6) == "fused"

    def test_only_uwb_above_threshold(self):
        assert select_source(0.8, 0.0) == "uwb"

    def test_only_cam_above_threshold(self):
        assert select_source(0.0, 0.7) == "camera"

    def test_both_below_threshold(self):
        assert select_source(0.0, 0.0) == "predicted"

    def test_threshold_boundary_uwb(self):
        # Exactly at threshold — should be treated as live
        assert select_source(0.15, 0.0, threshold=0.15) == "uwb"

    def test_threshold_boundary_below(self):
        assert select_source(0.14, 0.0, threshold=0.15) == "predicted"

    def test_custom_threshold(self):
        assert select_source(0.5, 0.0, threshold=0.6) == "predicted"
        assert select_source(0.5, 0.0, threshold=0.4) == "uwb"


class TestFusePositions:

    def test_equal_confidence_returns_midpoint(self):
        x, y = fuse_positions(0.0, 0.0, 1.0, 4.0, 4.0, 1.0)
        assert abs(x - 2.0) < 1e-9
        assert abs(y - 2.0) < 1e-9

    def test_full_uwb_weight_returns_uwb(self):
        x, y = fuse_positions(3.0, 1.0, 1.0, 0.0, 0.0, 0.0)
        assert abs(x - 3.0) < 1e-9

    def test_full_cam_weight_returns_cam(self):
        x, y = fuse_positions(0.0, 0.0, 0.0, -2.0, 5.0, 1.0)
        assert abs(x - (-2.0)) < 1e-9
        assert abs(y - 5.0) < 1e-9

    def test_weighted_blend(self):
        # UWB at (0,0) conf=3, camera at (4,0) conf=1 → x = 3/4*0 + 1/4*4 = 1
        x, y = fuse_positions(0.0, 0.0, 3.0, 4.0, 0.0, 1.0)
        assert abs(x - 1.0) < 1e-9

    def test_zero_total_returns_uwb_fallback(self):
        x, y = fuse_positions(7.0, 2.0, 0.0, 3.0, 1.0, 0.0)
        assert abs(x - 7.0) < 1e-9


class TestCompositeConfidence:

    def test_fused_source_high_confidence(self):
        conf = composite_confidence(0.9, 0.8, "fused", 0.05)
        assert conf > 0.7

    def test_predicted_source_capped(self):
        conf = composite_confidence(0.0, 0.0, "predicted", 0.1)
        assert conf <= 0.4

    def test_predicted_high_uncertainty_low_confidence(self):
        conf = composite_confidence(0.0, 0.0, "predicted", 3.0, max_kf_uncertainty_m=3.0)
        assert conf == pytest.approx(0.0)

    def test_uwb_only(self):
        conf = composite_confidence(0.8, 0.0, "uwb", 0.05)
        assert conf > 0.3

    def test_camera_only(self):
        conf = composite_confidence(0.0, 0.7, "camera", 0.05)
        assert conf > 0.2

    def test_high_kf_uncertainty_reduces_confidence(self):
        low_unc  = composite_confidence(0.9, 0.0, "uwb", 0.1)
        high_unc = composite_confidence(0.9, 0.0, "uwb", 2.9)
        assert low_unc > high_unc


class TestBearingAndRange:

    def test_straight_ahead(self):
        bearing, rng = bearing_and_range(2.0, 0.0)
        assert abs(bearing) < 1e-9
        assert abs(rng - 2.0) < 1e-9

    def test_left_of_robot(self):
        # +Y = left in base_link frame; bearing should be positive
        bearing, rng = bearing_and_range(0.0, 1.0)
        assert abs(bearing - math.pi / 2.0) < 1e-9
        assert abs(rng - 1.0) < 1e-9

    def test_right_of_robot(self):
        bearing, rng = bearing_and_range(0.0, -1.0)
        assert abs(bearing - (-math.pi / 2.0)) < 1e-9

    def test_diagonal(self):
        bearing, rng = bearing_and_range(1.0, 1.0)
        assert abs(bearing - math.pi / 4.0) < 1e-9
        assert abs(rng - math.sqrt(2.0)) < 1e-9

    def test_at_origin(self):
        bearing, rng = bearing_and_range(0.0, 0.0)
        assert rng == pytest.approx(0.0)
        assert math.isfinite(bearing)   # atan2(0,0) = 0 in most implementations

    def test_range_always_non_negative(self):
        for x, y in [(-1, 0), (0, -1), (-2, -3), (5, -5)]:
            _, rng = bearing_and_range(x, y)
            assert rng >= 0.0


# ─────────────────────────────────────────────────────────────────────────────
# Integration scenario tests
# ─────────────────────────────────────────────────────────────────────────────

class TestIntegrationScenarios:

    def test_euc_speed_velocity_tracking(self):
        """Verify KF can track EUC speed (8 m/s) within 0.5 m/s after warm-up."""
        kf = KalmanTracker(process_noise=3.0, meas_noise_uwb=0.20)
        kf.initialize(0.0, 0.0)
        dt = 1.0 / 10.0   # 10 Hz UWB rate
        speed = 8.0        # m/s (≈29 km/h)
        for i in range(60):
            t = i * dt
            kf.predict(dt)
            kf.update(speed * t, 0.0, "uwb")
        vx, vy = kf.velocity
        assert abs(vx - speed) < 0.6, f"vx={vx:.2f} expected≈{speed}"
        assert abs(vy) < 0.3

    def test_signal_loss_recovery(self):
        """
        After 1 s of signal loss the filter should still have a reasonable
        position estimate (not diverged to infinity).
        """
        kf = KalmanTracker(process_noise=3.0)
        kf.initialize(2.0, 0.5)
        # Warm up with 2 m/s x motion
        dt = 0.05
        for i in range(20):
            kf.predict(dt)
            kf.update(2.0 * (i + 1) * dt, 0.0, "uwb")
        # Coast for 1 second (20 × 50 ms) without measurements
        for _ in range(20):
            kf.predict(dt)
        x, y = kf.position
        assert math.isfinite(x) and math.isfinite(y)
        assert abs(x) < 20.0    # shouldn't have drifted more than 20 m

    def test_uwb_to_camera_handoff(self):
        """
        Simulate UWB going stale and camera taking over — Kalman should
        smoothly continue tracking without a jump.
        """
        kf = KalmanTracker(meas_noise_uwb=0.20, meas_noise_cam=0.12)
        kf.initialize(0.0, 0.0)
        dt = 0.05
        # Phase 1: UWB active
        for i in range(20):
            kf.predict(dt)
            kf.update(float(i) * 0.1, 0.0, "uwb")
        x_at_handoff, _ = kf.position

        # Phase 2: Camera takes over from same trajectory
        for i in range(20, 40):
            kf.predict(dt)
            kf.update(float(i) * 0.1, 0.0, "camera")
        x_after, _ = kf.position

        # Position should have continued progressing (not stuck or reset)
        assert x_after > x_at_handoff

    def test_confidence_degradation_during_coast(self):
        """Composite confidence should drop as KF uncertainty grows during coast."""
        kf = KalmanTracker(process_noise=3.0)
        kf.initialize(2.0, 0.0)

        # Fresh: tight uncertainty → high confidence
        unc_fresh = kf.position_uncertainty_m()
        conf_fresh = composite_confidence(0.0, 0.0, "predicted", unc_fresh)

        # After 2 s coast
        for _ in range(40):
            kf.predict(0.05)
        unc_stale = kf.position_uncertainty_m()
        conf_stale = composite_confidence(0.0, 0.0, "predicted", unc_stale)

        assert conf_fresh >= conf_stale

    def test_fused_source_confidence_weighted_position(self):
        """Fused position should sit between UWB and camera, closer to higher-conf source."""
        # UWB at x=0 with high conf, camera at x=10 with low conf
        uwb_c = 0.9
        cam_c = 0.1
        fx, fy = fuse_positions(0.0, 0.0, uwb_c, 10.0, 0.0, cam_c)
        # Should be much closer to UWB (0) than camera (10)
        assert fx < 3.0, f"fused_x={fx:.2f}"

    def test_select_source_transitions(self):
        """Verify correct source transitions as confidences change."""
        assert select_source(0.9, 0.8) == "fused"
        assert select_source(0.9, 0.0) == "uwb"
        assert select_source(0.0, 0.8) == "camera"
        assert select_source(0.0, 0.0) == "predicted"
