"""Tests for xyz_enu_matrix module - covariance matrix transformations."""

import numpy as np

from gri_utils.constants import ER_M
from gri_utils.conversion import enu_to_xyz_cov, xyz_to_enu_cov

# Type alias for readability
XYZ = tuple[float, float, float]


def test_xyz_to_enu_cov_single() -> None:
    """Test xyz_to_enu_cov with a single 3x3 matrix."""
    # XYZ diagonal matrix: X=100, Y=25, Z=1
    xyz = np.array(((100, 0, 0), (0, 25, 0), (0, 0, 1)))
    # At (ER_M, 0, 0): E=Y, N=Z, U=X -> ENU diagonal: E=25, N=1, U=100
    enu = np.array(((25, 0, 0), (0, 1, 0), (0, 0, 100)))
    ans = xyz_to_enu_cov((ER_M, 0, 0), xyz)
    assert np.allclose(enu, ans)


def test_enu_to_xyz_cov_single() -> None:
    """Test enu_to_xyz_cov with a single 3x3 matrix."""
    xyz = np.array(((100, 0, 0), (0, 25, 0), (0, 0, 1)))
    enu = np.array(((25, 0, 0), (0, 1, 0), (0, 0, 100)))
    ans = enu_to_xyz_cov((ER_M, 0, 0), enu)
    assert np.allclose(xyz, ans)


def test_xyz_to_enu_cov_vectorized() -> None:
    """Test xyz_to_enu_cov with stacked (N, 3, 3) matrices."""
    xyz_ref: XYZ = (ER_M, 0, 0)

    # Stack of 3 different XYZ covariance matrices
    xyz_matrices = np.array(
        [
            [[100, 0, 0], [0, 25, 0], [0, 0, 1]],
            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            [[4, 0, 0], [0, 9, 0], [0, 0, 16]],
        ],
    )

    # Expected ENU matrices (at equator/prime meridian: E=Y, N=Z, U=X)
    expected = np.array(
        [
            [[25, 0, 0], [0, 1, 0], [0, 0, 100]],
            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            [[9, 0, 0], [0, 16, 0], [0, 0, 4]],
        ],
    )

    enu_matrices = xyz_to_enu_cov(xyz_ref, xyz_matrices)

    assert enu_matrices.shape == (3, 3, 3)
    assert np.allclose(enu_matrices, expected)


def test_enu_to_xyz_cov_vectorized() -> None:
    """Test enu_to_xyz_cov with stacked (N, 3, 3) matrices."""
    xyz_ref: XYZ = (ER_M, 0, 0)

    # Stack of 3 different ENU covariance matrices
    enu_matrices = np.array(
        [
            [[25, 0, 0], [0, 1, 0], [0, 0, 100]],
            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            [[9, 0, 0], [0, 16, 0], [0, 0, 4]],
        ],
    )

    # Expected XYZ matrices (at equator/prime meridian: X=U, Y=E, Z=N)
    expected = np.array(
        [
            [[100, 0, 0], [0, 25, 0], [0, 0, 1]],
            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
            [[4, 0, 0], [0, 9, 0], [0, 0, 16]],
        ],
    )

    xyz_matrices = enu_to_xyz_cov(xyz_ref, enu_matrices)

    assert xyz_matrices.shape == (3, 3, 3)
    assert np.allclose(xyz_matrices, expected)


def test_xyz_enu_matrix_round_trip() -> None:
    """Test round-trip conversion XYZ->ENU->XYZ for matrices."""
    xyz_ref: XYZ = (ER_M, 0, 0)

    # Create some arbitrary covariance matrices
    xyz_matrices = np.array(
        [
            [[10, 1, 2], [1, 20, 3], [2, 3, 30]],
            [[5, 0.5, 0], [0.5, 8, 1], [0, 1, 12]],
            [[100, 10, 5], [10, 50, 2], [5, 2, 25]],
        ],
    )

    enu_matrices = xyz_to_enu_cov(xyz_ref, xyz_matrices)
    xyz_back = enu_to_xyz_cov(xyz_ref, enu_matrices)

    assert xyz_back.shape == xyz_matrices.shape
    assert np.allclose(xyz_back, xyz_matrices)


def test_xyz_enu_matrix_preserves_symmetry() -> None:
    """Test that matrix transformation preserves symmetry of covariance matrices."""
    xyz_ref: XYZ = (ER_M, 0, 0)

    # Symmetric covariance matrix
    xyz_cov = np.array([[10, 2, 3], [2, 20, 4], [3, 4, 30]])

    enu_cov = xyz_to_enu_cov(xyz_ref, xyz_cov)

    # Result should be symmetric
    assert np.allclose(enu_cov, enu_cov.T)
