"""Tests for rotation conversion utilities."""

import numpy as np

from sltk.data.rotations import (
    RotationConverter,
    axis_angle_to_matrix,
    matrix_to_6d,
    matrix_to_axis_angle,
    rotation_6d_to_matrix,
)


class TestAxisAngleToMatrix:
    """Tests for axis-angle to rotation matrix conversion."""

    def test_identity(self):
        """Zero axis-angle should give identity matrix."""
        aa = np.array([0.0, 0.0, 0.0])
        mat = axis_angle_to_matrix(aa)

        expected = np.eye(3)
        np.testing.assert_array_almost_equal(mat, expected, decimal=5)

    def test_90_degree_rotation(self):
        """Test 90 degree rotation around z-axis."""
        # 90 degrees = pi/2 radians around z-axis
        aa = np.array([0.0, 0.0, np.pi / 2])
        mat = axis_angle_to_matrix(aa)

        # Expected: rotate x to y, y to -x
        expected = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=np.float64)
        np.testing.assert_array_almost_equal(mat, expected, decimal=5)

    def test_batch(self):
        """Test batch conversion."""
        aa = np.random.randn(10, 3)
        mat = axis_angle_to_matrix(aa)

        assert mat.shape == (10, 3, 3)

    def test_roundtrip(self):
        """Test axis-angle -> matrix -> axis-angle roundtrip."""
        aa_original = np.array([0.1, 0.2, 0.3])
        mat = axis_angle_to_matrix(aa_original)
        aa_recovered = matrix_to_axis_angle(mat)

        # Should recover original (up to sign ambiguity for pi rotations)
        np.testing.assert_array_almost_equal(aa_original, aa_recovered, decimal=5)


class TestMatrixTo6D:
    """Tests for rotation matrix to 6D conversion."""

    def test_identity(self):
        """Identity matrix should give [1,0,0, 0,1,0]."""
        mat = np.eye(3)
        r6d = matrix_to_6d(mat)

        expected = np.array([1, 0, 0, 0, 1, 0])
        np.testing.assert_array_almost_equal(r6d, expected, decimal=5)

    def test_roundtrip(self):
        """Test matrix -> 6D -> matrix roundtrip."""
        # Create a valid rotation matrix
        aa = np.array([0.1, 0.2, 0.3])
        mat_original = axis_angle_to_matrix(aa)

        r6d = matrix_to_6d(mat_original)
        mat_recovered = rotation_6d_to_matrix(r6d)

        np.testing.assert_array_almost_equal(mat_original, mat_recovered, decimal=5)


class TestRotation6DToMatrix:
    """Tests for 6D to rotation matrix conversion."""

    def test_orthogonalization(self):
        """Test that output is a valid rotation matrix."""
        # Input slightly non-orthogonal 6D
        r6d = np.array([1.01, 0.01, 0.01, 0.01, 0.99, 0.01])
        mat = rotation_6d_to_matrix(r6d)

        # Check orthogonality: R^T @ R = I
        rtrt = mat.T @ mat
        np.testing.assert_array_almost_equal(rtrt, np.eye(3), decimal=5)

        # Check determinant = 1 (proper rotation, not reflection)
        det = np.linalg.det(mat)
        np.testing.assert_almost_equal(det, 1.0, decimal=5)


class TestRotationConverter:
    """Tests for RotationConverter class."""

    def test_to_6d(self):
        """Test conversion to 6D format."""
        converter = RotationConverter(target_format="rotation_6d")

        aa = np.array([0.1, 0.2, 0.3])
        r6d = converter.from_axis_angle(aa)

        assert r6d.shape == (6,)

    def test_to_matrix(self):
        """Test conversion to matrix format."""
        converter = RotationConverter(target_format="rotation_matrix")

        aa = np.array([0.1, 0.2, 0.3])
        mat = converter.from_axis_angle(aa)

        assert mat.shape == (3, 3)

    def test_convert_method(self):
        """Test the convert method."""
        converter = RotationConverter(target_format="rotation_6d")

        aa = np.array([0.1, 0.2, 0.3])
        r6d = converter.convert(aa, source_format="axis_angle")

        # Should be same as from_axis_angle
        expected = converter.from_axis_angle(aa)
        np.testing.assert_array_almost_equal(r6d, expected)
