"""
Test module for the vector math module.

Authors: Sebastián Rodríguez-Martínez
Contact: srodriguez@mbari.org
Last Update: 2024-03-25
"""

import numpy as np

from navlib import (
    derivative,
    difference,
    distance_traveled,
    max,
    mean,
    median,
    min,
    norm,
    normalize,
    print_stats,
    remove_mean,
    remove_median,
    remove_offset,
    resample,
    saturate,
    std,
    transpose,
    unwrap1pi,
    unwrap2pi,
    unwrap180,
    unwrap360,
    wrap1pi,
    wrap2pi,
    wrap180,
    wrap360,
    wrapunwrap,
    wrapunwrap360,
)
from navlib.math import derivative as derivative_relative
from navlib.math import difference as difference_relative
from navlib.math import distance_traveled as distance_traveled_relative
from navlib.math import max as max_relative
from navlib.math import mean as mean_relative
from navlib.math import median as median_relative
from navlib.math import min as min_relative
from navlib.math import norm as norm_relative
from navlib.math import normalize as normalize_relative
from navlib.math import print_stats as print_stats_relative
from navlib.math import remove_mean as remove_mean_relative
from navlib.math import remove_median as remove_median_relative
from navlib.math import remove_offset as remove_offset_relative
from navlib.math import resample as resample_relative
from navlib.math import saturate as saturate_relative
from navlib.math import std as std_relative
from navlib.math import transpose as transpose_relative
from navlib.math import unwrap1pi as unwrap1pi_relative
from navlib.math import unwrap2pi as unwrap2pi_relative
from navlib.math import unwrap180 as unwrap180_relative
from navlib.math import unwrap360 as unwrap360_relative
from navlib.math import wrap1pi as wrap1pi_relative
from navlib.math import wrap2pi as wrap2pi_relative
from navlib.math import wrap180 as wrap180_relative
from navlib.math import wrap360 as wrap360_relative
from navlib.math import wrapunwrap as wrapunwrap_relative
from navlib.math import wrapunwrap360 as wrapunwrap360_relative


def test_imports():
    # Checks that the absolute and relative imports work correctly
    assert derivative_relative.__code__.co_code == derivative.__code__.co_code
    assert difference_relative.__code__.co_code == difference.__code__.co_code
    assert (
        distance_traveled_relative.__code__.co_code
        == distance_traveled.__code__.co_code
    )
    assert max_relative.__code__.co_code == max.__code__.co_code
    assert mean_relative.__code__.co_code == mean.__code__.co_code
    assert median_relative.__code__.co_code == median.__code__.co_code
    assert min_relative.__code__.co_code == min.__code__.co_code
    assert norm_relative.__code__.co_code == norm.__code__.co_code
    assert normalize_relative.__code__.co_code == normalize.__code__.co_code
    assert print_stats_relative.__code__.co_code == print_stats.__code__.co_code
    assert remove_mean_relative.__code__.co_code == remove_mean.__code__.co_code
    assert remove_median_relative.__code__.co_code == remove_median.__code__.co_code
    assert remove_offset_relative.__code__.co_code == remove_offset.__code__.co_code
    assert resample_relative.__code__.co_code == resample.__code__.co_code
    assert saturate_relative.__code__.co_code == saturate.__code__.co_code
    assert std_relative.__code__.co_code == std.__code__.co_code
    assert transpose_relative.__code__.co_code == transpose.__code__.co_code
    assert unwrap1pi_relative.__code__.co_code == unwrap1pi.__code__.co_code
    assert unwrap2pi_relative.__code__.co_code == unwrap2pi.__code__.co_code
    assert unwrap180_relative.__code__.co_code == unwrap180.__code__.co_code
    assert unwrap360_relative.__code__.co_code == unwrap360.__code__.co_code
    assert wrap1pi_relative.__code__.co_code == wrap1pi.__code__.co_code
    assert wrap2pi_relative.__code__.co_code == wrap2pi.__code__.co_code
    assert wrap180_relative.__code__.co_code == wrap180.__code__.co_code
    assert wrap360_relative.__code__.co_code == wrap360.__code__.co_code
    assert wrapunwrap_relative.__code__.co_code == wrapunwrap.__code__.co_code
    assert wrapunwrap360_relative.__code__.co_code == wrapunwrap360.__code__.co_code


def test_norm():
    """
    Test the norm function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_norm = 3.7416573867739413
    assert norm(vec) - expected_norm < 1e-6

    # Test case 2: 2D vector with keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_norm = np.array([[1.73205081], [3.46410162], [5.19615242]])
    assert np.allclose(norm(vec, keepdims=True), expected_norm)

    # Test case 3: 2D vector without keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_norm = np.array([1.73205081, 3.46410162, 5.19615242])
    assert np.allclose(norm(vec, keepdims=False), expected_norm)

    # Test case 4: List input
    vec = [1, 2, 3]
    expected_norm = 3.7416573867739413
    assert norm(vec) - expected_norm < 1e-6

    # Test case 5: Wrong input type
    vec = "1, 2, 3"
    try:
        norm(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_normalize():
    """
    Test the normalize function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_normalized = np.array([0.26726124, 0.53452248, 0.80178373])
    assert np.allclose(normalize(vec), expected_normalized)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_normalized = np.array(
        [
            [0.57735027, 0.57735027, 0.57735027],
            [0.57735027, 0.57735027, 0.57735027],
            [0.57735027, 0.57735027, 0.57735027],
        ]
    )
    assert np.allclose(normalize(vec), expected_normalized)

    # Test case 3: List input
    vec = [1, 2, 3]
    expected_normalized = np.array([0.26726124, 0.53452248, 0.80178373])
    assert np.allclose(normalize(vec), expected_normalized)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    try:
        normalize(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_remove_offset():
    """
    Test the remove_offset function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    offset = 1
    expected_vec = np.array([0, 1, 2])
    assert np.allclose(remove_offset(vec, offset), expected_vec)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    offset = 1
    expected_vec = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])
    assert np.allclose(remove_offset(vec, offset), expected_vec)

    # Test case 3: List input
    vec = [1, 2, 3]
    offset = 1
    expected_vec = np.array([0, 1, 2])
    assert np.allclose(remove_offset(vec, offset), expected_vec)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    offset = 1
    try:
        remove_offset(vec, offset)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_remove_mean():
    """
    Test the remove_mean function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_vec = np.array([-1, 0, 1])
    assert np.allclose(remove_mean(vec), expected_vec)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_vec = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
    assert np.allclose(remove_mean(vec), expected_vec)

    # Test case 3: List input
    vec = [1, 2, 3]
    expected_vec = np.array([-1, 0, 1])
    assert np.allclose(remove_mean(vec), expected_vec)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    try:
        remove_mean(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_remove_median():
    """
    Test the remove_median function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_vec = np.array([-1, 0, 1])
    assert np.allclose(remove_median(vec), expected_vec)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_vec = np.array([[-1, -1, -1], [0, 0, 0], [1, 1, 1]])
    assert np.allclose(remove_median(vec), expected_vec)

    # Test case 3: List input
    vec = [1, 2, 3]
    expected_vec = np.array([-1, 0, 1])
    assert np.allclose(remove_median(vec), expected_vec)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    try:
        remove_median(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_mean():
    """
    Test the mean function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_mean = 2
    assert mean(vec) == expected_mean

    # Test case 2: 2D vector withou keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_mean = np.array([2, 2, 2])
    assert np.allclose(mean(vec), expected_mean)

    # Test case 3: 2D vector with keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_mean = np.array([[2, 2, 2]])
    assert np.allclose(mean(vec, keepdims=True), expected_mean)

    # Test case 4: List input
    vec = [1, 2, 3]
    expected_mean = 2
    assert mean(vec) == expected_mean

    # Test case 5: Wrong input type
    vec = "1, 2, 3"
    try:
        mean(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_median():
    """
    Test the median function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_median = 2
    assert median(vec) == expected_median

    # Test case 2: 2D vector withou keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_median = np.array([2, 2, 2])
    assert np.allclose(median(vec), expected_median)

    # Test case 3: 2D vector with keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_median = np.array([[2, 2, 2]])
    assert np.allclose(median(vec, keepdims=True), expected_median)

    # Test case 4: List input
    vec = [1, 2, 3]
    expected_median = 2
    assert median(vec) == expected_median

    # Test case 5: Wrong input type
    vec = "1, 2, 3"
    try:
        median(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_std():
    """
    Test the std function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_std = 0.816496580927726
    assert std(vec) - expected_std < 1e-6

    # Test case 2: 2D vector without keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_std = np.array([0.816496580927726, 0.816496580927726, 0.816496580927726])
    assert np.allclose(std(vec), expected_std)

    # Test case 3: 2D vector with keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_std = np.array([[0.816496580927726, 0.816496580927726, 0.816496580927726]])
    assert np.allclose(std(vec, keepdims=True), expected_std)

    # Test case 4: List input
    vec = [1, 2, 3]
    expected_std = 0.816496580927726
    assert std(vec) - expected_std < 1e-6

    # Test case 5: Wrong input type
    vec = "1, 2, 3"
    try:
        std(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_min():
    """
    Test the min function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_min = 1
    assert min(vec) == expected_min

    # Test case 2: 2D vector withou keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_min = np.array([1, 1, 1])
    assert np.allclose(min(vec), expected_min)

    # Test case 3: 2D vector with keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_min = np.array([[1, 1, 1]])
    assert np.allclose(min(vec, keepdims=True), expected_min)

    # Test case 4: List input
    vec = [1, 2, 3]
    expected_min = 1
    assert min(vec) == expected_min

    # Test case 5: Wrong input type
    vec = "1, 2, 3"
    try:
        min(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_max():
    """
    Test the max function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_max = 3
    assert max(vec) == expected_max

    # Test case 2: 2D vector withou keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_max = np.array([3, 3, 3])
    assert np.allclose(max(vec), expected_max)

    # Test case 3: 2D vector with keepdims
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_max = np.array([[3, 3, 3]])
    assert np.allclose(max(vec, keepdims=True), expected_max)

    # Test case 4: List input
    vec = [1, 2, 3]
    expected_max = 3
    assert max(vec) == expected_max

    # Test case 5: Wrong input type
    vec = "1, 2, 3"
    try:
        max(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_print_stats():
    """
    Test the print_stats function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    print_stats(vec)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    print_stats(vec)

    # Test case 3: List input
    vec = [1, 2, 3]
    print_stats(vec)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    try:
        print_stats(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_difference():
    """
    Test the difference function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    expected_diff = np.array([1, 1])
    assert np.allclose(difference(vec), expected_diff)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    expected_diff = np.array([[1, 1, 1], [1, 1, 1]])
    assert np.allclose(difference(vec), expected_diff)

    # Test case 3: List input
    vec = [1, 2, 3]
    expected_diff = np.array([1, 1])
    assert np.allclose(difference(vec), expected_diff)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    try:
        difference(vec)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_derivative():
    """
    Test the derivative function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    time = np.array([0.1, 0.2, 0.3])
    expected_deriv = np.array([10, 10])
    assert np.allclose(derivative(vec, time), expected_deriv)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    time = np.array([0.1, 0.2, 0.3])
    expected_deriv = np.array([[10, 10, 10], [10, 10, 10]])
    assert np.allclose(derivative(vec, time), expected_deriv)

    # Test case 3: List input
    vec = [1, 2, 3]
    time = [0.1, 0.2, 0.3]
    expected_deriv = np.array([10, 10])
    assert np.allclose(derivative(vec, time), expected_deriv)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    time = [0.1, 0.2, 0.3]
    try:
        derivative(vec, time)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_resample():
    """
    Test the resample function.
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3, 4])
    time = np.array([0.1, 0.2, 0.3, 0.4])
    new_time = np.array([0.15, 0.25, 0.35])
    expected_resampled = np.array([1.5, 2.5, 3.5])
    assert np.allclose(resample(new_time, time, vec), expected_resampled)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]])
    time = np.array([0.1, 0.2, 0.3, 0.4])
    new_time = np.array([0.15, 0.25, 0.35])
    expected_resampled = np.array([[1.5, 1.5, 1.5], [2.5, 2.5, 2.5], [3.5, 3.5, 3.5]])
    assert np.allclose(resample(new_time, time, vec), expected_resampled)

    # Test case 3: List input
    vec = [1, 2, 3, 4]
    time = [0.1, 0.2, 0.3, 0.4]
    new_time = [0.15, 0.25, 0.35]
    expected_resampled = np.array([1.5, 2.5, 3.5])
    assert np.allclose(resample(new_time, time, vec), expected_resampled)

    # Test case 4: Wrong input type
    vec_str = "1, 2, 3, 4"
    vec = [1, 2, 3, 4]
    time_str = "0.1, 0.2, 0.3, 0.4"
    time_array = np.array([[0.1, 0.2, 0.3, 0.4], [0.1, 0.2, 0.3, 0.4]])
    time_longer = [0.1, 0.2, 0.3, 0.4, 0.5]
    time_non_monotonic = [0.1, 0.2, 0.18, 0.4]
    time = [0.1, 0.2, 0.3, 0.4]
    new_time_str = "0.15, 0.25, 0.35"
    new_time = [0.15, 0.25, 0.35]
    try:
        resample(new_time, time, vec_str)
    except TypeError as e:
        assert str(e) == "Input fp must be a numpy array or a list."
    else:
        assert False
    try:
        resample(new_time_str, time, vec)
    except TypeError as e:
        assert str(e) == "Input x must be a numpy array or a list."
    else:
        assert False
    try:
        resample(new_time, time_str, vec)
    except TypeError as e:
        assert str(e) == "Input xp must be a numpy array or a list."
    else:
        assert False
    try:
        resample(new_time, time_array, vec)
    except ValueError as e:
        assert str(e) == "Input xp must be a 1-D sequence of n floats."
    else:
        assert False
    try:
        resample(new_time, time_longer, vec)
    except ValueError as e:
        assert (
            str(e)
            == "The number of samples in the old time vector and the signal must match."
        )
    else:
        assert False
    try:
        resample(new_time, time_non_monotonic, vec)
    except ValueError as e:
        assert str(e) == "Time series must be monotonically increasing."
    else:
        assert False


def test_distance_traveled():
    """
    Test the distance_traveled function
    """
    # Test case 1: Linear and Full 2D distance
    mat = np.array([[0, 0], [1, 1], [2, 2]])
    expected = np.linalg.norm(mat, axis=1)
    result = distance_traveled(mat, "2d", True, True)
    np.testing.assert_array_equal(result, expected)

    # Test case 2: Linear and non-full 2D distance
    mat = np.array([[0, 0], [1, 1], [2, 2]])
    expected = 2 * np.sqrt(2)
    result = distance_traveled(mat, "2d", True, False)
    assert result == expected

    # Test case 3: non-linear and full 2D distance
    mat = np.array([[0, 0], [1, 1], [2, 2]])
    expected = np.r_[0.0, np.cumsum(np.linalg.norm(np.diff(mat, axis=0), axis=1))]
    result = distance_traveled(mat, "2d", False, True)
    np.testing.assert_array_equal(result, expected)

    # Test case 4: non-linear and non-full 2D distance
    mat = np.array([[0, 0], [1, 1], [2, 2]])
    expected = 2 * np.sqrt(2)
    result = distance_traveled(mat, "2d", False, False)
    assert result == expected

    # Test case 5: Linear and Full 3D distance
    mat = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])
    expected = np.linalg.norm(mat, axis=1)
    result = distance_traveled(mat, "3d", True, True)
    np.testing.assert_array_equal(result, expected)

    # Test case 6: Linear and non-full 3D distance
    mat = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])
    expected = 2 * np.sqrt(3)
    result = distance_traveled(mat, "3d", True, False)
    assert result == expected

    # Test case 7: non-linear and full 3D distance
    mat = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])
    expected = np.r_[0.0, np.cumsum(np.linalg.norm(np.diff(mat, axis=0), axis=1))]
    result = distance_traveled(mat, "3d", False, True)
    np.testing.assert_array_equal(result, expected)

    # Test case 8: non-linear and non-full 3D distance
    mat = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])
    expected = 2 * np.sqrt(3)
    result = distance_traveled(mat, "3d", False, False)
    assert result == expected

    # Test case 9: Incorrect dimensions
    mat = np.array([[0, 0], [1, 1], [2, 2]])
    try:
        distance_traveled(mat, "4d")
    except ValueError as e:
        assert str(e) == "Dimensions must be either 2d or 3d."
    else:
        assert False

    # Test case 10: Insufficient columns
    mat = np.array([[0], [1], [2]])
    try:
        distance_traveled(mat, "2d")
    except ValueError as e:
        assert str(e) == "The input matrix must have at least two columns."
    else:
        assert False


def test_saturate():
    """
    Test the saturate function
    """
    # Test case 1: 1D vector
    vec = np.array([1, 2, 3])
    lower = 1.5
    upper = 2.5
    expected_saturated = np.array([1.5, 2, 2.5])
    assert np.allclose(saturate(vec, lower, upper), expected_saturated)

    # Test case 2: 2D vector
    vec = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    lower = 1.5
    upper = 2.5
    expected_saturated = np.array([[1.5, 1.5, 1.5], [2, 2, 2], [2.5, 2.5, 2.5]])
    assert np.allclose(saturate(vec, lower, upper), expected_saturated)

    # Test case 3: List input
    vec = [1, 2, 3]
    lower = 1.5
    upper = 2.5
    expected_saturated = np.array([1.5, 2, 2.5])
    assert np.allclose(saturate(vec, lower, upper), expected_saturated)

    # Test case 4: Wrong input type
    vec = "1, 2, 3"
    lower = 1.5
    upper = 2.5
    try:
        saturate(vec, lower, upper)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False


def test_wrap360():
    """
    Test the wrap360 function.
    """
    # Test case 1: Unidimensional array
    angles = np.array(
        [
            -400,
            -180,
            -120,
            -90,
            -60,
            -45,
            -30,
            -10,
            0,
            10,
            30,
            45,
            60,
            90,
            120,
            180,
            400,
        ]
    )
    expected_angles = np.array(
        [320, 180, 240, 270, 300, 315, 330, 350, 0, 10, 30, 45, 60, 90, 120, 180, 40]
    )
    assert np.allclose(wrap360(angles), expected_angles)
    assert np.allclose(wrap360(angles.tolist()), expected_angles)

    # Test case 2: Multidimensional array
    angles = np.array(
        [
            [
                -400,
                -180,
                -120,
                -90,
                -60,
                -45,
                -30,
                -10,
                0,
                10,
                30,
                45,
                60,
                90,
                120,
                180,
                400,
            ],
            [
                -400,
                -180,
                -120,
                -90,
                -60,
                -45,
                -30,
                -10,
                0,
                10,
                30,
                45,
                60,
                90,
                120,
                180,
                400,
            ],
        ]
    ).T
    expected_angles = np.array(
        [
            [
                320,
                180,
                240,
                270,
                300,
                315,
                330,
                350,
                0,
                10,
                30,
                45,
                60,
                90,
                120,
                180,
                40,
            ],
            [
                320,
                180,
                240,
                270,
                300,
                315,
                330,
                350,
                0,
                10,
                30,
                45,
                60,
                90,
                120,
                180,
                40,
            ],
        ]
    ).T
    assert np.allclose(wrap360(angles), expected_angles)
    assert np.allclose(wrap360(angles.tolist()), expected_angles)
    assert np.allclose(wrap360(angles.T), expected_angles.T)
    assert np.allclose(wrap360(angles.T.tolist()), expected_angles.T)


def test_wrap180():
    """
    Test the wrap180 function.
    """
    # Test case 1: Unidimensional array
    angles = np.array([-360, -270, -180, -90, 0, 90, 180, 270, 360])
    expected_angles = np.array([0, 90, -180, -90, 0, 90, -180, -90, 0])
    assert np.allclose(wrap180(angles), expected_angles)
    assert np.allclose(wrap180(angles.tolist()), expected_angles)

    # Test case 2: Multidimensional array
    angles = np.array(
        [
            [-360, -270, -180, -90, 0, 90, 180, 270, 360],
            [-360, -270, -180, -90, 0, 90, 180, 270, 360],
        ]
    ).T
    expected_angles = np.array(
        [
            [0, 90, -180, -90, 0, 90, -180, -90, 0],
            [0, 90, -180, -90, 0, 90, -180, -90, 0],
        ]
    ).T
    assert np.allclose(wrap180(angles), expected_angles)
    assert np.allclose(wrap180(angles.tolist()), expected_angles)
    assert np.allclose(wrap180(angles.T), expected_angles.T)
    assert np.allclose(wrap180(angles.T.tolist()), expected_angles.T)


def test_wrap2pi():
    """
    Test the wrap2pi function.
    """
    # Test case 1: Unidimensional array
    angles = np.deg2rad(
        np.array(
            [
                -400,
                -180,
                -120,
                -90,
                -60,
                -45,
                -30,
                -10,
                0,
                10,
                30,
                45,
                90,
                120,
                180,
                400,
            ]
        )
    )
    expected_angles = np.deg2rad(
        np.array(
            [320, 180, 240, 270, 300, 315, 330, 350, 0, 10, 30, 45, 90, 120, 180, 40]
        )
    )
    assert np.allclose(wrap2pi(angles), expected_angles)
    assert np.allclose(wrap2pi(angles.tolist()), expected_angles)

    # Test case 2: Multidimensional array
    angles = np.deg2rad(
        np.array(
            [
                [
                    -400,
                    -180,
                    -120,
                    -90,
                    -60,
                    -45,
                    -30,
                    -10,
                    0,
                    10,
                    30,
                    45,
                    90,
                    120,
                    180,
                ],
                [
                    -400,
                    -180,
                    -120,
                    -90,
                    -60,
                    -45,
                    -30,
                    -10,
                    0,
                    10,
                    30,
                    45,
                    90,
                    120,
                    180,
                ],
            ]
        ).T
    )
    expected_angles = np.deg2rad(
        np.array(
            [
                [320, 180, 240, 270, 300, 315, 330, 350, 0, 10, 30, 45, 90, 120, 180],
                [320, 180, 240, 270, 300, 315, 330, 350, 0, 10, 30, 45, 90, 120, 180],
            ]
        ).T
    )
    assert np.allclose(wrap2pi(angles), expected_angles)
    assert np.allclose(wrap2pi(angles.tolist()), expected_angles)
    assert np.allclose(wrap2pi(angles.T), expected_angles.T)
    assert np.allclose(wrap2pi(angles.T.tolist()), expected_angles.T)


def test_wrap1pi():
    """
    Test the wrap1pi function.
    """
    # Test case 1: Unidimensional array
    angles = np.deg2rad(np.array([-360, -270, -180, -90, 0, 90, 180, 270, 360]))
    expected_angles = np.deg2rad(np.array([0, 90, -180, -90, 0, 90, -180, -90, 0]))
    assert np.allclose(wrap1pi(angles), expected_angles)
    assert np.allclose(wrap1pi(angles.tolist()), expected_angles)

    # Test case 2: Multidimensional array
    angles = np.deg2rad(
        np.array(
            [
                [-360, -270, -180, -90, 0, 90, 180, 270, 360],
                [-360, -270, -180, -90, 0, 90, 180, 270, 360],
            ]
        ).T
    )
    expected_angles = np.deg2rad(
        np.array(
            [
                [0, 90, -180, -90, 0, 90, -180, -90, 0],
                [0, 90, -180, -90, 0, 90, -180, -90, 0],
            ]
        ).T
    )
    assert np.allclose(wrap1pi(angles), expected_angles)
    assert np.allclose(wrap1pi(angles.tolist()), expected_angles)
    assert np.allclose(wrap1pi(angles.T), expected_angles.T)
    assert np.allclose(wrap1pi(angles.T.tolist()), expected_angles.T)


def test_unwrap2pi():
    """
    Test the unwrap2pi function.
    """
    # Test case 1: No discontinuities
    angles = np.deg2rad(
        np.array([0, 90, 180, 270, 360, 90, 180, 270, 360, 90, 180, 270, 360])
    )
    expected_angles = np.deg2rad(
        np.array([0, 90, 180, 270, 360, 450, 540, 630, 720, 810, 900, 990, 1080])
    )
    assert np.allclose(unwrap2pi(angles), expected_angles)
    assert np.allclose(unwrap2pi(angles.tolist()), expected_angles)

    # Test case 2: Discontinuities
    angles = np.deg2rad(np.array([0, 90, 180, 270, 360, 90, 320, 350, 380]))
    expected_angles = np.deg2rad(np.array([0, 90, 180, 270, 360, 450, 320, 350, 380]))
    assert np.allclose(unwrap2pi(angles), expected_angles)
    assert np.allclose(unwrap2pi(angles.tolist()), expected_angles)

    # Test case 3: Multidimensional array
    angles = np.deg2rad(
        np.array(
            [
                [0, 90, 180, 270, 360, 90, 180, 270, 360, 90, 180, 270, 360],
                [0, 90, 180, 270, 360, 90, 180, 270, 360, 90, 180, 270, 360],
            ]
        ).T
    )
    expected_angles = np.deg2rad(
        np.array(
            [
                [0, 90, 180, 270, 360, 450, 540, 630, 720, 810, 900, 990, 1080],
                [0, 90, 180, 270, 360, 450, 540, 630, 720, 810, 900, 990, 1080],
            ]
        ).T
    )
    assert np.allclose(unwrap2pi(angles), expected_angles)
    assert np.allclose(unwrap2pi(angles.tolist()), expected_angles)


def test_unwrap1pi():
    """
    Test the unwrap1pi function.
    """
    # Test case 1: No discontinuities
    angles = np.deg2rad(
        np.array([0, 45, 90, 135, 180, 45, 90, 135, 180, 45, 90, 135, 180])
    )
    expected_angles = np.deg2rad(
        np.array([0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540])
    )
    assert np.allclose(unwrap1pi(angles), expected_angles)
    assert np.allclose(unwrap1pi(angles.tolist()), expected_angles)

    # Test case 2: Discontinuities
    angles = np.deg2rad(np.array([0, 45, 90, 135, 180, 45, 145, 170, 180]))
    expected_angles = np.deg2rad(np.array([0, 45, 90, 135, 180, 225, 145, 170, 180]))
    assert np.allclose(unwrap1pi(angles), expected_angles)
    assert np.allclose(unwrap1pi(angles.tolist()), expected_angles)

    # Test case 3: Multidimensional array
    angles = np.deg2rad(
        np.array(
            [
                [0, 45, 90, 135, 180, 45, 90, 135, 180, 45, 90, 135, 180],
                [0, 45, 90, 135, 180, 45, 90, 135, 180, 45, 90, 135, 180],
            ]
        ).T
    )
    expected_angles = np.deg2rad(
        np.array(
            [
                [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
                [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
            ]
        ).T
    )
    assert np.allclose(unwrap1pi(angles), expected_angles)
    assert np.allclose(unwrap1pi(angles.tolist()), expected_angles)


def test_unwrap360():
    """
    Test the unwrap360 function.
    """
    # Test case 1: No discontinuities
    angles = np.array([0, 90, 180, 270, 360, 90, 180, 270, 360, 90, 180, 270, 360])
    expected_angles = np.array(
        [0, 90, 180, 270, 360, 450, 540, 630, 720, 810, 900, 990, 1080]
    )
    assert np.allclose(unwrap360(angles), expected_angles)
    assert np.allclose(unwrap360(angles.tolist()), expected_angles)

    # Test case 2: Discontinuities
    angles = np.array([0, 90, 180, 270, 360, 90, 320, 350, 380])
    expected_angles = np.array([0, 90, 180, 270, 360, 450, 320, 350, 380])
    assert np.allclose(unwrap360(angles), expected_angles)
    assert np.allclose(unwrap360(angles.tolist()), expected_angles)

    # Test case 3: Multidimensional array
    angles = np.array(
        [
            [0, 90, 180, 270, 360, 90, 180, 270, 360, 90, 180, 270, 360],
            [0, 90, 180, 270, 360, 90, 180, 270, 360, 90, 180, 270, 360],
        ]
    ).T
    expected_angles = np.array(
        [
            [0, 90, 180, 270, 360, 450, 540, 630, 720, 810, 900, 990, 1080],
            [0, 90, 180, 270, 360, 450, 540, 630, 720, 810, 900, 990, 1080],
        ]
    ).T
    assert np.allclose(unwrap360(angles), expected_angles)
    assert np.allclose(unwrap360(angles.tolist()), expected_angles)


def test_unwrap180():
    """
    Test the unwrap180 function.
    """
    # Test case 1: No discontinuities
    angles = np.array([0, 45, 90, 135, 180, 45, 90, 135, 180, 45, 90, 135, 180])
    expected_angles = np.array(
        [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540]
    )
    assert np.allclose(unwrap180(angles), expected_angles)
    assert np.allclose(unwrap180(angles.tolist()), expected_angles)

    # Test case 2: Discontinuities
    angles = np.array([0, 45, 90, 135, 180, 45, 145, 170, 180])
    expected_angles = np.array([0, 45, 90, 135, 180, 225, 145, 170, 180])
    assert np.allclose(unwrap180(angles), expected_angles)
    assert np.allclose(unwrap180(angles.tolist()), expected_angles)

    # Test case 3: Multidimensional array
    angles = np.array(
        [
            [0, 45, 90, 135, 180, 45, 90, 135, 180, 45, 90, 135, 180],
            [0, 45, 90, 135, 180, 45, 90, 135, 180, 45, 90, 135, 180],
        ]
    ).T
    expected_angles = np.array(
        [
            [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
            [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
        ]
    ).T
    assert np.allclose(unwrap180(angles), expected_angles)
    assert np.allclose(unwrap180(angles.tolist()), expected_angles)


def test_wrapunwrap():
    """
    Test the wrapunwrap function
    """
    # Test case
    angles = np.deg2rad(
        np.array([0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540])
    )
    expected_angles = np.deg2rad(
        np.array([0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540])
    )
    assert np.allclose(wrapunwrap(angles), expected_angles)
    assert np.allclose(wrapunwrap(angles.tolist()), expected_angles)

    # Test case 2: Discontinuities
    angles = np.deg2rad(np.array([0, 45, 90, 135, 180, 225, 145, 170, 180]))
    expected_angles = np.deg2rad(np.array([0, 45, 90, 135, 180, 225, 145, 170, 180]))
    assert np.allclose(wrapunwrap(angles), expected_angles)
    assert np.allclose(wrapunwrap(angles.tolist()), expected_angles)

    # Test case 3: Multidimensional array
    angles = np.deg2rad(
        np.array(
            [
                [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
                [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
            ]
        ).T
    )
    expected_angles = np.deg2rad(
        np.array(
            [
                [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
                [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
            ]
        ).T
    )
    assert np.allclose(wrapunwrap(angles), expected_angles)
    assert np.allclose(wrapunwrap(angles.tolist()), expected_angles)


def test_wrapunwrap360():
    """
    Test the wrapunwrap360 function
    """
    # Test case
    angles = np.array([0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540])
    expected_angles = np.array(
        [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540]
    )
    assert np.allclose(wrapunwrap360(angles), expected_angles)
    assert np.allclose(wrapunwrap360(angles.tolist()), expected_angles)

    # Test case 2: Discontinuities
    angles = np.array([0, 45, 90, 135, 180, 225, 145, 170, 180])
    expected_angles = np.array([0, 45, 90, 135, 180, 225, 145, 170, 180])
    assert np.allclose(wrapunwrap360(angles), expected_angles)
    assert np.allclose(wrapunwrap360(angles.tolist()), expected_angles)

    # Test case 3: Multidimensional array
    angles = np.array(
        [
            [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
            [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
        ]
    ).T
    expected_angles = np.array(
        [
            [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
            [0, 45, 90, 135, 180, 225, 270, 315, 360, 405, 450, 495, 540],
        ]
    ).T
    assert np.allclose(wrapunwrap360(angles), expected_angles)
    assert np.allclose(wrapunwrap360(angles.tolist()), expected_angles)


def test_transpose():
    """
    Test the transpose function.
    """
    # Test case 1: 2D matrix
    mat = np.array([[1, 2], [3, 4]])
    expected_transpose = np.array([[1, 3], [2, 4]])
    assert np.array_equal(transpose(mat), expected_transpose)

    # Test case 2: 3D matrix
    mat = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
    expected_transpose = np.array([[[1, 3], [2, 4]], [[5, 7], [6, 8]]])
    assert np.array_equal(transpose(mat), expected_transpose)

    # Test case 3: 1D vector
    vec = np.array([1, 2, 3])
    expected_transpose = np.array([1, 2, 3])  # Transpose of a 1D vector is itself
    assert np.array_equal(transpose(vec), expected_transpose)

    # Test case 4: List input
    mat = [[1, 2], [3, 4]]
    expected_transpose = np.array([[1, 3], [2, 4]])
    assert np.array_equal(transpose(mat), expected_transpose)

    # Test case 5: Wrong input type
    mat = "1, 2, 3"
    try:
        transpose(mat)
    except TypeError as e:
        assert str(e) == "Input must be a numpy array or a list."
    else:
        assert False

    # Test case 6: Dimension greater than 3
    mat = np.random.rand(2, 2, 2, 2)
    try:
        transpose(mat)
    except ValueError as e:
        assert str(e) == "The dimension of the array must be up to 3D"
    else:
        assert False
