"""
Test module for the SE3 module.

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

import numpy as np
import pytest

from navlib import (
    adjoint,
    axis_ang6,
    distance_to_se3,
    matrix2xyzrph,
    matrix_exp6,
    matrix_inverse,
    matrix_log6,
    pose_diff,
    project_to_se3,
    se3_to_vec,
    vec_to_se3,
    vec_to_so3,
    xyzrph2matrix,
)
from navlib.math import adjoint as adjoint_relative
from navlib.math import axis_ang6 as axis_ang6_relative
from navlib.math import distance_to_se3 as distance_to_se3_relative
from navlib.math import matrix2xyzrph as matrix2xyzrph_relative
from navlib.math import matrix_exp6 as matrix_exp6_relative
from navlib.math import matrix_inverse as matrix_inverse_relative
from navlib.math import matrix_log6 as matrix_log6_relative
from navlib.math import pose_diff as pose_diff_relative
from navlib.math import project_to_se3 as project_to_se3_relative
from navlib.math import se3_to_vec as se3_to_vec_relative
from navlib.math import vec_to_se3 as vec_to_se3_relative
from navlib.math import vec_to_so3 as vec_to_so3_relative
from navlib.math import xyzrph2matrix as xyzrph2matrix_relative


def test_imports():
    # Checks that the absolute and relative imports work correctly
    assert adjoint_relative.__code__.co_code == adjoint.__code__.co_code
    assert axis_ang6_relative.__code__.co_code == axis_ang6.__code__.co_code
    assert distance_to_se3_relative.__code__.co_code == distance_to_se3.__code__.co_code
    assert matrix2xyzrph_relative.__code__.co_code == matrix2xyzrph.__code__.co_code
    assert matrix_exp6_relative.__code__.co_code == matrix_exp6.__code__.co_code
    assert matrix_inverse_relative.__code__.co_code == matrix_inverse.__code__.co_code
    assert matrix_log6_relative.__code__.co_code == matrix_log6.__code__.co_code
    assert pose_diff_relative.__code__.co_code == pose_diff.__code__.co_code
    assert project_to_se3_relative.__code__.co_code == project_to_se3.__code__.co_code
    assert se3_to_vec_relative.__code__.co_code == se3_to_vec.__code__.co_code
    assert vec_to_se3_relative.__code__.co_code == vec_to_se3.__code__.co_code
    assert vec_to_so3_relative.__code__.co_code == vec_to_so3.__code__.co_code
    assert xyzrph2matrix_relative.__code__.co_code == xyzrph2matrix.__code__.co_code


def test_xyzrph2matrix():
    """
    Test the rph2rot function.
    """
    # Test case 1: Zero rotation
    rph = np.array([0, 0, 0])
    xyz = np.array([1.0, 2.0, 3.0])
    xyzrph = np.concatenate((xyz, rph))
    expected_t_mat = np.eye(4)
    expected_t_mat[:3, 3] = xyz
    assert np.allclose(xyzrph2matrix(xyzrph), expected_t_mat)
    assert np.allclose(xyzrph2matrix(xyzrph.tolist()), expected_t_mat)

    # Test case 2: Rotation around X-axis
    rph = np.array([np.pi / 2, 0, 0])
    xyz = np.array([1.0, 2.0, 3.0])
    xyzrph = np.concatenate((xyz, rph))
    expected_t_mat = np.eye(4)
    expected_t_mat[:3, :3] = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    expected_t_mat[:3, 3] = xyz
    assert np.allclose(xyzrph2matrix(xyzrph), expected_t_mat)
    assert np.allclose(xyzrph2matrix(xyzrph.tolist()), expected_t_mat)

    # Test case 3: Rotation around Y-axis
    rph = np.array([0, np.pi / 2, 0])
    xyz = np.array([1.0, 2.0, 3.0])
    xyzrph = np.concatenate((xyz, rph))
    expected_t_mat = np.eye(4)
    expected_t_mat[:3, :3] = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    expected_t_mat[:3, 3] = xyz
    assert np.allclose(xyzrph2matrix(xyzrph), expected_t_mat)
    assert np.allclose(xyzrph2matrix(xyzrph.tolist()), expected_t_mat)

    # Test case 4: Rotation around Z-axis
    rph = np.array([0, 0, np.pi / 2])
    xyz = np.array([1.0, 2.0, 3.0])
    xyzrph = np.concatenate((xyz, rph))
    expected_t_mat = np.eye(4)
    expected_t_mat[:3, :3] = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    expected_t_mat[:3, 3] = xyz
    assert np.allclose(xyzrph2matrix(xyzrph), expected_t_mat)
    assert np.allclose(xyzrph2matrix(xyzrph.tolist()), expected_t_mat)

    # Test case 5: Rotation around X, Y, and Z axes
    rph = np.array([np.pi / 2, np.pi / 2, np.pi / 2])
    xyz = np.array([1.0, 2.0, 3.0])
    xyzrph = np.concatenate((xyz, rph))
    expected_t_mat = np.eye(4)
    expected_t_mat[:3, :3] = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    expected_t_mat[:3, 3] = xyz
    assert np.allclose(xyzrph2matrix(xyzrph), expected_t_mat)
    assert np.allclose(xyzrph2matrix(xyzrph.tolist()), expected_t_mat)

    # Test case 6: Wrong input shape
    with pytest.raises(ValueError):
        xyzrph2matrix(np.array([1, 2, 3, 0, 0]))


def test_matrix2xyzrph():
    """
    Test the matrix2xyzrph function.
    """
    # Test case 1: Zero rotation
    xyz = np.array([1.0, 2.0, 3.0])
    t_mat = np.eye(4)
    t_mat[:3, 3] = xyz
    expected_xyzrph = np.concatenate((xyz, np.array([0, 0, 0])))
    assert np.allclose(matrix2xyzrph(t_mat), expected_xyzrph)

    # Test case 2: Rotation around X-axis
    xyz = np.array([1.0, 2.0, 3.0])
    t_mat = np.eye(4)
    t_mat[:3, :3] = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    t_mat[:3, 3] = xyz
    expected_xyzrph = np.concatenate((xyz, np.array([np.pi / 2, 0, 0])))
    assert np.allclose(matrix2xyzrph(t_mat), expected_xyzrph)

    # Test case 3: Rotation around Y-axis
    xyz = np.array([1.0, 2.0, 3.0])
    t_mat = np.eye(4)
    t_mat[:3, :3] = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    t_mat[:3, 3] = xyz
    expected_xyzrph = np.concatenate((xyz, np.array([0, np.pi / 2, 0])))
    assert np.allclose(matrix2xyzrph(t_mat), expected_xyzrph)

    # Test case 4: Rotation around Z-axis
    xyz = np.array([1.0, 2.0, 3.0])
    t_mat = np.eye(4)
    t_mat[:3, :3] = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    t_mat[:3, 3] = xyz
    expected_xyzrph = np.concatenate((xyz, np.array([0, 0, np.pi / 2])))
    assert np.allclose(matrix2xyzrph(t_mat), expected_xyzrph)

    # Test case 5: Rotation around X, Y, and Z axes
    xyz = np.array([1.0, 2.0, 3.0])
    t_mat = np.eye(4)
    t_mat[:3, :3] = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    t_mat[:3, 3] = xyz
    expected_xyzrph = np.concatenate((xyz, np.array([0, np.pi / 2, 0])))
    print(matrix2xyzrph(t_mat))
    print(expected_xyzrph)
    assert np.allclose(matrix2xyzrph(t_mat), expected_xyzrph)

    # Test case 6: Wrong input shape
    with pytest.raises(ValueError):
        matrix2xyzrph(np.eye(3))

    # Test case 7: Wrong input instance
    with pytest.raises(ValueError):
        matrix2xyzrph(np.eye(4).tolist())


def test_invert_matrix():
    """
    Test the matrix_inverse function.
    """
    # Test case 1: Identity matrix
    T = np.eye(4)
    expected_result = np.eye(4)
    np.testing.assert_array_equal(matrix_inverse(T), expected_result)

    # Test case 2: Translation matrix
    T = np.eye(4)
    T[:3, 3] = [1, 2, 3]
    expected_result = np.eye(4)
    expected_result[:3, 3] = [-1, -2, -3]
    np.testing.assert_array_equal(matrix_inverse(T), expected_result)

    # Test case 3: Rotation matrix
    T = np.eye(4)
    T[:3, :3] = [[0, -1, 0], [1, 0, 0], [0, 0, 1]]  # 90 degree rotation around Z-axis
    expected_result = np.eye(4)
    expected_result[:3, :3] = [
        [0, 1, 0],
        [-1, 0, 0],
        [0, 0, 1],
    ]  # -90 degree rotation around Z-axis
    np.testing.assert_array_almost_equal(matrix_inverse(T), expected_result, decimal=8)

    # Test case 4: Wrong input shape
    with pytest.raises(ValueError):
        matrix_inverse(np.eye(3))

    # Test case 5: Wrong input instance
    with pytest.raises(ValueError):
        matrix_inverse(np.eye(4).tolist())


def test_pose_diff():
    """
    Test the pose_diff function.
    """
    # Test case 1: Identity matrices
    T1 = np.eye(4)
    T2 = np.eye(4)
    expected_result = np.eye(4)
    np.testing.assert_array_equal(pose_diff(T1, T2), expected_result)

    # Test case 2: Translation matrices
    T1 = np.eye(4)
    T1[:3, 3] = [1, 2, 3]
    T2 = np.eye(4)
    T2[:3, 3] = [4, 5, 6]
    expected_result = np.eye(4)
    expected_result[:3, 3] = [3, 3, 3]
    np.testing.assert_array_equal(pose_diff(T1, T2), expected_result)

    # Test case 3: Rotation matrices
    T1 = np.eye(4)
    T1[:3, :3] = [[0, -1, 0], [1, 0, 0], [0, 0, 1]]  # 90 degree rotation around Z-axis
    T2 = np.eye(4)
    T2[:3, :3] = [[0, 1, 0], [-1, 0, 0], [0, 0, 1]]  # -90 degree rotation around Z-axis
    expected_result = np.eye(4)
    expected_result[:3, :3] = [
        [-1, 0, 0],
        [0, -1, 0],
        [0, 0, 1],
    ]  # 180 degree rotation around Z-axis
    np.testing.assert_array_almost_equal(pose_diff(T1, T2), expected_result, decimal=8)

    # Test case 4: Not a numpy array
    with pytest.raises(ValueError):
        pose_diff("not a numpy array", np.eye(4))
    with pytest.raises(ValueError):
        pose_diff(np.eye(4), "not a numpy array")

    # Test case 5: Wrong shape
    with pytest.raises(ValueError):
        pose_diff(np.eye(3), np.eye(4))
    with pytest.raises(ValueError):
        pose_diff(np.eye(4), np.eye(3))


def test_adjoint():
    """
    Test the adjoint function.
    """
    # Test case 1: Identity matrix
    T = np.eye(4)
    expected_result = np.eye(6)
    np.testing.assert_array_equal(adjoint(T), expected_result)

    # Test case 2: Translation matrix
    T = np.eye(4)
    T[:3, 3] = [1, 2, 3]
    expected_result = np.eye(6)
    expected_result[3:, :3] = vec_to_so3([1, 2, 3])
    np.testing.assert_array_equal(adjoint(T), expected_result)

    # Test case 3: Rotation matrix
    T = np.eye(4)
    T[:3, :3] = [[0, -1, 0], [1, 0, 0], [0, 0, 1]]  # 90 degree rotation around Z-axis
    expected_result = np.zeros((6, 6))
    expected_result[:3, :3] = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    expected_result[3:, 3:] = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    np.testing.assert_array_almost_equal(adjoint(T), expected_result)

    # Test case 4: Not a numpy array
    with pytest.raises(ValueError):
        adjoint("not a numpy array")

    # Test case 5: Wrong shape
    with pytest.raises(ValueError):
        adjoint(np.eye(3))


def test_vec_to_se3():
    """
    Test the vec_to_se3 function.
    """
    # Test case 1: Zero vector
    vec = np.zeros(6)
    expected_result = np.zeros((4, 4))
    np.testing.assert_array_equal(vec_to_se3(vec), expected_result)

    # Test case 2: Non-zero vector
    vec = np.array([1, 2, 3, 4, 5, 6])
    expected_result = np.zeros((4, 4))
    expected_result[:3, :3] = vec_to_so3(vec[:3])
    expected_result[:3, 3] = vec[3:]
    np.testing.assert_array_equal(vec_to_se3(vec), expected_result)

    # Test case 3: Not a numpy array or a list
    with pytest.raises(ValueError):
        vec_to_se3("not a numpy array or a list")

    # Test case 4: Wrong number of elements
    with pytest.raises(ValueError):
        vec_to_se3(np.zeros(5))


def test_se3_to_vec():
    """
    Test the se3_to_vec function.
    """
    # Test case 1: Zero matrix
    mat = np.zeros((4, 4))
    expected_result = np.zeros(6)
    np.testing.assert_array_equal(se3_to_vec(mat), expected_result)

    # Test case 2: Non-zero matrix
    mat = np.array([[0, -3, 2, 4], [3, 0, -1, 5], [-2, 1, 0, 6], [0, 0, 0, 0]])
    expected_result = np.array([1, 2, 3, 4, 5, 6])
    np.testing.assert_array_equal(se3_to_vec(mat), expected_result)

    # Test case 3: Not a numpy array
    with pytest.raises(ValueError):
        se3_to_vec("not a numpy array")

    # Test case 4: Wrong shape
    with pytest.raises(ValueError):
        se3_to_vec(np.eye(3))


def test_axis_ang6():
    """
    Test the axis_ang6 function.
    """
    # Test case 1: Zero vector
    vec = np.array([0, 0, 0, 1, 1, 1])
    expected_axis = vec / np.linalg.norm(np.array([1, 1, 1]))
    expected_theta = np.linalg.norm(np.array([1, 1, 1]))
    axis, theta = axis_ang6(vec)
    np.testing.assert_array_equal(axis, expected_axis)
    assert theta == expected_theta

    # Test case 2: Non-zero vector
    vec = np.array([1, 2, 3, 4, 5, 6])
    expected_axis = vec / np.linalg.norm(vec[:3])
    expected_theta = np.linalg.norm(vec[:3])
    axis, theta = axis_ang6(vec)
    np.testing.assert_array_almost_equal(axis, expected_axis, decimal=8)
    assert np.isclose(theta, expected_theta, atol=1e-8)

    # Test case 3: Not a numpy array or a list
    with pytest.raises(ValueError):
        axis_ang6("not a numpy array or a list")

    # Test case 4: Wrong number of elements
    with pytest.raises(ValueError):
        axis_ang6(np.zeros(5))


def test_matrix_exp6():
    """
    Test the matrix_exp6 function.
    """
    # Test case 1: Non-zero matrix
    se3mat = np.array(
        [
            [0, 0, 0, 0],
            [0, 0, -1.57079632, 2.35619449],
            [0, 1.57079632, 0, 2.35619449],
            [0, 0, 0, 0],
        ]
    )
    expected_result = np.array(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, -1.0, 0.0],
            [0.0, 1.0, 0.0, 3.0],
            [0, 0, 0, 1],
        ]
    )
    np.testing.assert_array_almost_equal(
        matrix_exp6(se3mat), expected_result, decimal=8
    )

    # Test case 3: Not a numpy array
    with pytest.raises(ValueError):
        matrix_exp6("not a numpy array")

    # Test case 4: Wrong shape
    with pytest.raises(ValueError):
        matrix_exp6(np.eye(3))


def test_matrix_log6():
    """
    Test the matrix_log6 function.
    """
    # Test case 1: Identity matrix
    mat = np.eye(4)
    expected_result = np.zeros((4, 4))
    np.testing.assert_array_equal(matrix_log6(mat), expected_result)

    # Test case 2: Non-identity matrix
    mat = np.array(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, -1.0, 0.0],
            [0.0, 1.0, 0.0, 3.0],
            [0, 0, 0, 1],
        ]
    )
    expected_result = np.array(
        [
            [0, 0, 0, 0],
            [0, 0, -1.57079632, 2.35619449],
            [0, 1.57079632, 0, 2.35619449],
            [0, 0, 0, 0],
        ]
    )
    np.testing.assert_array_almost_equal(matrix_log6(mat), expected_result, decimal=8)

    # Test case 3: Not a numpy array
    with pytest.raises(ValueError):
        matrix_log6("not a numpy array")

    # Test case 4: Wrong shape
    with pytest.raises(ValueError):
        matrix_log6(np.eye(3))


def test_project_to_SE3():
    """
    Test the project_to_SE3 function.
    """
    # Test case 1: Non-zero matrix
    mat = np.array(
        [
            [0.675, 0.150, 0.720, 1.2],
            [0.370, 0.771, -0.511, 5.4],
            [-0.630, 0.619, 0.472, 3.6],
            [0.003, 0.002, 0.010, 0.9],
        ]
    )
    expected_result = np.array(
        [
            [0.67901136, 0.14894516, 0.71885945, 1.2],
            [0.37320708, 0.77319584, -0.51272279, 5.4],
            [-0.63218672, 0.61642804, 0.46942137, 3.6],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    np.testing.assert_array_almost_equal(
        project_to_se3(mat), expected_result, decimal=8
    )

    # Test case 3: Not a numpy array
    with pytest.raises(ValueError):
        project_to_se3("not a numpy array")

    # Test case 4: Wrong shape
    with pytest.raises(ValueError):
        project_to_se3(np.eye(3))


def test_distance_to_SE3():
    """
    Test the distance_to_SE3 function.
    """
    # Test case: Given example
    mat = np.array(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 0.1, -0.95, 0.0],
            [0.0, 1.0, 0.1, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    expected_result = 0.08835
    assert np.isclose(distance_to_se3(mat), expected_result, atol=1e-5)

    # Test case 1: Not a numpy array
    with pytest.raises(ValueError):
        distance_to_se3("not a numpy array")

    # Test case 2: Wrong shape
    with pytest.raises(ValueError):
        distance_to_se3(np.eye(3))
