"""
test_landmark_smoother.py — Unit tests for LandmarkKalmanFilter (no ROS2 required).

Covers:
  - Initialisation: first call returns observations unchanged
  - Smoothing: output tracks input over multiple steps
  - Unit vector: direction is preserved
  - Shape / dtype
  - dt handling: zero vs positive dt
  - Noise reduction: repeated identical observations converge
  - Track key logic tested via helper
  - Fallback for short landmark arrays (pass-through guard)
"""

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_social_face._landmark_kalman import LandmarkKalmanFilter


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

def _random_lm(n=68, seed=0):
    rng = np.random.default_rng(seed)
    return rng.uniform(0.1, 0.9, size=(n, 2)).astype(np.float32)


def _zero_lm(n=68):
    return np.zeros((n, 2), dtype=np.float32)


# ── Shape and dtype ───────────────────────────────────────────────────────────

class TestOutputShape:

    def test_output_shape_68(self):
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm(68)
        out = kf.update(obs, dt=0.0)
        assert out.shape == (68, 2)

    def test_output_dtype_float32(self):
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm(68)
        out = kf.update(obs, dt=0.0)
        assert out.dtype == np.float32

    def test_custom_n_points(self):
        kf  = LandmarkKalmanFilter(n_points=5)
        obs = _random_lm(5)
        out = kf.update(obs, dt=0.0)
        assert out.shape == (5, 2)

    def test_wrong_shape_raises(self):
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm(5)          # wrong number of points
        with pytest.raises(ValueError):
            kf.update(obs, dt=0.0)


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

class TestInitialisation:

    def test_first_update_returns_obs(self):
        """On first call the filter has not yet built momentum, so output ≈ input."""
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm(68)
        out = kf.update(obs, dt=0.0)
        np.testing.assert_allclose(out, obs, atol=1e-5)

    def test_not_initialised_before_update(self):
        kf = LandmarkKalmanFilter(n_points=68)
        assert not kf.is_initialised

    def test_initialised_after_first_update(self):
        kf = LandmarkKalmanFilter(n_points=68)
        kf.update(_random_lm(), dt=0.0)
        assert kf.is_initialised


# ── Smoothing behaviour ───────────────────────────────────────────────────────

class TestSmoothingBehaviour:

    def test_stationary_obs_converges(self):
        """Repeated identical observations should produce output ≈ input."""
        kf  = LandmarkKalmanFilter(n_points=68, obs_noise_std=0.005)
        obs = _random_lm(68, seed=1)
        for _ in range(30):
            out = kf.update(obs, dt=1.0 / 30.0)
        np.testing.assert_allclose(out, obs, atol=0.01)

    def test_output_stays_bounded(self):
        """Output should remain within [0, 1] when inputs are in [0, 1]."""
        kf  = LandmarkKalmanFilter(n_points=68)
        rng = np.random.default_rng(42)
        for _ in range(50):
            obs = rng.uniform(0.05, 0.95, size=(68, 2)).astype(np.float32)
            out = kf.update(obs, dt=1.0 / 30.0)
        # Allow a small margin beyond [0,1] due to velocity extrapolation
        assert out.min() > -0.1
        assert out.max() <  1.1

    def test_smoothing_reduces_noise(self):
        """
        With low process noise the filter should dampen high-frequency jitter.
        Average per-frame change in output should be less than in noisy input.
        """
        kf  = LandmarkKalmanFilter(n_points=68, process_noise_std=0.002, obs_noise_std=0.01)
        rng = np.random.default_rng(7)
        base  = rng.uniform(0.2, 0.8, size=(68, 2)).astype(np.float32)
        noise = 0.05  # std-dev of added noise
        prev_obs = None
        prev_out = None
        obs_deltas = []
        out_deltas = []
        for _ in range(40):
            obs = (base + rng.normal(0, noise, size=(68, 2))).astype(np.float32)
            out = kf.update(obs, dt=1.0 / 30.0)
            if prev_obs is not None:
                obs_deltas.append(np.abs(obs - prev_obs).mean())
                out_deltas.append(np.abs(out - prev_out).mean())
            prev_obs = obs.copy()
            prev_out = out.copy()
        assert np.mean(out_deltas) < np.mean(obs_deltas)

    def test_velocity_prediction_follows_linear_motion(self):
        """
        Points moving linearly should be predicted ahead-of-time so the
        filter output lags by less than the raw measurement noise.
        """
        kf  = LandmarkKalmanFilter(n_points=68, process_noise_std=0.005, obs_noise_std=0.001)
        dt  = 1.0 / 30.0
        vel = 0.002  # normalised coords/frame
        obs = np.tile([0.3, 0.4], (68, 1)).astype(np.float32)
        # Warm up the filter
        for i in range(20):
            obs_step = obs + vel * i
            kf.update(obs_step, dt=dt)
        # After warm-up, output should track the moving target
        out = kf.update(obs + vel * 20, dt=dt)
        target = obs + vel * 20
        np.testing.assert_allclose(out, target, atol=0.01)


# ── dt handling ───────────────────────────────────────────────────────────────

class TestDtHandling:

    def test_zero_dt_on_first_call_ok(self):
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm()
        out = kf.update(obs, dt=0.0)   # must not crash
        assert out.shape == (68, 2)

    def test_large_dt_does_not_crash(self):
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm()
        kf.update(obs, dt=0.0)
        out = kf.update(obs, dt=5.0)   # 5 second gap
        assert out.shape == (68, 2)

    def test_multiple_calls_increase_smoothness(self):
        """Calling update repeatedly should not raise exceptions."""
        kf  = LandmarkKalmanFilter(n_points=68)
        obs = _random_lm()
        for i in range(100):
            kf.update(obs, dt=1.0 / 30.0)


# ── Process and observation noise parameters ──────────────────────────────────

class TestNoiseParameters:

    def test_high_obs_noise_smooths_more(self):
        """Higher obs noise → output moves less per step (more smoothing)."""
        rng  = np.random.default_rng(99)
        base = rng.uniform(0.2, 0.8, size=(68, 2)).astype(np.float32)

        kf_low  = LandmarkKalmanFilter(n_points=68, obs_noise_std=0.001)
        kf_high = LandmarkKalmanFilter(n_points=68, obs_noise_std=0.1)

        for i in range(10):
            obs = (base + rng.normal(0, 0.02, size=(68, 2))).astype(np.float32)
            kf_low.update(obs,  dt=1.0 / 30.0)
            kf_high.update(obs, dt=1.0 / 30.0)

        obs_new = (base + rng.normal(0, 0.02, size=(68, 2))).astype(np.float32)
        prev_low  = kf_low._x[:, :2].copy()
        prev_high = kf_high._x[:, :2].copy()
        kf_low.update(obs_new,  dt=1.0 / 30.0)
        kf_high.update(obs_new, dt=1.0 / 30.0)

        delta_low  = np.abs(kf_low._x[:, :2]  - prev_low).mean()
        delta_high = np.abs(kf_high._x[:, :2] - prev_high).mean()
        assert delta_high < delta_low

    def test_known_stationary_value(self):
        """With very low noise a stationary signal should converge closely."""
        kf  = LandmarkKalmanFilter(n_points=1, process_noise_std=1e-4, obs_noise_std=1e-4)
        obs = np.array([[0.5, 0.5]], dtype=np.float32)
        for _ in range(50):
            out = kf.update(obs, dt=1.0 / 30.0)
        np.testing.assert_allclose(out, obs, atol=0.001)


# ── Track key helper (reimplemented to mirror node logic without importing ROS) ─

def _track_key(person_id: int, face_id: int) -> int:
    """Mirror of LandmarkSmootherNode._track_key logic."""
    return person_id if person_id >= 0 else -(face_id + 1)


class TestTrackKeyLogic:

    def test_positive_person_id_used(self):
        assert _track_key(person_id=3, face_id=7) == 3

    def test_zero_person_id_used(self):
        assert _track_key(person_id=0, face_id=7) == 0

    def test_minus_one_person_id_uses_face_id(self):
        key = _track_key(person_id=-1, face_id=2)
        assert key == -(2 + 1)   # == -3

    def test_different_face_ids_give_different_keys(self):
        k0 = _track_key(person_id=-1, face_id=0)
        k1 = _track_key(person_id=-1, face_id=1)
        assert k0 != k1

    def test_person_key_does_not_collide_with_face_key(self):
        """Person ID 3 and face_id 3 should not produce the same key."""
        kp = _track_key(person_id=3,  face_id=99)
        kf = _track_key(person_id=-1, face_id=3)
        assert kp != kf


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