"""Tests for rotation utility functions.

Tests verify that rotation_to_quat_array correctly converts various rotation
representations to quaternion arrays for use in AOA computations.
"""

import numpy as np
from scipy.spatial.transform import Rotation

from gri_utils.observables import rotation_to_quat_array

# Converting scipy Rotation objects


def test_single_rotation_object_to_array() -> None:
    """Single scipy Rotation converts to shape (4,) array."""
    rotation = Rotation.from_euler("z", 45, degrees=True)

    quat_array = rotation_to_quat_array(rotation)

    assert quat_array.shape == (4,)
    assert np.allclose(quat_array, rotation.as_quat(), atol=1e-10)


def test_batch_rotation_object_to_array() -> None:
    """Batch scipy Rotation converts to shape (n, 4) array."""
    rotations = Rotation.from_euler("z", [0, 45, 90, 180], degrees=True)

    quat_array = rotation_to_quat_array(rotations)

    assert quat_array.shape == (4, 4)
    assert np.allclose(quat_array, rotations.as_quat(), atol=1e-10)


# Passing through quaternion arrays


def test_single_quaternion_array_passthrough() -> None:
    """Single quaternion array (4,) passes through unchanged."""
    quat = np.array([0.0, 0.0, 0.7071068, 0.7071068])

    result = rotation_to_quat_array(quat)

    assert result.shape == (4,)
    assert np.allclose(result, quat, atol=1e-10)


def test_batch_quaternion_array_passthrough() -> None:
    """Batch quaternion array (n, 4) passes through unchanged."""
    quats = np.array(
        [
            [0.0, 0.0, 0.0, 1.0],
            [0.0, 0.0, 0.7071068, 0.7071068],
            [0.0, 0.7071068, 0.0, 0.7071068],
        ],
    )

    result = rotation_to_quat_array(quats)

    assert result.shape == (3, 4)
    assert np.allclose(result, quats, atol=1e-10)


# Converting list inputs


def test_list_quaternion_to_array() -> None:
    """Python list converts to numpy array."""
    quat_list = [0.0, 0.0, 0.0, 1.0]

    result = rotation_to_quat_array(quat_list)

    assert isinstance(result, np.ndarray)
    assert result.shape == (4,)
    assert np.allclose(result, quat_list, atol=1e-10)


def test_nested_list_quaternions_to_array() -> None:
    """Nested list converts to 2D numpy array."""
    quat_list = [
        [0.0, 0.0, 0.0, 1.0],
        [0.0, 0.0, 0.7071068, 0.7071068],
    ]

    result = rotation_to_quat_array(quat_list)

    assert isinstance(result, np.ndarray)
    assert result.shape == (2, 4)


# Higher-dimensional arrays for broadcasting


def test_multidimensional_quaternion_array() -> None:
    """Higher-dimensional quaternion array preserves shape."""
    quats = np.zeros((2, 3, 4))
    quats[..., 3] = 1.0  # Identity quaternions

    result = rotation_to_quat_array(quats)

    assert result.shape == (2, 3, 4)
    assert np.allclose(result, quats, atol=1e-10)


# Consistency between Rotation and array inputs


def test_rotation_and_array_give_same_result() -> None:
    """Rotation object and its quaternion array produce identical results."""
    rotation = Rotation.from_euler("xyz", [30, 45, 60], degrees=True)
    quat = rotation.as_quat()

    result_from_rotation = rotation_to_quat_array(rotation)
    result_from_array = rotation_to_quat_array(quat)

    assert np.allclose(result_from_rotation, result_from_array, atol=1e-10)
