"""
Test module for the SO3 module.

Authors: Sebastián Rodríguez-Martínez
Contact: srodriguez@mbari.org
Last Update: 2023-10-22
"""

import autograd.numpy as np_ag
import numpy as np
import pytest

from navlib import (
    axis_ang3,
    distance_to_so3,
    hpr2quat,
    hpr2rot,
    matrix_exp3,
    matrix_log3,
    project_to_so3,
    quat2hpr,
    quat2rot,
    quat2rph,
    rot2hpr,
    rot2quat,
    rot2rph,
    rot_diff,
    rot_inv,
    rph2quat,
    rph2rot,
    so3_optimization,
    so3_to_vec,
    vec_to_so3,
)
from navlib.math import axis_ang3 as axis_ang3_relative
from navlib.math import distance_to_so3 as distance_to_so3_relative
from navlib.math import hpr2quat as hpr2quat_relative
from navlib.math import hpr2rot as hpr2rot_relative
from navlib.math import matrix_exp3 as matrix_exp3_relative
from navlib.math import matrix_log3 as matrix_log3_relative
from navlib.math import project_to_so3 as project_to_so3_relative
from navlib.math import quat2hpr as quat2hpr_relative
from navlib.math import quat2rot as quat2rot_relative
from navlib.math import quat2rph as quat2rph_relative
from navlib.math import rot2hpr as rot2hpr_relative
from navlib.math import rot2quat as rot2quat_relative
from navlib.math import rot2rph as rot2rph_relative
from navlib.math import rot_diff as rot_diff_relative
from navlib.math import rot_inv as rot_inv_relative
from navlib.math import rph2quat as rph2quat_relative
from navlib.math import rph2rot as rph2rot_relative
from navlib.math import so3_optimization as so3_optimization_relative
from navlib.math import so3_to_vec as so3_to_vec_relative
from navlib.math import vec_to_so3 as vec_to_so3_relative


def test_imports():
    # Checks that the absolute and relative imports work correctly
    assert axis_ang3_relative.__code__.co_code == axis_ang3.__code__.co_code
    assert distance_to_so3_relative.__code__.co_code == distance_to_so3.__code__.co_code
    assert hpr2quat_relative.__code__.co_code == hpr2quat.__code__.co_code
    assert hpr2rot_relative.__code__.co_code == hpr2rot.__code__.co_code
    assert matrix_exp3_relative.__code__.co_code == matrix_exp3.__code__.co_code
    assert matrix_log3_relative.__code__.co_code == matrix_log3.__code__.co_code
    assert project_to_so3_relative.__code__.co_code == project_to_so3.__code__.co_code
    assert quat2hpr_relative.__code__.co_code == quat2hpr.__code__.co_code
    assert quat2rot_relative.__code__.co_code == quat2rot.__code__.co_code
    assert quat2rph_relative.__code__.co_code == quat2rph.__code__.co_code
    assert rot2hpr_relative.__code__.co_code == rot2hpr.__code__.co_code
    assert rot2quat_relative.__code__.co_code == rot2quat.__code__.co_code
    assert rot2rph_relative.__code__.co_code == rot2rph.__code__.co_code
    assert rot_diff_relative.__code__.co_code == rot_diff.__code__.co_code
    assert rot_inv_relative.__code__.co_code == rot_inv.__code__.co_code
    assert rph2quat_relative.__code__.co_code == rph2quat.__code__.co_code
    assert rph2rot_relative.__code__.co_code == rph2rot.__code__.co_code
    assert so3_to_vec_relative.__code__.co_code == so3_to_vec.__code__.co_code
    assert vec_to_so3_relative.__code__.co_code == vec_to_so3.__code__.co_code
    assert (
        so3_optimization_relative.__code__.co_code == so3_optimization.__code__.co_code
    )


def test_rot_inv():
    """
    Test the rot_inv function.
    """
    # Test case 1: Zero rotation
    rot_mat = np.eye(3)
    expected_rot_mat = np.eye(3)
    assert np.allclose(rot_inv(rot_mat), expected_rot_mat)
    assert np.allclose(rot_inv(rot_mat.tolist()), expected_rot_mat)

    # Test case 2: Random rotation
    rot_mat = np.array(
        [
            [0.93629336, -0.27509585, 0.21835066],
            [0.28962948, 0.95642509, -0.03695701],
            [-0.19866933, 0.0978434, 0.97517033],
        ]
    )
    expected_rot_mat = np.array(
        [
            [0.93629336, 0.28962948, -0.19866933],
            [-0.27509585, 0.95642509, 0.0978434],
            [0.21835066, -0.03695701, 0.97517033],
        ]
    )
    assert np.allclose(rot_inv(rot_mat), expected_rot_mat)
    assert np.allclose(rot_inv(rot_mat.tolist()), expected_rot_mat)


def test_rph2quat():
    """
    Test the rph2quat function.
    """
    # Test case 1: Zero rotation
    rph = np.array([0, 0, 0])
    expected_quat = np.array([0, 0, 0, 1])
    assert np.allclose(rph2quat(rph), expected_quat)
    assert np.allclose(rph2quat(rph.tolist()), expected_quat)

    # Test case 2: Rotation around X-axis
    rph = np.array([np.pi / 2, 0, 0])
    expected_quat = np.array([0.70710678, 0, 0, 0.70710678])
    assert np.allclose(rph2quat(rph), expected_quat)
    assert np.allclose(rph2quat(rph.tolist()), expected_quat)

    # Test case 3: Rotation around Y-axis
    rph = np.array([0, np.pi / 2, 0])
    expected_quat = np.array([0, 0.70710678, 0, 0.70710678])
    assert np.allclose(rph2quat(rph), expected_quat)
    assert np.allclose(rph2quat(rph.tolist()), expected_quat)

    # Test case 4: Rotation around Z-axis
    rph = np.array([0, 0, np.pi / 2])
    expected_quat = np.array([0, 0, 0.70710678, 0.70710678])
    assert np.allclose(rph2quat(rph), expected_quat)
    assert np.allclose(rph2quat(rph.tolist()), expected_quat)

    # Test case 5: Rotation around X, Y, and Z axes
    rph = np.array([np.pi / 2, np.pi / 2, np.pi / 2])
    expected_quat = np.array([0.5, 0.5, 0.5, 0.5])


def test_hpr2quat():
    """
    Test the hpr2quat function.
    """
    # Test case 1: Zero rotation
    hpr = np.array([0, 0, 0])
    expected_quat = np.array([0, 0, 0, 1])
    assert np.allclose(hpr2quat(hpr), expected_quat)
    assert np.allclose(hpr2quat(hpr.tolist()), expected_quat)

    # Test case 2: Rotation around X-axis
    hpr = np.array([0, 0, np.pi / 2])
    expected_quat = np.array([0.70710678, 0, 0, 0.70710678])
    assert np.allclose(hpr2quat(hpr), expected_quat)
    assert np.allclose(hpr2quat(hpr.tolist()), expected_quat)

    # Test case 3: Rotation around Y-axis
    hpr = np.array([0, np.pi / 2, 0])
    expected_quat = np.array([0, 0.70710678, 0, 0.70710678])
    assert np.allclose(hpr2quat(hpr), expected_quat)
    assert np.allclose(hpr2quat(hpr.tolist()), expected_quat)

    # Test case 4: Rotation around Z-axis
    hpr = np.array([np.pi / 2, 0, 0])
    expected_quat = np.array([0, 0, 0.70710678, 0.70710678])
    assert np.allclose(hpr2quat(hpr), expected_quat)
    assert np.allclose(hpr2quat(hpr.tolist()), expected_quat)

    # Test case 5: Rotation around X, Y, and Z axes
    hpr = np.array([np.pi / 2, np.pi / 2, np.pi / 2])
    expected_quat = np.array([0, 0.70710678, 0, 0.70710678])
    assert np.allclose(hpr2quat(hpr), expected_quat)
    assert np.allclose(hpr2quat(hpr.tolist()), expected_quat)


def test_quat2rph():
    """
    Test the quat2rph function.
    """
    # Test case 1: Zero rotation
    quat = np.array([0, 0, 0, 1])
    expected_rph = np.array([0, 0, 0])
    assert np.allclose(quat2rph(quat), expected_rph)
    assert np.allclose(quat2rph(quat.tolist()), expected_rph)

    # Test case 2: Rotation around X-axis
    quat = np.array([0.3826834, 0, 0, 0.9238795])
    expected_rph = np.array([np.pi / 4, 0, 0])
    assert np.allclose(quat2rph(quat), expected_rph)
    assert np.allclose(quat2rph(quat.tolist()), expected_rph)

    # Test case 3: Rotation around Y-axis
    quat = np.array([0, 0.3826834, 0, 0.9238795])
    expected_rph = np.array([0, np.pi / 4, 0])
    assert np.allclose(quat2rph(quat), expected_rph)
    assert np.allclose(quat2rph(quat.tolist()), expected_rph)

    # Test case 4: Rotation around Z-axis
    quat = np.array([0, 0, 0.3826834, 0.9238795])
    expected_rph = np.array([0, 0, np.pi / 4])
    assert np.allclose(quat2rph(quat), expected_rph)
    assert np.allclose(quat2rph(quat.tolist()), expected_rph)

    # Test case 5: Rotation around X, Y, and Z axes
    quat = np.array([0.1913417, 0.4619398, 0.1913417, 0.8446232])
    expected_rph = np.array([np.pi / 4, np.pi / 4, np.pi / 4])
    assert np.allclose(quat2rph(quat), expected_rph)
    assert np.allclose(quat2rph(quat.tolist()), expected_rph)


def test_quat2hpr():
    """
    Test the quat2hpr function.
    """
    # Test case 1: Zero rotation
    quat = np.array([0, 0, 0, 1])
    expected_hpr = np.array([0, 0, 0])
    assert np.allclose(quat2hpr(quat), expected_hpr)
    assert np.allclose(quat2hpr(quat.tolist()), expected_hpr)

    # Test case 2: Rotation around X-axis
    quat = np.array([0.3826834, 0, 0, 0.9238795])
    expected_hpr = np.array([0, 0, np.pi / 4])
    assert np.allclose(quat2hpr(quat), expected_hpr)
    assert np.allclose(quat2hpr(quat.tolist()), expected_hpr)

    # Test case 3: Rotation around Y-axis
    quat = np.array([0, 0.3826834, 0, 0.9238795])
    expected_hpr = np.array([0, np.pi / 4, 0])
    assert np.allclose(quat2hpr(quat), expected_hpr)
    assert np.allclose(quat2hpr(quat.tolist()), expected_hpr)

    # Test case 4: Rotation around Z-axis
    quat = np.array([0, 0, 0.3826834, 0.9238795])
    expected_hpr = np.array([np.pi / 4, 0, 0])
    assert np.allclose(quat2hpr(quat), expected_hpr)
    assert np.allclose(quat2hpr(quat.tolist()), expected_hpr)

    # Test case 5: Rotation around X, Y, and Z axes
    quat = np.array([0.1913417, 0.4619398, 0.1913417, 0.8446232])
    expected_hpr = np.array([np.pi / 4, np.pi / 4, np.pi / 4])
    assert np.allclose(quat2hpr(quat), expected_hpr)
    assert np.allclose(quat2hpr(quat.tolist()), expected_hpr)


def test_quat2rot():
    """
    Test the quat2rot function.
    """
    # Test case 1: Zero rotation
    quat = np.array([0, 0, 0, 1])
    expected_rot_mat = np.eye(3)
    assert np.allclose(quat2rot(quat), expected_rot_mat)
    assert np.allclose(quat2rot(quat.tolist()), expected_rot_mat)

    # Test case 2: Rotation around X-axis
    quat = np.array([0.3826834, 0, 0, 0.9238795])
    expected_rot_mat = np.array(
        [[1, 0, 0], [0, 0.7071068, -0.7071068], [0, 0.7071068, 0.7071068]]
    )
    assert np.allclose(quat2rot(quat), expected_rot_mat)
    assert np.allclose(quat2rot(quat.tolist()), expected_rot_mat)

    # Test case 3: Rotation around Y-axis
    quat = np.array([0, 0.3826834, 0, 0.9238795])
    expected_rot_mat = np.array(
        [[0.7071068, 0, 0.7071068], [0, 1, 0], [-0.7071068, 0, 0.7071068]]
    )
    assert np.allclose(quat2rot(quat), expected_rot_mat)
    assert np.allclose(quat2rot(quat.tolist()), expected_rot_mat)

    # Test case 4: Rotation around Z-axis
    quat = np.array([0, 0, 0.3826834, 0.9238795])
    expected_rot_mat = np.array(
        [[0.7071068, -0.7071068, 0], [0.7071068, 0.7071068, 0], [0, 0, 1]]
    )
    assert np.allclose(quat2rot(quat), expected_rot_mat)
    assert np.allclose(quat2rot(quat.tolist()), expected_rot_mat)

    # Test case 5: Rotation around X, Y, and Z axes
    quat = np.array([0.1913417, 0.4619398, 0.1913417, 0.8446232])
    expected_rot_mat = np.array(
        [
            [0.50000000, -0.1464466, 0.8535534],
            [0.50000000, 0.8535534, -0.1464466],
            [-0.7071068, 0.5000000, 0.5000000],
        ]
    )
    assert np.allclose(quat2rot(quat), expected_rot_mat)
    assert np.allclose(quat2rot(quat.tolist()), expected_rot_mat)


def test_rot2quat():
    """
    Test the rot2quat function.
    """
    # Test case 1: Zero rotation
    rot_mat = np.eye(3)
    expected_quat = np.array([0, 0, 0, 1])
    assert np.allclose(rot2quat(rot_mat), expected_quat)
    assert np.allclose(rot2quat(rot_mat.tolist()), expected_quat)

    # Test case 2: Rotation around X-axis
    rot_mat = np.array(
        [[1, 0, 0], [0, 0.7071068, -0.7071068], [0, 0.7071068, 0.7071068]]
    )
    expected_quat = np.array([0.3826834, 0, 0, 0.9238795])
    assert np.allclose(rot2quat(rot_mat), expected_quat)
    assert np.allclose(rot2quat(rot_mat.tolist()), expected_quat)

    # Test case 3: Rotation around Y-axis
    rot_mat = np.array(
        [[0.7071068, 0, 0.7071068], [0, 1, 0], [-0.7071068, 0, 0.7071068]]
    )
    expected_quat = np.array([0, 0.3826834, 0, 0.9238795])
    assert np.allclose(rot2quat(rot_mat), expected_quat)
    assert np.allclose(rot2quat(rot_mat.tolist()), expected_quat)

    # Test case 4: Rotation around Z-axis
    rot_mat = np.array(
        [[0.7071068, -0.7071068, 0], [0.7071068, 0.7071068, 0], [0, 0, 1]]
    )
    expected_quat = np.array([0, 0, 0.3826834, 0.9238795])
    assert np.allclose(rot2quat(rot_mat), expected_quat)
    assert np.allclose(rot2quat(rot_mat.tolist()), expected_quat)

    # Test case 5: Rotation around X, Y, and Z axes
    rot_mat = np.array(
        [
            [0.50000000, -0.1464466, 0.8535534],
            [0.50000000, 0.8535534, -0.1464466],
            [-0.7071068, 0.5000000, 0.5000000],
        ]
    )
    expected_quat = np.array([0.1913417, 0.4619398, 0.1913417, 0.8446232])
    assert np.allclose(rot2quat(rot_mat), expected_quat)
    assert np.allclose(rot2quat(rot_mat.tolist()), expected_quat)


def test_rph2rot():
    """
    Test the rph2rot function.
    """
    # Test case 1: Zero rotation
    rph = np.array([0, 0, 0])
    expected_rot_mat = np.eye(3)
    assert np.allclose(rph2rot(rph), expected_rot_mat)
    assert np.allclose(rph2rot(rph.tolist()), expected_rot_mat)

    # Test case 2: Rotation around X-axis
    rph = np.array([np.pi / 2, 0, 0])
    expected_rot_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    assert np.allclose(rph2rot(rph), expected_rot_mat)
    assert np.allclose(rph2rot(rph.tolist()), expected_rot_mat)

    # Test case 3: Rotation around Y-axis
    rph = np.array([0, np.pi / 2, 0])
    expected_rot_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    assert np.allclose(rph2rot(rph), expected_rot_mat)
    assert np.allclose(rph2rot(rph.tolist()), expected_rot_mat)

    # Test case 4: Rotation around Z-axis
    rph = np.array([0, 0, np.pi / 2])
    expected_rot_mat = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    assert np.allclose(rph2rot(rph), expected_rot_mat)
    assert np.allclose(rph2rot(rph.tolist()), expected_rot_mat)

    # Test case 5: Rotation around X, Y, and Z axes
    rph = np.array([np.pi / 2, np.pi / 2, np.pi / 2])
    expected_rot_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    assert np.allclose(rph2rot(rph), expected_rot_mat)
    assert np.allclose(rph2rot(rph.tolist()), expected_rot_mat)


def test_hpr2rot():
    """
    Test the hpr2rot function.
    """
    # Test case 1: Zero rotation
    hpr = np.array([0, 0, 0])
    expected_rot_mat = np.eye(3)
    assert np.allclose(hpr2rot(hpr), expected_rot_mat)
    assert np.allclose(hpr2rot(hpr.tolist()), expected_rot_mat)

    # Test case 2: Rotation around X-axis
    hpr = np.array([0, 0, np.pi / 2])
    expected_rot_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    assert np.allclose(hpr2rot(hpr), expected_rot_mat)
    assert np.allclose(hpr2rot(hpr.tolist()), expected_rot_mat)

    # Test case 3: Rotation around Y-axis
    hpr = np.array([0, np.pi / 2, 0])
    expected_rot_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    assert np.allclose(hpr2rot(hpr), expected_rot_mat)
    assert np.allclose(hpr2rot(hpr.tolist()), expected_rot_mat)

    # Test case 4: Rotation around Z-axis
    hpr = np.array([np.pi / 2, 0, 0])
    expected_rot_mat = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    assert np.allclose(hpr2rot(hpr), expected_rot_mat)
    assert np.allclose(hpr2rot(hpr.tolist()), expected_rot_mat)

    # Test case 5: Rotation around X, Y, and Z axes
    hpr = np.array([np.pi / 2, np.pi / 2, np.pi / 2])
    expected_rot_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    assert np.allclose(hpr2rot(hpr), expected_rot_mat)
    assert np.allclose(hpr2rot(hpr.tolist()), expected_rot_mat)


def test_rot2rph():
    """
    Test the rot2rph function.
    """
    # Test case 1: Zero rotation
    rot_mat = np.eye(3)
    expected_rph = np.array([0, 0, 0])
    assert np.allclose(rot2rph(rot_mat), expected_rph)
    assert np.allclose(rot2rph(rot_mat.tolist()), expected_rph)

    # Test case 2: Rotation around X-axis
    rot_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    expected_rph = np.array([np.pi / 2, 0, 0])
    assert np.allclose(rot2rph(rot_mat), expected_rph)
    assert np.allclose(rot2rph(rot_mat.tolist()), expected_rph)

    # Test case 3: Rotation around Y-axis
    rot_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    expected_rph = np.array([0, np.pi / 2, 0])
    assert np.allclose(rot2rph(rot_mat), expected_rph)
    assert np.allclose(rot2rph(rot_mat.tolist()), expected_rph)

    # Test case 4: Rotation around Z-axis
    rot_mat = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    expected_rph = np.array([0, 0, np.pi / 2])
    assert np.allclose(rot2rph(rot_mat), expected_rph)
    assert np.allclose(rot2rph(rot_mat.tolist()), expected_rph)

    # Test case 5: Rotation around X, Y, and Z axes
    rot_mat = np.array(
        [
            [0.93629336, -0.27509585, 0.21835066],
            [0.28962948, 0.95642509, -0.03695701],
            [-0.19866933, 0.0978434, 0.97517033],
        ]
    )
    expected_rph = np.array([0.1, 0.2, 0.3])
    assert np.allclose(rot2rph(rot_mat), expected_rph)
    assert np.allclose(rot2rph(rot_mat.tolist()), expected_rph)


def test_rot2hpr():
    """
    Test the rot2hpr function.
    """
    # Test case 1: Zero rotation
    rot_mat = np.eye(3)
    expected_hpr = np.array([0, 0, 0])
    assert np.allclose(rot2hpr(rot_mat), expected_hpr)
    assert np.allclose(rot2hpr(rot_mat.tolist()), expected_hpr)

    # Test case 2: Rotation around X-axis
    rot_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
    expected_hpr = np.array([0, 0, np.pi / 2])
    assert np.allclose(rot2hpr(rot_mat), expected_hpr)
    assert np.allclose(rot2hpr(rot_mat.tolist()), expected_hpr)

    # Test case 3: Rotation around Y-axis
    rot_mat = np.array([[0, 0, 1], [0, 1, 0], [-1, 0, 0]])
    expected_hpr = np.array([0, np.pi / 2, 0])
    assert np.allclose(rot2hpr(rot_mat), expected_hpr)
    assert np.allclose(rot2hpr(rot_mat.tolist()), expected_hpr)

    # Test case 4: Rotation around Z-axis
    rot_mat = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
    expected_hpr = np.array([np.pi / 2, 0, 0])
    assert np.allclose(rot2hpr(rot_mat), expected_hpr)
    assert np.allclose(rot2hpr(rot_mat.tolist()), expected_hpr)

    # Test case 5: Rotation around X, Y, and Z axes
    rot_mat = np.array(
        [
            [0.93629336, -0.27509585, 0.21835066],
            [0.28962948, 0.95642509, -0.03695701],
            [-0.19866933, 0.0978434, 0.97517033],
        ]
    )
    expected_hpr = np.array([0.3, 0.2, 0.1])
    assert np.allclose(rot2hpr(rot_mat), expected_hpr)
    assert np.allclose(rot2hpr(rot_mat.tolist()), expected_hpr)


def test_rot_diff():
    """
    Test the rot_diff function.
    """
    # Test case 1: Zero rotation
    rot_mat_1 = np.eye(3)
    rot_mat_2 = np.eye(3)
    expected_rot_diff = np.eye(3)
    assert np.allclose(rot_diff(rot_mat_1, rot_mat_2), expected_rot_diff)

    # Test case 2: Combined rotation
    rot_mat_1 = np.array(
        [
            [0.93629336, -0.27509585, 0.21835066],
            [0.28962948, 0.95642509, -0.03695701],
            [-0.19866933, 0.0978434, 0.97517033],
        ]
    )
    rot_mat_2 = np.array(
        [
            [0.82298387, -0.41998535, 0.38175244],
            [0.56975033, 0.81758362, -0.08209422],
            [-0.0, 0.3939193, 0.91914503],
        ]
    )
    expected_rot_diff = np.array(
        [
            [0.93557083, -0.23469286, 0.15104944],
            [0.31852406, 0.93603611, -0.09360322],
            [0.1586428, 0.26221888, 0.98271281],
        ]
    )
    assert np.allclose(rot_diff(rot_mat_1, rot_mat_2), expected_rot_diff)


def test_vec2so3():
    """
    Test the skew_symmetric function.
    """
    # Test case 1: Zero vector
    vec = np.array([0, 0, 0])
    expected_skew_sym = np.zeros((3, 3))
    assert np.allclose(vec_to_so3(vec), expected_skew_sym)
    assert np.allclose(vec_to_so3(vec.tolist()), expected_skew_sym)

    # Test case 2: Non-zero vector
    vec = np.array([1, 2, 3])
    expected_skew_sym = np.array([[0, -3, 2], [3, 0, -1], [-2, 1, 0]])
    assert np.allclose(vec_to_so3(vec), expected_skew_sym)
    assert np.allclose(vec_to_so3(vec.tolist()), expected_skew_sym)


def test_so3_to_vec():
    """
    Test the so3_to_vec function.
    """
    # Test case 1: Zero vector
    so3_mat = np.zeros((3, 3))
    expected_vec = np.array([0, 0, 0])
    assert np.allclose(so3_to_vec(so3_mat), expected_vec)
    assert np.allclose(so3_to_vec(so3_mat.tolist()), expected_vec)

    # Test case 2: Non-zero vector
    so3_mat = np.array([[0, -3, 2], [3, 0, -1], [-2, 1, 0]])
    expected_vec = np.array([1, 2, 3])
    assert np.allclose(so3_to_vec(so3_mat), expected_vec)
    assert np.allclose(so3_to_vec(so3_mat.tolist()), expected_vec)


def test_axis_ang3():
    """
    Test the axis_ang3 function.
    """
    # Test case 1: Zero vector
    vector3 = np.array([0, 0, 0])
    expected_vec = np.array([0, 0, 0])
    expected_ang = 0
    vec, ang = axis_ang3(vector3)
    assert np.allclose(vec, expected_vec)
    assert np.allclose(ang, expected_ang)

    vec, ang = axis_ang3(vector3.tolist())
    assert np.allclose(vec, expected_vec)
    assert np.allclose(ang, expected_ang)

    # Test case 2: Non-zero vector
    vector3 = np.array([1, 2, 3])
    expected_vec = np.array([0.26726124, 0.53452248, 0.80178373])
    expected_ang = 3.7416573867739413
    vec, ang = axis_ang3(vector3)
    assert np.allclose(vec, expected_vec)
    assert np.allclose(ang, expected_ang)

    vec, ang = axis_ang3(vector3.tolist())
    assert np.allclose(vec, expected_vec)
    assert np.allclose(ang, expected_ang)


def test_matrix_exp3():
    """
    Test the matrix_exp3 function.
    """
    # Test case 1: Zero vector
    so3mat = np.zeros((3, 3))
    expected_mat = np.eye(3)
    assert np.allclose(matrix_exp3(so3mat), expected_mat)
    assert np.allclose(matrix_exp3(so3mat.tolist()), expected_mat)

    # Test case 2: Non-zero vector
    so3mat = np.array([[0, -3, 2], [3, 0, -1], [-2, 1, 0]])
    expected_mat = np.array(
        [
            [-0.69492056, 0.71352099, 0.08929286],
            [-0.19200697, -0.30378504, 0.93319235],
            [0.69297817, 0.63134970, 0.34810748],
        ]
    )
    assert np.allclose(matrix_exp3(so3mat), expected_mat)
    assert np.allclose(matrix_exp3(so3mat.tolist()), expected_mat)


def test_matrix_log3():
    """
    Test the matrix_log3 function.
    """
    # Test case 1: Identity matrix
    so3mat = np.eye(3)
    expected_mat = np.zeros((3, 3))
    assert np.allclose(matrix_log3(so3mat), expected_mat)
    assert np.allclose(matrix_log3(so3mat.tolist()), expected_mat)

    # Test case 2: Non-zero vector
    so3mat = np.array(
        [
            [0.93629336, -0.27509585, 0.21835066],
            [0.28962948, 0.95642509, -0.03695701],
            [-0.19866933, 0.0978434, 0.97517033],
        ]
    )
    expected_mat = np.array(
        [
            [0.0, -0.28874894, 0.21322592],
            [0.28874894, 0.0, -0.06892461],
            [-0.21322592, 0.06892461, 0.0],
        ]
    )
    assert np.allclose(matrix_log3(so3mat), expected_mat)
    assert np.allclose(matrix_log3(so3mat.tolist()), expected_mat)


def test_project_to_SO3():
    """
    Test the project_to_SO3 function.
    """
    # Test case
    so3mat = np.array(
        [[0.675, 0.150, 0.720], [0.370, 0.771, -0.511], [-0.630, 0.619, 0.472]]
    )
    expected_mat = np.array(
        [
            [0.67901136, 0.14894516, 0.71885945],
            [0.37320708, 0.77319584, -0.51272279],
            [-0.63218672, 0.61642804, 0.46942137],
        ]
    )
    assert np.allclose(project_to_so3(so3mat), expected_mat)
    assert np.allclose(project_to_so3(so3mat.tolist()), expected_mat)


def test_distance_to_so3():
    """
    Test the distance_to_so3 function.
    """
    # Test case: Zero rotation
    so3mat = np.eye(3)
    expected_dist = 0.0
    assert np.allclose(distance_to_so3(so3mat), expected_dist)
    assert np.allclose(distance_to_so3(so3mat.tolist()), expected_dist)

    # Test case: Non-zero rotation with negative determinant
    so3mat = -np.eye(3)
    expected_dist = 1.00e10
    assert np.allclose(distance_to_so3(so3mat), expected_dist)
    assert np.allclose(distance_to_so3(so3mat.tolist()), expected_dist)

    # Test case: Non-zero rotation with positive determinant
    so3mat = np.array([[1.0, 0.0, 0.0], [0.0, 0.1, -0.95], [0.0, 1.0, 0.1]])
    expected_dist = 0.088352985
    assert np.allclose(distance_to_so3(so3mat), expected_dist)
    assert np.allclose(distance_to_so3(so3mat.tolist()), expected_dist)


def test_so3_optimization_success():
    """
    Test that so3_optimization returns a valid SO(3) matrix when
    optimizing a simple objective function.
    """
    # Define target as identity so that the optimal rotation is identity.
    target = np.eye(3)

    def objective(x, target):
        # Minimize squared Frobenius norm of (x - target)
        return np_ag.linalg.norm(x - target) ** 2

    result = so3_optimization(objective, target)

    # Check that result is a 3x3 numpy array
    assert isinstance(result, np.ndarray)
    assert result.shape == (3, 3)

    # Verify that result is a rotation matrix (orthonormal with determinant +1)
    assert np.allclose(result.T @ result, np.eye(3), atol=1e-5)
    assert np.allclose(np.linalg.det(result), 1.0, atol=1e-5)


def test_so3_optimization_extra_args():
    """
    Test that so3_optimization correctly passes positional and keyword
    arguments to the objective function.
    """

    def objective(x, scale, offset):
        # Objective scales the squared distance from identity and adds an offset.
        return scale * np_ag.linalg.norm(x - np_ag.eye(3)) ** 2 + offset

    result = so3_optimization(objective, 1.0, offset=0.0)

    # Check that result is a 3x3 numpy array in SO(3)
    assert isinstance(result, np.ndarray)
    assert result.shape == (3, 3)
    assert np.allclose(result.T @ result, np.eye(3), atol=1e-5)
    assert np.allclose(np.linalg.det(result), 1.0, atol=1e-5)


def test_so3_optimization_non_callable():
    """
    Test that using a non-callable objective_function raises a TypeError.
    """
    with pytest.raises(TypeError) as excinfo:
        so3_optimization("not a function")
    assert "The objective_function must be callable." in str(excinfo.value)


def quat_mult(q, p):
    """
    Multiply two quaternions q and p (both in [x, y, z, w] order).
    Returns q ⊗ p.
    """
    x1, y1, z1, w1 = q
    x2, y2, z2, w2 = p
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    return np.array([x, y, z, w], dtype=np.float64)
