#!/usr/bin/env python3
"""Test position estimator and EKF."""

import pytest
import numpy as np


def test_ekf_initialization():
    from filterpy.kalman import ExtendedKalmanFilter
    
    ekf = ExtendedKalmanFilter(dim_x=9, dim_z=6)
    ekf.x = np.zeros(9)
    ekf.P = np.eye(9)
    ekf.Q = np.eye(9) * 0.1
    ekf.R = np.eye(6) * 0.1
    
    assert ekf.x.shape == (9,)
    assert ekf.P.shape == (9, 9)


def test_state_transition():
    dt = 0.02
    x = np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.0, 0.0, 0.0])
    
    F = np.eye(9)
    F[0, 3] = dt
    F[1, 4] = dt
    F[2, 5] = dt
    
    x_new = F @ x
    
    assert np.isclose(x_new[0], 1.0 + 0.1 * dt)
    assert np.isclose(x_new[1], 2.0 + 0.2 * dt)
    assert np.isclose(x_new[2], 3.0 + 0.3 * dt)


def test_measurement_update():
    from filterpy.kalman import ExtendedKalmanFilter
    
    ekf = ExtendedKalmanFilter(dim_x=6, dim_z=3)
    ekf.x = np.zeros(6)
    ekf.P = np.eye(6)
    ekf.Q = np.eye(6) * 0.1
    ekf.R = np.eye(3) * 0.1
    ekf.F = np.eye(6)
    
    H = np.zeros((3, 6))
    H[0, 0] = 1.0
    H[1, 1] = 1.0
    H[2, 2] = 1.0
    
    z = np.array([1.0, 2.0, 3.0])
    
    ekf.update(z, HJacobian=lambda x: H, Hx=lambda x: H @ x)
    
    assert ekf.x[0] > 0


def test_weighted_average():
    weights = [0.6, 0.3, 0.1]
    values = [np.array([1, 0, 0]), np.array([0, 1, 0]), np.array([0, 0, 1])]
    
    total = sum([w * v for w, v in zip(weights, values)])
    
    assert np.isclose(total[0], 0.6)
    assert np.isclose(total[1], 0.3)
    assert np.isclose(total[2], 0.1)


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