"""
Test module for the Kalman Filter estimation module.

Authors: Sebastián Rodríguez-Martínez
Contact: srodriguez@mbari.org
"""

import inspect

import numpy as np

from navlib import UUV_EKF, kf_lti_discretize, kf_predict, kf_update
from navlib.nav import UUV_EKF as UUV_EKF_relative
from navlib.nav import kf_lti_discretize as kf_lti_discretize_relative
from navlib.nav import kf_predict as kf_predict_relative
from navlib.nav import kf_update as kf_update_relative


def test_imports():
    # Checks that the absolute and relative imports work correctly for functions
    assert (
        kf_lti_discretize_relative.__code__.co_code
        == kf_lti_discretize.__code__.co_code
    )
    assert kf_predict_relative.__code__.co_code == kf_predict.__code__.co_code
    assert kf_update_relative.__code__.co_code == kf_update.__code__.co_code

    # For classes, compare their method code objects
    methods = [
        "initialize_state",
        "predict",
        "correction_position_xy",
        "correction_position_xyz",
        "correction_depth",
        "correction_attitude",
        "correction_angular_rate",
        "correction_velocity",
        "correction_bottom_height",
        "_state_transition_jacobian",
        "_process_noise_covariance",
    ]
    for method in methods:
        rel_method = getattr(UUV_EKF_relative, method)
        abs_method = getattr(UUV_EKF, method)
        assert inspect.getsource(rel_method) == inspect.getsource(abs_method)

    # Compare property names
    rel_props = {
        k for k, v in vars(UUV_EKF_relative).items() if isinstance(v, property)
    }
    abs_props = {k for k, v in vars(UUV_EKF).items() if isinstance(v, property)}
    assert rel_props == abs_props


def test_kf_lti_discretize():
    """
    Test the kf_lti_discretize function.
    """
    # Test case 1: Valid input
    Ac = np.array([[0, 1], [-1, -1]])
    Bc = np.array([[0], [1]])
    Qc = np.array([[1, 0], [0, 1]])
    dt = 0.1
    Ad, Bd, Qd = kf_lti_discretize(Ac, Bc, Qc, dt)
    assert Ad.shape == (2, 2)
    assert Bd.shape == (2, 1)
    assert Qd.shape == (2, 2)

    # Test case 2: List input
    Ad, Bd, Qd = kf_lti_discretize(Ac.tolist(), Bc.tolist(), Qc.tolist(), dt)
    assert Ad.shape == (2, 2)
    assert Bd.shape == (2, 1)
    assert Qd.shape == (2, 2)

    # Test case 3: Default Bc and Qc
    Ad, Bd, Qd = kf_lti_discretize(Ac, dt=dt)
    assert Ad.shape == (2, 2)
    assert Bd.shape == (2, 1)
    assert Qd.shape == (2, 2)

    # Test case 4: Invalid Ac type
    try:
        kf_lti_discretize("invalid_type", Bc, Qc, dt)
    except TypeError as e:
        assert str(e) == "Ac must be a numpy array"
    else:
        assert False

    # Test case 5: Invalid Bc type
    try:
        kf_lti_discretize(Ac, "invalid_type", Qc, dt)
    except TypeError as e:
        assert str(e) == "Bc must be a numpy array"
    else:
        assert False

    # Test case 6: Invalid Qc type
    try:
        kf_lti_discretize(Ac, Bc, "invalid_type", dt)
    except TypeError as e:
        assert str(e) == "Qc must be a numpy array"
    else:
        assert False

    # Test case 7: Invalid dt type
    try:
        kf_lti_discretize(Ac, Bc, Qc, "invalid_type")
    except TypeError as e:
        assert str(e) == "dt must be a number"
    else:
        assert False

    # Test case 8: Invalid Ac shape
    try:
        kf_lti_discretize(np.array([1, 2, 3, 4]), Bc, Qc, dt)
    except ValueError as e:
        assert str(e) == "Ac must be a 2D matrix"
    else:
        assert False

    # Test case 9: Invalid Bc shape
    try:
        kf_lti_discretize(Ac, np.array([1, 2, 3]), Qc, dt)
    except ValueError as e:
        assert str(e) == "Ac and Bc must have the same number of rows"
    else:
        assert False

    # Test case 10: Invalid Qc shape (not square)
    try:
        kf_lti_discretize(Ac, Bc, np.array([[1, 0, 0], [0, 1, 0]]), dt)
    except ValueError as e:
        assert str(e) == "Qc must be a square matrix"
    else:
        assert False

    # Test case 11: Invalid Qc shape (mismatch with Ac)
    try:
        kf_lti_discretize(Ac, Bc, np.array([[1, 0, 0], [0, 1, 0], [0, 1, 0]]), dt)
    except ValueError as e:
        assert str(e) == "Qc must have the same number of rows as Ac"
    else:
        assert False


def test_uuv_ekf_state_consistency():
    # Test that two instances (relative and absolute) behave the same for basic operations
    ekf_abs = UUV_EKF()
    ekf_rel = UUV_EKF_relative()
    # Initialize both
    timestamp = 0.0
    pos = [1.0, 2.0, 3.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.1, 0.2, 0.3]
    ang_vel = [0.01, 0.02, 0.03]
    ekf_abs.initialize_state(timestamp, pos, quat, vel, ang_vel)
    ekf_rel.initialize_state(timestamp, pos, quat, vel, ang_vel)
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)

    # Predict step
    new_timestamp = 0.1
    omega_body = np.array([0.01, 0.02, 0.03])
    ekf_abs.predict(new_timestamp, omega_body)
    ekf_rel.predict(new_timestamp, omega_body)
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)

    # Correction step (position)
    d2_abs, scaled_abs = ekf_abs.correction_position_xy(1.0, 1.5, 2.5)
    d2_rel, scaled_rel = ekf_rel.correction_position_xy(1.0, 1.5, 2.5)
    assert np.isclose(d2_abs, d2_rel)
    assert scaled_abs == scaled_rel
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)


def test_uuv_ekf_predict_and_correction_depth_consistency():
    ekf_abs = UUV_EKF()
    ekf_rel = UUV_EKF_relative()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf_abs.initialize_state(timestamp, pos, quat, vel, ang_vel)
    ekf_rel.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_timestamp = 0.2
    omega_body = np.array([0.0, 0.0, 0.0])
    ekf_abs.predict(new_timestamp, omega_body)
    ekf_rel.predict(new_timestamp, omega_body)
    d2_abs, scaled_abs = ekf_abs.correction_depth(10.0)
    d2_rel, scaled_rel = ekf_rel.correction_depth(10.0)
    assert np.isclose(d2_abs, d2_rel)
    assert scaled_abs == scaled_rel
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)


def test_uuv_ekf_correction_ahrs_consistency():
    ekf_abs = UUV_EKF()
    ekf_rel = UUV_EKF_relative()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf_abs.initialize_state(timestamp, pos, quat, vel, ang_vel)
    ekf_rel.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_quat = [0.0, 0.0, 1.0, 0.0]
    new_quat = np.array(new_quat) / np.linalg.norm(new_quat)
    new_ang_vel = [0.1, 0.2, 0.3]
    ekf_abs.correction_attitude(new_quat)
    ekf_abs.correction_angular_rate(new_ang_vel)
    ekf_rel.correction_attitude(new_quat)
    ekf_rel.correction_angular_rate(new_ang_vel)
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)


def test_uuv_ekf_correction_velocity_consistency():
    ekf_abs = UUV_EKF()
    ekf_rel = UUV_EKF_relative()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf_abs.initialize_state(timestamp, pos, quat, vel, ang_vel)
    ekf_rel.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2_abs, scaled_abs = ekf_abs.correction_velocity([1.0, 0.0, 0.0])
    d2_rel, scaled_rel = ekf_rel.correction_velocity([1.0, 0.0, 0.0])
    assert np.isclose(d2_abs, d2_rel)
    assert scaled_abs == scaled_rel
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)


def test_uuv_ekf_correction_bottom_height_consistency():
    ekf_abs = UUV_EKF(bottom_height=True)
    ekf_rel = UUV_EKF_relative(bottom_height=True)
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf_abs.initialize_state(timestamp, pos, quat, vel, ang_vel, bottom_height=0.0)
    ekf_rel.initialize_state(timestamp, pos, quat, vel, ang_vel, bottom_height=0.0)
    d2_abs, scaled_abs = ekf_abs.correction_bottom_height(10.0)
    d2_rel, scaled_rel = ekf_rel.correction_bottom_height(10.0)
    assert np.isclose(d2_abs, d2_rel)
    assert scaled_abs == scaled_rel
    np.testing.assert_allclose(ekf_abs.current_state, ekf_rel.current_state)
    np.testing.assert_allclose(ekf_abs.current_covariance, ekf_rel.current_covariance)


def test_kf_predict():
    """
    Test the kf_predict function.
    """
    # Test case 1: Valid input
    x = np.array([1, 2])
    P = np.array([[1, 0], [0, 1]])
    A = np.array([[1, 1], [0, 1]])
    Q = np.array([[0.1, 0], [0, 0.1]])
    B = np.array([[0.5], [1]])
    u = np.array([1])
    x_pred, P_pred = kf_predict(x, P, A, Q, B, u)
    assert x_pred.shape == (2,)
    assert P_pred.shape == (2, 2)

    # Test case 2: List input
    x_pred, P_pred = kf_predict(
        x.tolist(), P.tolist(), A.tolist(), Q.tolist(), B.tolist(), u.tolist()
    )
    assert x_pred.shape == (2,)
    assert P_pred.shape == (2, 2)

    # Test case 3: Default B and u
    x_pred, P_pred = kf_predict(x, P, A, Q)
    assert x_pred.shape == (2,)
    assert P_pred.shape == (2, 2)

    # Test case 4: Invalid x type
    try:
        kf_predict("invalid_type", P, A, Q, B, u)
    except TypeError as e:
        assert str(e) == "x must be a numpy array"
    else:
        assert False

    # Test case 5: Invalid P type
    try:
        kf_predict(x, "invalid_type", A, Q, B, u)
    except TypeError as e:
        assert str(e) == "P must be a numpy array"
    else:
        assert False

    # Test case 6: Invalid A type
    try:
        kf_predict(x, P, "invalid_type", Q, B, u)
    except TypeError as e:
        assert str(e) == "A must be a numpy array"
    else:
        assert False

    # Test case 7: Invalid Q type
    try:
        kf_predict(x, P, A, "invalid_type", B, u)
    except TypeError as e:
        assert str(e) == "Q must be a numpy array"
    else:
        assert False

    # Test case 8: Invalid B type
    try:
        kf_predict(x, P, A, Q, "invalid_type", u)
    except TypeError as e:
        assert str(e) == "B must be a numpy array"
    else:
        assert False

    # Test case 9: Invalid u type
    try:
        kf_predict(x, P, A, Q, B, "invalid_type")
    except TypeError as e:
        assert str(e) == "u must be a numpy array"
    else:
        assert False

    # Test case 10: Invalid x shape
    try:
        kf_predict(np.array([1, 2, 3]), P, A, Q, B, u)
    except ValueError as e:
        assert str(e) == "x and A must have the same number of rows"
    else:
        assert False

    # Test case 11: Invalid P shape
    try:
        kf_predict(x, np.array([1, 2, 3]), A, Q, B, u)
    except ValueError as e:
        assert str(e) == "P must be a square matrix"
    else:
        assert False

    # Test case 12: Invalid A shape
    try:
        kf_predict(x, P, np.array([1, 2, 3]), Q, B, u)
    except ValueError as e:
        assert str(e) == "A must be a 2D matrix"
    else:
        assert False

    # Test case 13: Invalid Q shape
    try:
        kf_predict(x, P, A, np.array([1, 2, 3]), B, u)
    except ValueError as e:
        assert str(e) == "Q must be a square matrix"
    else:
        assert False

    # Test case 14: Invalid B shape
    try:
        kf_predict(x, P, A, Q, np.array([1, 2, 3]), u)
    except ValueError as e:
        assert str(e) == "B must be a 2D matrix"
    else:
        assert False

    # Test case 15: Invalid u shape
    try:
        kf_predict(x, P, A, Q, B, np.array([1, 2, 3]))
    except ValueError as e:
        assert (
            str(e)
            == "The number of columns in B must be equal to the number of rows in u"
        )
    else:
        assert False


def test_kf_update():
    """
    Test the kf_update function.
    """
    # Test case 1: Valid input
    x = np.array([1, 2])
    P = np.array([[1, 0], [0, 1]])
    y = np.array([1, 2])
    H = np.array([[1, 0], [0, 1]])
    R = np.array([[0.1, 0], [0, 0.1]])
    x_upd, P_upd, K, dy, S = kf_update(x, P, y, H, R)
    assert x_upd.shape == (2,)
    assert P_upd.shape == (2, 2)
    assert K.shape == (2, 2)
    assert dy.shape == (2,)
    assert S.shape == (2, 2)

    # Test case 2: List input
    x_upd, P_upd, K, dy, S = kf_update(
        x.tolist(), P.tolist(), y.tolist(), H.tolist(), R.tolist()
    )
    assert x_upd.shape == (2,)
    assert P_upd.shape == (2, 2)
    assert K.shape == (2, 2)
    assert dy.shape == (2,)
    assert S.shape == (2, 2)

    # Test case 3: Invalid x type
    try:
        kf_update("invalid_type", P, y, H, R)
    except TypeError as e:
        assert str(e) == "x must be a numpy array"
    else:
        assert False

    # Test case 4: Invalid P type
    try:
        kf_update(x, "invalid_type", y, H, R)
    except TypeError as e:
        assert str(e) == "P must be a numpy array"
    else:
        assert False

    # Test case 5: Invalid y type
    try:
        kf_update(x, P, "invalid_type", H, R)
    except TypeError as e:
        assert str(e) == "y must be a numpy array"
    else:
        assert False

    # Test case 6: Invalid H type
    try:
        kf_update(x, P, y, "invalid_type", R)
    except TypeError as e:
        assert str(e) == "H must be a numpy array"
    else:
        assert False

    # Test case 7: Invalid R type
    try:
        kf_update(x, P, y, H, "invalid_type")
    except TypeError as e:
        assert str(e) == "R must be a numpy array"
    else:
        assert False

    # Test case 10: Invalid P shape
    try:
        kf_update(x, np.array([1, 2, 3]), y, H, R)
    except ValueError as e:
        assert str(e) == "P must be a square matrix"
    else:
        assert False

    # Test case 11: Invalid H shape
    try:
        kf_update(x, P, y, np.array([1, 2, 3]), R)
    except ValueError as e:
        assert str(e) == "H must be a 2D matrix"
    else:
        assert False

    # Test case 12: Invalid R shape
    try:
        kf_update(x, P, y, H, np.array([1, 2, 3]))
    except ValueError as e:
        assert str(e) == "R must be a square matrix"
    else:
        assert False

    # Test case 13: Dimension mismatch between x and P
    try:
        kf_update(np.array([1, 2, 3]), np.array([[1, 0], [0, 1]]), y, H, R)
    except ValueError as e:
        assert str(e) == "P must have the same number of rows as x"
    else:
        assert False


def test_uuv_ekf_init_and_properties():
    # Test default initialization
    ekf = UUV_EKF()
    assert ekf.state_dim == 13
    assert ekf.current_state.shape == (14,)
    assert ekf.current_covariance.shape == (13, 13)
    assert ekf.H_position.shape == (2, 13)
    assert ekf.R_position.shape == (2, 2)
    assert ekf.H_depth.shape == (1, 13)
    assert ekf.R_depth.shape == (1, 1)
    assert ekf.H_attitude.shape == (3, 13)
    assert ekf.R_attitude.shape == (3, 3)
    assert ekf.H_angrate.shape == (3, 13)
    assert ekf.R_angrate.shape == (3, 3)
    assert ekf.H_vel.shape == (3, 13)
    assert ekf.R_vel.shape == (3, 3)
    assert ekf.H_bottom_height.shape == (1, 13)
    assert ekf.R_bottom_height.shape == (1, 1)

    # Test with bottom_height=True
    ekf2 = UUV_EKF(bottom_height=True)
    assert ekf2.state_dim == 14
    assert ekf2.current_state.shape == (15,)
    assert ekf2.current_covariance.shape == (14, 14)
    assert ekf2.H_position.shape == (2, 14)
    assert ekf2.H_depth.shape == (1, 14)
    assert ekf2.H_attitude.shape == (3, 14)
    assert ekf2.H_angrate.shape == (3, 14)
    assert ekf2.H_vel.shape == (3, 14)
    assert ekf2.H_bottom_height.shape == (1, 14)


def test_uuv_ekf_initialize_state_and_setters():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [1.0, 2.0, 3.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.1, 0.2, 0.3]
    ang_vel = [0.01, 0.02, 0.03]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    np.testing.assert_allclose(ekf.current_state[0], timestamp)
    np.testing.assert_allclose(ekf.current_state[1:4], pos)
    np.testing.assert_allclose(ekf.current_state[4:8], quat)
    np.testing.assert_allclose(ekf.current_state[8:11], vel)
    np.testing.assert_allclose(ekf.current_state[11:14], ang_vel)


def test_uuv_ekf_predict():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [1.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_timestamp = 1.0
    omega_body = np.array([0.0, 0.0, 0.0])
    ekf.predict(new_timestamp, omega_body)
    # Position should have advanced by velocity * dt
    np.testing.assert_allclose(
        ekf.current_state[1:4],
        np.array(pos) + np.array(vel) * (new_timestamp - timestamp),
    )
    # Quaternion should remain unchanged for zero angular velocity
    np.testing.assert_allclose(ekf.current_state[4:8], quat, atol=1e-6)


def test_uuv_ekf_correction_position_xy():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2, scaled = ekf.correction_position_xy(1.0, 2.0)
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, position should move closer to measurement
    assert np.abs(ekf.current_state[1] - 1.0) < 1.0
    assert np.abs(ekf.current_state[2] - 2.0) < 2.0

    # Custom covariance for position
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2, scaled = ekf.correction_position_xy(
        1.0, 2.0, measurement_covariance=np.eye(2) * 10.0
    )
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, position should move closer to measurement
    assert np.abs(ekf.current_state[1] - 1.0) < 1.0
    assert np.abs(ekf.current_state[2] - 2.0) < 2.0


def test_uuv_ekf_correction_position_xyz():
    ekf = UUV_EKF(covariance_depth=10.0)
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2, scaled = ekf.correction_position_xyz(1.0, 2.0, 1.0)
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, position should move closer to measurement
    assert np.abs(ekf.current_state[1] - 1.0) < 1.0
    assert np.abs(ekf.current_state[2] - 2.0) < 2.0
    assert np.abs(ekf.current_state[3] - 1.0) < 1.0

    # Custom covariance for position
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2, scaled = ekf.correction_position_xyz(
        1.0, 2.0, 1.0, measurement_covariance=np.eye(3) * 10.0
    )
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, position should move closer to measurement
    assert np.abs(ekf.current_state[1] - 1.0) < 1.0
    assert np.abs(ekf.current_state[2] - 2.0) < 2.0
    assert np.abs(ekf.current_state[3] - 1.0) < 1.0


def test_uuv_ekf_correction_depth():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    timestamp = 1.0
    d2, scaled = ekf.correction_depth(5.0)
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, depth should move closer to measurement
    assert np.abs(ekf.current_state[3] - 5.0) < 5.0

    # Custom covariance for depth
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    timestamp = 1.0
    d2, scaled = ekf.correction_depth(5.0, measurement_covariance=np.eye(1) * 0.001)
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, depth should move closer to measurement
    assert np.abs(ekf.current_state[3] - 5.0) < 5.0


def test_uuv_ekf_correction_attitude():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_quat = [0.0, 0.0, 1.0, 0.0]
    new_quat = np.array(new_quat) / np.linalg.norm(new_quat)
    ekf.correction_attitude(new_quat)
    # Quaternion should be normalized
    np.testing.assert_allclose(np.linalg.norm(ekf.current_state[4:8]), 1.0, atol=1e-6)

    # Custom covariance for attitude
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_quat = [0.0, 0.0, 1.0, 0.0]
    new_quat = np.array(new_quat) / np.linalg.norm(new_quat)
    ekf.correction_attitude(new_quat, measurement_covariance=np.eye(3) * 0.001)
    # Quaternion should be normalized
    np.testing.assert_allclose(np.linalg.norm(ekf.current_state[4:8]), 1.0, atol=1e-6)


def test_uuv_ekf_correction_angular_rate():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_ang_vel = [0.1, 0.2, 0.3]
    ekf.correction_angular_rate(new_ang_vel)
    # Quaternion should be normalized
    np.testing.assert_allclose(np.linalg.norm(ekf.current_state[4:8]), 1.0, atol=1e-6)

    # Custom covariance for angular rate
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    new_ang_vel = [0.1, 0.2, 0.3]
    ekf.correction_angular_rate(new_ang_vel, measurement_covariance=np.eye(3) * 0.001)
    # Quaternion should be normalized
    np.testing.assert_allclose(np.linalg.norm(ekf.current_state[4:8]), 1.0, atol=1e-6)


def test_uuv_ekf_correction_velocity():
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2, scaled = ekf.correction_velocity([1.0, 0.0, 0.0])
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, velocity should move closer to measurement
    assert np.abs(ekf.current_state[8] - 1.0) < 1.0

    # Custom covariance for velocity
    ekf = UUV_EKF()
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel)
    d2, scaled = ekf.correction_velocity(
        [1.0, 0.0, 0.0], measurement_covariance=np.eye(3) * 0.001
    )
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, velocity should move closer to measurement
    assert np.abs(ekf.current_state[8] - 1.0) < 1.0


def test_uuv_ekf_correction_bottom_height():
    ekf = UUV_EKF(bottom_height=True)
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel, bottom_height=0.0)
    d2, scaled = ekf.correction_bottom_height(10.0)
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, bottom height should move closer to measurement
    assert np.abs(ekf.current_state[-1] - 10.0) < 10.0

    # Custom covariance for bottom height
    ekf = UUV_EKF(bottom_height=True)
    timestamp = 0.0
    pos = [0.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    vel = [0.0, 0.0, 0.0]
    ang_vel = [0.0, 0.0, 0.0]
    ekf.initialize_state(timestamp, pos, quat, vel, ang_vel, bottom_height=0.0)
    d2, scaled = ekf.correction_bottom_height(
        10.0, measurement_covariance=np.eye(1) * 0.001
    )
    assert isinstance(d2, float)
    assert isinstance(scaled, bool)
    # After correction, bottom height should move closer to measurement
    assert np.abs(ekf.current_state[-1] - 10.0) < 10.0
