"""
test_obstacle_velocity.py — Unit tests for obstacle velocity estimation (no ROS2).

Covers:
  KalmanTrack — construction:
    - initial position matches centroid
    - initial velocity is zero
    - initial confidence is 0.0
    - initial alive=True, coasting=0, age=0

  KalmanTrack — predict():
    - predict(dt=0) leaves position unchanged
    - predict(dt=1) with zero velocity leaves position unchanged
    - predict() increments coasting by 1
    - alive stays True until coasting > max_coasting
    - alive becomes False exactly when coasting > max_coasting

  KalmanTrack — update():
    - update() resets coasting to 0
    - update() increments age
    - update() shifts position toward measurement
    - update() reduces position uncertainty (trace of P decreases)
    - confidence increases with age, caps at 1.0
    - confidence never exceeds 1.0
    - metadata stored (width, depth, point_count)

  KalmanTrack — velocity convergence:
    - constant-velocity target: vx converges toward true velocity after N steps
    - stationary target: speed stays near zero

  associate():
    - empty tracks → all clusters unmatched
    - empty clusters → all tracks unmatched
    - both empty → empty matches
    - single perfect match below max_dist
    - single pair above max_dist → no match
    - two tracks two clusters: diagonal nearest wins
    - more tracks than clusters: extra tracks unmatched
    - more clusters than tracks: extra clusters unmatched

  ObstacleTracker:
    - empty centroids → no tracks created
    - single cluster → one track, confidence=0 initially
    - same position twice → track maintained, confidence increases
    - cluster disappears: coasts then deleted after max_coasting+1 frames
    - new cluster after deletion gets new track_id
    - two clusters → two tracks, correct IDs
    - moving cluster: velocity direction correct after convergence
    - track_id is monotonically increasing
    - metadata (width, depth, point_count) stored on track
"""

import sys
import os
import math

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from saltybot_bringup._obstacle_velocity import (
    KalmanTrack,
    ObstacleTracker,
    associate,
)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _c(x: float, y: float) -> np.ndarray:
    return np.array([x, y], dtype=np.float64)


def _track(x: float = 0.0, y: float = 0.0, **kw) -> KalmanTrack:
    return KalmanTrack(1, _c(x, y), **kw)


# ── KalmanTrack — construction ────────────────────────────────────────────────

class TestKalmanTrackInit:

    def test_initial_position(self):
        t = _track(3.0, -1.5)
        assert t.position == pytest.approx([3.0, -1.5], abs=1e-9)

    def test_initial_velocity_zero(self):
        t = _track(1.0, 2.0)
        assert t.velocity == pytest.approx([0.0, 0.0], abs=1e-9)

    def test_initial_speed_zero(self):
        t = _track(1.0, 0.0)
        assert t.speed == pytest.approx(0.0, abs=1e-9)

    def test_initial_confidence_zero(self):
        t = KalmanTrack(1, _c(0, 0), n_init_frames=3)
        assert t.confidence == pytest.approx(0.0, abs=1e-9)

    def test_initial_alive(self):
        assert _track().alive is True

    def test_initial_coasting_zero(self):
        assert _track().coasting == 0

    def test_initial_age_zero(self):
        assert _track().age == 0


# ── KalmanTrack — predict ─────────────────────────────────────────────────────

class TestKalmanTrackPredict:

    def test_predict_dt_zero_no_position_change(self):
        t = _track(2.0, 3.0)
        t.predict(0.0)
        assert t.position == pytest.approx([2.0, 3.0], abs=1e-9)

    def test_predict_dt_positive_zero_velocity_no_position_change(self):
        t = _track(1.0, 1.0)
        t.predict(1.0)
        assert t.position == pytest.approx([1.0, 1.0], abs=1e-6)

    def test_predict_negative_dt_clamped_to_zero(self):
        t = _track(1.0, 0.0)
        t.predict(-5.0)
        assert t.position == pytest.approx([1.0, 0.0], abs=1e-9)

    def test_predict_increments_coasting(self):
        t = _track()
        assert t.coasting == 0
        t.predict(0.1)
        assert t.coasting == 1
        t.predict(0.1)
        assert t.coasting == 2

    def test_alive_before_max_coasting(self):
        t = KalmanTrack(1, _c(0, 0), max_coasting=3)
        for _ in range(3):
            t.predict(0.1)
        assert t.alive is True   # coast=3, 3 > 3 is False → still alive

    def test_alive_false_after_exceeding_max_coasting(self):
        t = KalmanTrack(1, _c(0, 0), max_coasting=3)
        for _ in range(4):
            t.predict(0.1)
        assert t.alive is False  # coast=4, 4 > 3 → dead

    def test_predict_advances_position_when_velocity_set(self):
        """After seeding velocity via update, predict advances x."""
        t = KalmanTrack(1, _c(0.0, 0.0), r_pos=0.001, n_init_frames=1)
        # Drive velocity to ~(1,0) by updating with advancing centroids
        for i in range(1, 8):
            t.predict(1.0)
            t.update(_c(float(i), 0.0))
        # Now predict one more step — position should advance in +x
        x0 = t.position[0]
        t.predict(1.0)
        assert t.position[0] > x0


# ── KalmanTrack — update ──────────────────────────────────────────────────────

class TestKalmanTrackUpdate:

    def test_update_resets_coasting(self):
        t = _track()
        t.predict(0.1)   # coast=1
        t.predict(0.1)   # coast=2
        t.update(_c(0, 0))
        assert t.coasting == 0

    def test_update_increments_age(self):
        t = _track()
        t.update(_c(1, 0))
        assert t.age == 1
        t.update(_c(2, 0))
        assert t.age == 2

    def test_update_shifts_position_toward_measurement(self):
        t = _track(0.0, 0.0)
        t.update(_c(5.0, 0.0))
        # Position should have moved in +x direction
        assert t.position[0] > 0.0

    def test_update_reduces_position_covariance(self):
        t = KalmanTrack(1, _c(0, 0))
        p_before = float(t._P[0, 0])
        t.update(_c(0, 0))
        p_after = float(t._P[0, 0])
        assert p_after < p_before

    def test_confidence_increases_with_age(self):
        t = KalmanTrack(1, _c(0, 0), n_init_frames=4)
        assert t.confidence == pytest.approx(0.0, abs=1e-9)
        t.update(_c(0, 0))
        assert t.confidence == pytest.approx(0.25)
        t.update(_c(0, 0))
        assert t.confidence == pytest.approx(0.50)
        t.update(_c(0, 0))
        assert t.confidence == pytest.approx(0.75)
        t.update(_c(0, 0))
        assert t.confidence == pytest.approx(1.0)

    def test_confidence_caps_at_one(self):
        t = KalmanTrack(1, _c(0, 0), n_init_frames=2)
        for _ in range(10):
            t.update(_c(0, 0))
        assert t.confidence == pytest.approx(1.0)

    def test_update_stores_metadata(self):
        t = _track()
        t.update(_c(1, 1), width=0.3, depth=0.5, point_count=7)
        assert t.last_width       == pytest.approx(0.3)
        assert t.last_depth       == pytest.approx(0.5)
        assert t.last_point_count == 7


# ── KalmanTrack — velocity convergence ───────────────────────────────────────

class TestKalmanTrackVelocityConvergence:

    def test_constant_velocity_converges(self):
        """
        Target at vx=1 m/s: observations at x=0,1,2,...
        After 15 predict+update cycles with low noise, vx should be near 1.0.
        """
        t = KalmanTrack(1, _c(0.0, 0.0), r_pos=0.01, q_pos=0.001, q_vel=0.1)
        for i in range(1, 16):
            t.predict(1.0)
            t.update(_c(float(i), 0.0))
        assert t.velocity[0] == pytest.approx(1.0, abs=0.15)
        assert t.velocity[1] == pytest.approx(0.0, abs=0.10)

    def test_diagonal_velocity_converges(self):
        """Target moving at (vx=1, vy=1) m/s."""
        t = KalmanTrack(1, _c(0.0, 0.0), r_pos=0.01, q_pos=0.001, q_vel=0.1)
        for i in range(1, 16):
            t.predict(1.0)
            t.update(_c(float(i), float(i)))
        assert t.velocity[0] == pytest.approx(1.0, abs=0.15)
        assert t.velocity[1] == pytest.approx(1.0, abs=0.15)

    def test_stationary_target_speed_near_zero(self):
        """Target at fixed position: estimated speed should stay small."""
        t = KalmanTrack(1, _c(2.0, 3.0), r_pos=0.001)
        for _ in range(15):
            t.predict(0.1)
            t.update(_c(2.0, 3.0))
        assert t.speed < 0.05


# ── associate ─────────────────────────────────────────────────────────────────

class TestAssociate:

    def test_empty_tracks(self):
        pos  = np.empty((0, 2))
        cent = np.array([[1.0, 0.0]])
        m, ut, uc = associate(pos, cent, 1.0)
        assert m  == []
        assert ut == []
        assert uc == [0]

    def test_empty_clusters(self):
        pos  = np.array([[1.0, 0.0]])
        cent = np.empty((0, 2))
        m, ut, uc = associate(pos, cent, 1.0)
        assert m  == []
        assert ut == [0]
        assert uc == []

    def test_both_empty(self):
        m, ut, uc = associate(np.empty((0, 2)), np.empty((0, 2)), 1.0)
        assert m  == []
        assert ut == []
        assert uc == []

    def test_single_match_below_threshold(self):
        pos  = np.array([[0.0, 0.0]])
        cent = np.array([[0.1, 0.0]])
        m, ut, uc = associate(pos, cent, 0.5)
        assert m  == [(0, 0)]
        assert ut == []
        assert uc == []

    def test_single_pair_above_threshold_no_match(self):
        pos  = np.array([[0.0, 0.0]])
        cent = np.array([[2.0, 0.0]])
        m, ut, uc = associate(pos, cent, 0.5)
        assert m  == []
        assert ut == [0]
        assert uc == [0]

    def test_two_tracks_two_clusters_diagonal(self):
        pos  = np.array([[0.0, 0.0], [3.0, 0.0]])
        cent = np.array([[0.1, 0.0], [3.1, 0.0]])
        m, ut, uc = associate(pos, cent, 0.5)
        assert (0, 0) in m
        assert (1, 1) in m
        assert ut == []
        assert uc == []

    def test_more_tracks_than_clusters(self):
        pos  = np.array([[0.0, 0.0], [5.0, 0.0]])
        cent = np.array([[0.1, 0.0]])
        m, ut, uc = associate(pos, cent, 0.5)
        assert len(m)  == 1
        assert len(ut) == 1   # one track unmatched
        assert uc == []

    def test_more_clusters_than_tracks(self):
        pos  = np.array([[0.0, 0.0]])
        cent = np.array([[0.1, 0.0], [5.0, 0.0]])
        m, ut, uc = associate(pos, cent, 0.5)
        assert len(m)  == 1
        assert ut == []
        assert len(uc) == 1   # one cluster unmatched

    def test_nearest_wins(self):
        """Track 0 is closest to cluster 0; ensure it's matched to cluster 0."""
        pos  = np.array([[0.0, 0.0], [10.0, 0.0]])
        cent = np.array([[0.05, 0.0], [10.2, 0.0]])
        m, ut, uc = associate(pos, cent, 1.0)
        assert (0, 0) in m
        assert (1, 1) in m

    def test_threshold_strictly_less_than(self):
        """Distance exactly equal to max_dist should NOT match."""
        pos  = np.array([[0.0, 0.0]])
        cent = np.array([[0.5, 0.0]])
        m, ut, uc = associate(pos, cent, 0.5)   # dist=0.5, max_dist=0.5
        assert m == []


# ── ObstacleTracker ───────────────────────────────────────────────────────────

class TestObstacleTracker:

    def test_empty_input_no_tracks(self):
        tracker = ObstacleTracker()
        result  = tracker.update([], timestamp=0.0)
        assert result == []

    def test_single_cluster_creates_track(self):
        tracker = ObstacleTracker()
        tracks  = tracker.update([_c(1, 0)], timestamp=0.0)
        assert len(tracks) == 1

    def test_track_id_starts_at_one(self):
        tracker = ObstacleTracker()
        tracks  = tracker.update([_c(0, 0)], timestamp=0.0)
        assert tracks[0].track_id == 1

    def test_same_position_twice_single_track(self):
        tracker = ObstacleTracker()
        t1 = tracker.update([_c(1, 0)], timestamp=0.0)
        t2 = tracker.update([_c(1, 0)], timestamp=0.1)
        assert len(t2) == 1
        assert t2[0].track_id == t1[0].track_id

    def test_confidence_increases_with_updates(self):
        tracker = ObstacleTracker(n_init_frames=3)
        tracks  = tracker.update([_c(0, 0)], timestamp=0.0)
        c0      = tracks[0].confidence
        tracks  = tracker.update([_c(0, 0)], timestamp=0.1)
        assert tracks[0].confidence > c0

    def test_track_coasts_then_dies(self):
        tracker = ObstacleTracker(max_coasting_frames=2)
        tracker.update([_c(0, 0)], timestamp=0.0)   # create
        t1 = tracker.update([], timestamp=0.1)      # coast 1
        assert len(t1) == 1
        t2 = tracker.update([], timestamp=0.2)      # coast 2
        assert len(t2) == 1
        t3 = tracker.update([], timestamp=0.3)      # coast 3 > 2 → dead
        assert len(t3) == 0

    def test_new_cluster_after_deletion_new_id(self):
        tracker = ObstacleTracker(max_coasting_frames=1)
        t0 = tracker.update([_c(0, 0)], timestamp=0.0)
        old_id = t0[0].track_id
        tracker.update([], timestamp=0.1)  # coast 1
        tracker.update([], timestamp=0.2)  # coast 2 > 1 → dead
        t1 = tracker.update([_c(0, 0)], timestamp=0.3)
        assert len(t1) == 1
        assert t1[0].track_id != old_id

    def test_two_clusters_two_tracks(self):
        tracker = ObstacleTracker()
        tracks  = tracker.update([_c(0, 0), _c(5, 0)], timestamp=0.0)
        assert len(tracks) == 2
        ids = {t.track_id for t in tracks}
        assert len(ids) == 2

    def test_track_ids_monotonically_increasing(self):
        tracker = ObstacleTracker()
        tracker.update([_c(0, 0)],    timestamp=0.0)
        tracker.update([_c(10, 10)],  timestamp=0.1)   # far → new track
        all_ids = [t.track_id for t in tracker.tracks.values()]
        assert all_ids == sorted(all_ids)

    def test_moving_cluster_velocity_direction(self):
        """
        Cluster moves in +x at 1 m/s; after convergence vx should be positive.
        Use low noise and many steps for reliable convergence.
        """
        tracker = ObstacleTracker(
            n_init_frames=1,
            r_pos=0.01,
            q_pos=0.001,
            q_vel=0.1,
        )
        for i in range(20):
            tracker.update([_c(float(i) * 0.1, 0.0)], timestamp=float(i) * 0.1)
        tracks = list(tracker.tracks.values())
        assert len(tracks) == 1
        assert tracks[0].velocity[0] > 0.05

    def test_metadata_stored_on_track(self):
        tracker = ObstacleTracker()
        tracker.update(
            [_c(1, 1)],
            timestamp    = 0.0,
            widths       = [0.4],
            depths       = [0.6],
            point_counts = [9],
        )
        t = list(tracker.tracks.values())[0]
        assert t.last_width       == pytest.approx(0.4)
        assert t.last_depth       == pytest.approx(0.6)
        assert t.last_point_count == 9

    def test_far_cluster_creates_new_track(self):
        """Cluster beyond max_association_dist creates a second track."""
        tracker = ObstacleTracker(max_association_dist_m=0.5)
        tracker.update([_c(0, 0)], timestamp=0.0)
        tracks = tracker.update([_c(10, 0)], timestamp=0.1)
        # Original track coasts, new track spawned for far cluster
        assert len(tracks) == 2

    def test_empty_to_single_and_back(self):
        tracker = ObstacleTracker(max_coasting_frames=0)
        t1 = tracker.update([_c(1, 0)], timestamp=0.0)
        assert len(t1) == 1
        t2 = tracker.update([], timestamp=0.1)   # coast=1 > 0 → dead
        assert len(t2) == 0
        t3 = tracker.update([_c(1, 0)], timestamp=0.2)
        assert len(t3) == 1

    def test_constant_velocity_estimate(self):
        """
        Target moves at vx=0.1 m/s (0.1 m per 1-second step).
        Each step is well within the default max_association_dist_m=0.5 m,
        so the track is continuously matched.  After many updates, estimated
        speed should be close to 0.1 m/s.
        """
        tracker = ObstacleTracker(
            n_init_frames=1, r_pos=0.001, q_pos=0.0001, q_vel=0.05
        )
        for i in range(30):
            tracker.update([_c(float(i) * 0.1, 0.0)], timestamp=float(i))
        t = list(tracker.tracks.values())[0]
        assert t.speed == pytest.approx(0.1, abs=0.03)


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