"""
Test module for the DVL calibration module.

Authors: Sebastián Rodríguez-Martínez
Contact: srodriguez@mbari.org
Last Update: 2025-04-15
"""

import numpy as np
import pytest

from navlib import cal_dvl_acc_so3, cal_dvl_correct_by_sound_velocity, cal_dvl_jck_so3
from navlib.cal import cal_dvl_acc_so3 as cal_dvl_acc_so3_relative
from navlib.cal import (
    cal_dvl_correct_by_sound_velocity as cal_dvl_correct_by_sound_velocity_relative,
)
from navlib.cal import cal_dvl_jck_so3 as cal_dvl_jck_so3_relative


def test_import():
    # Test if both import methods work
    assert cal_dvl_acc_so3_relative.__code__.co_code == cal_dvl_acc_so3.__code__.co_code
    assert (
        cal_dvl_correct_by_sound_velocity_relative.__code__.co_code
        == cal_dvl_correct_by_sound_velocity.__code__.co_code
    )
    assert cal_dvl_jck_so3_relative.__code__.co_code == cal_dvl_jck_so3.__code__.co_code


def dummy_inputs():
    # Create valid dummy inputs
    t = np.linspace(0, 1, 10)
    valid_time_dvl = t.copy()
    valid_time_ahrs = t.copy()
    valid_time_depth = t.copy()
    valid_velocity_dvl = np.ones((10, 3))
    valid_acceleration_ahrs = np.ones((10, 3))
    valid_rph_ahrs = np.ones((10, 3))
    valid_angular_rate_ahrs = np.ones((10, 3))
    valid_depth = t.copy()
    valid_lever_arm = np.array([0, 0, 0])

    return {
        "time_dvl": valid_time_dvl,
        "time_ahrs": valid_time_ahrs,
        "time_depth": valid_time_depth,
        "velocity_dvl": valid_velocity_dvl,
        "acceleration_ahrs": valid_acceleration_ahrs,
        "rph_ahrs": valid_rph_ahrs,
        "angular_rate_ahrs": valid_angular_rate_ahrs,
        "depth": valid_depth,
        "low_pass_filter": True,
        "remove_gravity": True,
        "lever_arm": valid_lever_arm,
        "method": "ls",
        "force_so3": True,
    }


def test_invalid_time_ahrs_type():
    params = dummy_inputs()
    params["time_ahrs"] = "not a numpy array"
    with pytest.raises(TypeError):
        cal_dvl_acc_so3(**params)


def test_invalid_low_pass_filter_type():
    params = dummy_inputs()
    params["low_pass_filter"] = "True"
    with pytest.raises(TypeError):
        cal_dvl_acc_so3(**params)


def test_invalid_time_ahrs_dimensions():
    params = dummy_inputs()
    # Provide a wrong dimension: 2-dimensional array that is not (n,) or (n,1) or (1,n)
    params["time_ahrs"] = np.ones((5, 5))
    with pytest.raises(ValueError):
        cal_dvl_acc_so3(**params)


def test_velocity_dvl_wrong_shape():
    params = dummy_inputs()
    # Provide velocity_dvl with wrong shape; should be 10x3 or 3x10 but here use 10x2.
    params["velocity_dvl"] = np.ones((10, 2))
    with pytest.raises(ValueError):
        cal_dvl_acc_so3(**params)


def test_mismatched_time_and_rph():
    params = dummy_inputs()
    # Make time_ahrs length different from rph_ahrs rows.
    params["time_ahrs"] = np.linspace(0, 1, 8)
    with pytest.raises(ValueError):
        cal_dvl_acc_so3(**params)


def test_mismatched_time_and_velocity():
    params = dummy_inputs()
    # Make time_dvl length different from velocity_dvl rows.
    params["time_dvl"] = np.linspace(0, 1, 12)
    with pytest.raises(ValueError):
        cal_dvl_acc_so3(**params)


def test_invalid_lever_arm_shape():
    params = dummy_inputs()
    # Provide lever_arm with wrong dimension; must have 3 elements.
    params["lever_arm"] = np.array([0, 0])
    with pytest.raises(ValueError):
        cal_dvl_acc_so3(**params)


def get_valid_inputs():
    # velocity_dvl: valid 10x3 array
    velocity_dvl = np.ones((10, 3))
    # conductivity, temperature, depth, latitude, longitude: matching 1D arrays of length 10
    conductivity = np.ones(10)
    temperature = np.ones(10) * 10  # arbitrary valid values
    depth = -100 * np.ones(10)  # negative depth values are required
    latitude = np.zeros(10)
    longitude = np.zeros(10)
    return velocity_dvl, conductivity, temperature, depth, latitude, longitude


def test_invalid_velocity_dvl_type():
    # velocity_dvl must be a list or numpy array; passing an int should trigger TypeError.
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    with pytest.raises(TypeError):
        cal_dvl_correct_by_sound_velocity(
            123, conductivity, temperature, depth, latitude, longitude
        )


def test_invalid_conductivity_type():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    with pytest.raises(TypeError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, "not an array", temperature, depth, latitude, longitude
        )


def test_invalid_temperature_type():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    with pytest.raises(TypeError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, conductivity, {"temp": 10}, depth, latitude, longitude
        )


def test_invalid_depth_type():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    with pytest.raises(TypeError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, conductivity, temperature, "depth", latitude, longitude
        )


def test_invalid_latitude_type():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    # latitude must be a numpy array or a float; passing a dict will raise.
    with pytest.raises(TypeError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl,
            conductivity,
            temperature,
            depth,
            latitude={"lat": 0},
            longitude=longitude,
        )


def test_invalid_longitude_type():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    with pytest.raises(TypeError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl,
            conductivity,
            temperature,
            depth,
            latitude,
            longitude={"lon": 0},
        )


def test_velocity_shape_error():
    # velocity_dvl must be 2D and have one dimension size equal to 3.
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    bad_velocity = np.ones((5, 5))  # Both dims are not 3.
    with pytest.raises(ValueError):
        cal_dvl_correct_by_sound_velocity(
            bad_velocity, conductivity, temperature, depth, latitude, longitude
        )


def test_shape_mismatch_error():
    # Mismatch in number of measurements between velocity_dvl and other parameters.
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    bad_conductivity = np.ones(5)  # Length does not match velocity_dvl (10)
    with pytest.raises(ValueError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, bad_conductivity, temperature, depth, latitude, longitude
        )


def test_positive_depth_error():
    # Depth must be negative.
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    bad_depth = np.abs(depth)  # Make them positive
    with pytest.raises(ValueError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, conductivity, temperature, bad_depth, latitude, longitude
        )


def test_invalid_latitude_range():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    # Set latitude out of valid range (< -90 or > 90)
    bad_latitude = -100 * np.ones_like(latitude)
    with pytest.raises(ValueError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, conductivity, temperature, depth, bad_latitude, longitude
        )


def test_invalid_longitude_range():
    (
        velocity_dvl,
        conductivity,
        temperature,
        depth,
        latitude,
        longitude,
    ) = get_valid_inputs()
    bad_longitude = 400 * np.ones_like(longitude)  # Out of valid range (-360,360)
    with pytest.raises(ValueError):
        cal_dvl_correct_by_sound_velocity(
            velocity_dvl, conductivity, temperature, depth, latitude, bad_longitude
        )


def get_valid_inputs_jck():
    """
    Generate a valid set of inputs for cal_dvl_jck.
    All time signals are 1D arrays of equal length. The velocity, rph, angular_rate
    and position arrays are shaped (n,3).
    """
    n = 10
    time_dvl = np.linspace(0, 9, n)
    time_ahrs = np.linspace(0, 9, n)
    time_position = np.linspace(0, 9, n)
    velocity_dvl = np.ones((n, 3))
    rph_ahrs = np.zeros((n, 3))
    angular_rate_ahrs = np.zeros((n, 3))
    position = 10 * np.ones((n, 3))
    return (
        time_dvl,
        time_ahrs,
        time_position,
        velocity_dvl,
        rph_ahrs,
        angular_rate_ahrs,
        position,
    )


def test_valid_ls_solution():
    """Test that a valid input returns a (3,3) matrix (LS method)."""
    inputs = get_valid_inputs_jck()
    R = cal_dvl_jck_so3(
        *inputs,
        low_pass_filter=False,
        method="ls",
        force_so3=False,
        zero_mean_sensor=False,
        zero_mean_equation=False,
    )
    assert R.shape == (3, 3)


def test_valid_svd_solution():
    """Test that a valid input returns a (3,3) matrix (SVD method)."""
    inputs = get_valid_inputs_jck()
    R = cal_dvl_jck_so3(
        *inputs,
        low_pass_filter=False,
        method="svd",
        force_so3=False,
        zero_mean_sensor=False,
        zero_mean_equation=False,
    )
    assert R.shape == (3, 3)


def test_invalid_time_dvl_type():
    """Test that a wrong type for time_dvl raises a TypeError."""
    inputs = list(get_valid_inputs_jck())
    inputs[0] = "not an array"  # time_dvl must be an array
    with pytest.raises(TypeError):
        cal_dvl_jck_so3(*inputs)


def test_invalid_velocity_shape():
    """Test that an invalid shape for velocity_dvl raises a ValueError."""
    inputs = list(get_valid_inputs_jck())
    # Force an incorrect shape (e.g. 5x5 instead of 3xN or Nx3)
    inputs[3] = np.ones((5, 5))
    with pytest.raises(ValueError):
        cal_dvl_jck_so3(*inputs)


def test_invalid_rph_ahrs_type():
    """Test that a wrong type for rph_ahrs raises a TypeError."""
    inputs = list(get_valid_inputs_jck())
    inputs[4] = "invalid"
    with pytest.raises(TypeError):
        cal_dvl_jck_so3(*inputs)


def test_invalid_position_shape():
    """Test that an invalid shape for position raises a ValueError."""
    inputs = list(get_valid_inputs_jck())
    # Provide a wrong shape (e.g. 10x10 instead of (n,3))
    inputs[6] = np.ones((10, 10))
    with pytest.raises(ValueError):
        cal_dvl_jck_so3(*inputs)
