"""Tests for Kalman filter suite -- KalmanFilter and UnscentedKF."""

import math
import pytest

from horizon._horizon import KalmanFilter, UnscentedKF


# ===========================================================================
# KalmanFilter -- Construction
# ===========================================================================


class TestKalmanFilterConstruction:
    def test_basic_creation(self):
        kf = KalmanFilter(2, 1)
        assert len(kf.state()) == 2
        assert len(kf.covariance()) == 2

    def test_initial_state_is_zero(self):
        kf = KalmanFilter(3, 2)
        assert kf.state() == [0.0, 0.0, 0.0]

    def test_initial_covariance_is_identity(self):
        kf = KalmanFilter(2, 1)
        cov = kf.covariance()
        assert cov[0][0] == pytest.approx(1.0)
        assert cov[1][1] == pytest.approx(1.0)
        assert cov[0][1] == pytest.approx(0.0)

    def test_zero_state_dim_raises(self):
        with pytest.raises(ValueError):
            KalmanFilter(0, 1)

    def test_zero_obs_dim_raises(self):
        with pytest.raises(ValueError):
            KalmanFilter(1, 0)

    def test_too_large_dim_raises(self):
        with pytest.raises(ValueError):
            KalmanFilter(65, 1)

    def test_too_large_obs_dim_raises(self):
        with pytest.raises(ValueError):
            KalmanFilter(1, 65)

    def test_max_allowed_dim(self):
        kf = KalmanFilter(64, 64)
        assert len(kf.state()) == 64

    def test_repr(self):
        kf = KalmanFilter(2, 1)
        r = repr(kf)
        assert "KalmanFilter" in r
        assert "2" in r


# ===========================================================================
# KalmanFilter -- Setters
# ===========================================================================


class TestKalmanFilterSetters:
    def test_set_transition(self):
        kf = KalmanFilter(2, 1)
        kf.set_transition([[1.0, 1.0], [0.0, 1.0]])

    def test_set_transition_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_transition([[1.0]])

    def test_set_transition_non_finite_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_transition([[float("nan"), 0.0], [0.0, 1.0]])

    def test_set_observation(self):
        kf = KalmanFilter(2, 1)
        kf.set_observation([[1.0, 0.0]])

    def test_set_observation_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_observation([[1.0, 0.0], [0.0, 1.0]])

    def test_set_process_noise(self):
        kf = KalmanFilter(2, 1)
        kf.set_process_noise([[0.1, 0.0], [0.0, 0.1]])

    def test_set_process_noise_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_process_noise([[0.1]])

    def test_set_measurement_noise(self):
        kf = KalmanFilter(2, 1)
        kf.set_measurement_noise([[0.5]])

    def test_set_measurement_noise_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_measurement_noise([[0.5, 0.0], [0.0, 0.5]])

    def test_set_state(self):
        kf = KalmanFilter(2, 1)
        kf.set_state([5.0, 3.0])
        assert kf.state() == [5.0, 3.0]

    def test_set_state_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_state([1.0])

    def test_set_state_non_finite_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.set_state([float("inf"), 0.0])

    def test_set_covariance(self):
        kf = KalmanFilter(2, 1)
        kf.set_covariance([[2.0, 0.0], [0.0, 2.0]])
        cov = kf.covariance()
        assert cov[0][0] == pytest.approx(2.0)

    def test_set_process_noise_matrix(self):
        kf = KalmanFilter(2, 1)
        kf.set_process_noise([[0.5, 0.0], [0.0, 0.5]])
        kf.predict()
        cov = kf.covariance()
        assert cov[0][0] > 1.0  # Grew from identity by Q

    def test_set_measurement_noise_matrix(self):
        kf = KalmanFilter(2, 1)
        kf.set_measurement_noise([[0.5]])
        kf.predict()
        state = kf.update([1.0])
        assert len(state) == 2


# ===========================================================================
# KalmanFilter -- Predict
# ===========================================================================


class TestKalmanFilterPredict:
    def test_predict_returns_state(self):
        kf = KalmanFilter(2, 1)
        pred = kf.predict()
        assert len(pred) == 2

    def test_predict_with_identity_transition(self):
        kf = KalmanFilter(1, 1)
        kf.set_state([5.0])
        pred = kf.predict()
        assert pred[0] == pytest.approx(5.0)

    def test_predict_covariance_grows(self):
        kf = KalmanFilter(1, 1)
        cov_before = kf.covariance()[0][0]
        kf.predict()
        cov_after = kf.covariance()[0][0]
        # P_pred = F*P*F' + Q >= P
        assert cov_after >= cov_before

    def test_predict_with_transition_matrix(self):
        kf = KalmanFilter(2, 1)
        kf.set_transition([[1.0, 1.0], [0.0, 1.0]])
        kf.set_state([0.0, 2.0])
        pred = kf.predict()
        # position = 0 + 1*2 = 2, velocity stays 2
        assert pred[0] == pytest.approx(2.0)
        assert pred[1] == pytest.approx(2.0)

    def test_predict_multiple_steps(self):
        kf = KalmanFilter(1, 1)
        kf.set_state([0.0])
        kf.set_process_noise([[0.01]])
        for _ in range(10):
            kf.predict()
        cov = kf.covariance()[0][0]
        # After 10 predictions, covariance should have grown substantially
        assert cov > 1.0


# ===========================================================================
# KalmanFilter -- Update
# ===========================================================================


class TestKalmanFilterUpdate:
    def test_update_returns_state(self):
        kf = KalmanFilter(1, 1)
        kf.predict()
        state = kf.update([3.0])
        assert len(state) == 1

    def test_update_moves_toward_measurement(self):
        kf = KalmanFilter(1, 1)
        kf.set_state([0.0])
        kf.predict()
        state = kf.update([10.0])
        assert state[0] > 0.0  # moved toward 10

    def test_update_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        kf.predict()
        with pytest.raises(ValueError):
            kf.update([1.0, 2.0])  # obs_dim is 1

    def test_update_non_finite_raises(self):
        kf = KalmanFilter(1, 1)
        kf.predict()
        with pytest.raises(ValueError):
            kf.update([float("nan")])

    def test_covariance_shrinks_after_update(self):
        kf = KalmanFilter(1, 1)
        kf.predict()
        cov_pred = kf.covariance()[0][0]
        kf.update([5.0])
        cov_upd = kf.covariance()[0][0]
        assert cov_upd <= cov_pred

    def test_convergence_on_constant_signal(self):
        kf = KalmanFilter(1, 1)
        for _ in range(50):
            kf.predict()
            kf.update([7.0])
        assert kf.state()[0] == pytest.approx(7.0, abs=0.5)

    def test_update_with_multivariate_obs(self):
        kf = KalmanFilter(2, 2)
        kf.predict()
        state = kf.update([1.0, 2.0])
        assert len(state) == 2


# ===========================================================================
# KalmanFilter -- Log Likelihood
# ===========================================================================


class TestKalmanFilterLogLikelihood:
    def test_log_likelihood_is_negative(self):
        kf = KalmanFilter(1, 1)
        kf.predict()
        ll = kf.log_likelihood([0.0])
        assert ll <= 0.0

    def test_log_likelihood_higher_for_closer_measurement(self):
        kf = KalmanFilter(1, 1)
        kf.predict()
        ll_close = kf.log_likelihood([0.0])  # state is ~0
        ll_far = kf.log_likelihood([100.0])
        assert ll_close > ll_far

    def test_log_likelihood_wrong_size_raises(self):
        kf = KalmanFilter(2, 1)
        with pytest.raises(ValueError):
            kf.log_likelihood([1.0, 2.0])

    def test_log_likelihood_is_finite(self):
        kf = KalmanFilter(1, 1)
        kf.predict()
        ll = kf.log_likelihood([5.0])
        assert math.isfinite(ll)


# ===========================================================================
# KalmanFilter -- Tracking scenario
# ===========================================================================


class TestKalmanFilterTracking:
    def test_position_velocity_model(self):
        """2D state (position, velocity), 1D observation (position)."""
        kf = KalmanFilter(2, 1)
        dt = 1.0
        kf.set_transition([[1.0, dt], [0.0, 1.0]])
        kf.set_observation([[1.0, 0.0]])
        kf.set_process_noise([[0.01, 0.0], [0.0, 0.01]])
        kf.set_measurement_noise([[1.0]])

        # Feed linearly increasing measurements (velocity ~ 1)
        for t in range(20):
            kf.predict()
            kf.update([float(t)])

        state = kf.state()
        # Position should track
        assert state[0] == pytest.approx(19.0, abs=2.0)
        # Velocity should converge toward 1
        assert state[1] == pytest.approx(1.0, abs=0.5)

    def test_noisy_measurement_still_converges(self):
        """Filter should track through noise."""
        kf = KalmanFilter(1, 1)
        kf.set_measurement_noise([[4.0]])
        kf.set_process_noise([[0.01]])
        # Measurements oscillate around 10
        measurements = [10.0 + ((-1) ** i) * 2.0 for i in range(40)]
        for m in measurements:
            kf.predict()
            kf.update([m])
        assert kf.state()[0] == pytest.approx(10.0, abs=3.0)


# ===========================================================================
# UnscentedKF -- Construction
# ===========================================================================


class TestUnscentedKFConstruction:
    def test_basic_creation(self):
        ukf = UnscentedKF(2, 1)
        assert len(ukf.state()) == 2

    def test_custom_params(self):
        ukf = UnscentedKF(2, 1, alpha=0.1, beta=2.0, kappa=1.0)
        assert len(ukf.state()) == 2

    def test_zero_state_dim_raises(self):
        with pytest.raises(ValueError):
            UnscentedKF(0, 1)

    def test_zero_obs_dim_raises(self):
        with pytest.raises(ValueError):
            UnscentedKF(1, 0)

    def test_too_large_dim_raises(self):
        with pytest.raises(ValueError):
            UnscentedKF(65, 1)

    def test_alpha_zero_raises(self):
        with pytest.raises(ValueError):
            UnscentedKF(2, 1, alpha=0.0)

    def test_alpha_negative_raises(self):
        with pytest.raises(ValueError):
            UnscentedKF(2, 1, alpha=-0.1)

    def test_repr(self):
        ukf = UnscentedKF(2, 1)
        r = repr(ukf)
        assert "UnscentedKF" in r

    def test_initial_state_is_zero(self):
        ukf = UnscentedKF(3, 2)
        assert ukf.state() == [0.0, 0.0, 0.0]


# ===========================================================================
# UnscentedKF -- Operations
# ===========================================================================


class TestUnscentedKFOperations:
    def test_predict(self):
        ukf = UnscentedKF(1, 1)
        pred = ukf.predict()
        assert len(pred) == 1

    def test_update(self):
        ukf = UnscentedKF(1, 1)
        ukf.predict()
        state = ukf.update([5.0])
        assert len(state) == 1

    def test_update_wrong_size_raises(self):
        ukf = UnscentedKF(2, 1)
        ukf.predict()
        with pytest.raises(ValueError):
            ukf.update([1.0, 2.0])

    def test_update_non_finite_raises(self):
        ukf = UnscentedKF(1, 1)
        ukf.predict()
        with pytest.raises(ValueError):
            ukf.update([float("inf")])

    def test_convergence(self):
        ukf = UnscentedKF(1, 1)
        for _ in range(50):
            ukf.predict()
            ukf.update([3.0])
        assert ukf.state()[0] == pytest.approx(3.0, abs=1.0)

    def test_covariance(self):
        ukf = UnscentedKF(2, 1)
        cov = ukf.covariance()
        assert len(cov) == 2
        assert len(cov[0]) == 2

    def test_log_likelihood(self):
        ukf = UnscentedKF(1, 1)
        ukf.predict()
        ll = ukf.log_likelihood([0.0])
        assert ll <= 0.0

    def test_set_state(self):
        ukf = UnscentedKF(2, 1)
        ukf.set_state([10.0, 5.0])
        assert ukf.state() == [10.0, 5.0]

    def test_setters(self):
        ukf = UnscentedKF(2, 1)
        ukf.set_transition([[1.0, 1.0], [0.0, 1.0]])
        ukf.set_observation([[1.0, 0.0]])
        ukf.set_process_noise([[0.01, 0.0], [0.0, 0.01]])
        ukf.set_measurement_noise([[1.0]])
        ukf.set_covariance([[2.0, 0.0], [0.0, 2.0]])

    def test_covariance_shrinks_after_update(self):
        ukf = UnscentedKF(1, 1)
        ukf.predict()
        cov_before = ukf.covariance()[0][0]
        ukf.update([1.0])
        cov_after = ukf.covariance()[0][0]
        assert cov_after <= cov_before

    def test_predict_covariance_grows(self):
        ukf = UnscentedKF(1, 1)
        cov_before = ukf.covariance()[0][0]
        ukf.predict()
        cov_after = ukf.covariance()[0][0]
        assert cov_after >= cov_before
