import numpy as np
import pytest

from gri_utils import constants
from gri_utils.conversion import enu_xyz_rotation, xyz_enu_rotation

# Type aliases for readability
XYZ = tuple[float, float, float]
Matrix3x3 = tuple[tuple[float, float, float], ...]


@pytest.mark.parametrize(
    ("xyz", "expected_enu", "expected_xyz"),
    [
        # On X axis
        (
            (constants.ER_M, 0, 0),
            ((0, 0, 1), (1, 0, 0), (0, 1, 0)),
            ((0, 1, 0), (0, 0, 1), (1, 0, 0)),
        ),
    ],
)
def test_transform_matrix(
    xyz: XYZ,
    expected_enu: Matrix3x3,
    expected_xyz: Matrix3x3,
) -> None:
    """Test ENU/XYZ transform matrices with both array and sequence inputs."""
    # Test with sequence input
    tenu_seq = enu_xyz_rotation(xyz)
    assert np.allclose(tenu_seq, np.array(expected_enu))

    # Test with array input
    tenu_arr = enu_xyz_rotation(np.array(xyz))
    assert np.allclose(tenu_arr, np.array(expected_enu))

    # Test xyz_enu_transform (geocentric)
    txyz = xyz_enu_rotation(xyz, geocentric=True)
    assert np.allclose(txyz, np.array(expected_xyz))

    # Verify transpose relationship
    assert np.allclose(txyz, tenu_arr.T)


@pytest.mark.parametrize(
    ("xyz", "expected"),
    [
        # On X axis
        ((constants.ER_M, 0, 0), ((0, 0, 1), (1, 0, 0), (0, 1, 0))),
        # On Y axis
        ((0, constants.ER_M, 0), ((-1, 0, 0), (0, 0, 1), (0, 1, 0))),
        # On Z axis (pole) - east norm is zero, handled specially
        ((0, 0, constants.ER_M), ((1, 0, 0), (0, 1, 0), (0, 0, 1))),
    ],
)
def test_pole_and_axis_positions(xyz: XYZ, expected: Matrix3x3) -> None:
    """Test transform at poles and along axes where east vector norm is zero."""
    tenu = enu_xyz_rotation(np.array(xyz))
    assert np.allclose(tenu, np.array(expected))
