"""Tests for PDOA to AOA conversion.

Tests verify that pdoa_to_aoa correctly converts phase differences back
to direction cosines, forming a round-trip with get_pdoa.
"""

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

from gri_utils.constants import C
from gri_utils.conversion import xyz_enu_rotation
from gri_utils.observables.pdoa import get_pdoa
from gri_utils.observables.pdoa_to_aoa import pdoa_to_aoa

# Round-trip tests: get_pdoa -> pdoa_to_aoa should recover direction


def test_round_trip_cardinal_directions() -> None:
    """Round-trip test for cardinal directions."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity

    frequency_hz = C / 0.3  # 1 GHz
    # L-shaped array in XY plane (like in the email)
    element_positions = np.array(
        [
            [0.075, 0.0, 0.0],  # Along +X
            [0.15, 0.0, 0.0],  # Further along +X
            [0.0, 0.075, 0.0],  # Along +Y
        ],
    )

    # Test cardinal directions
    directions = [
        np.array([1.0, 0.0, 0.0]),  # +X
        np.array([0.0, 1.0, 0.0]),  # +Y
        np.array([0.0, 0.0, 1.0]),  # +Z
        np.array([1.0, 1.0, 0.0]) / np.sqrt(2),  # +X+Y diagonal
    ]

    for direction in directions:
        # Create source far away in given direction
        source_xyz = collector_xyz + 100000.0 * direction

        # Forward: position -> pdoa
        pdoa = get_pdoa(
            collector_xyz,
            source_xyz,
            quat,
            element_positions,
            frequency_hz,
        )

        # Inverse: pdoa -> aoa (returns (K, 2), take principal solution)
        aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
        aoa = aoas[0]

        # Compare direction cosines
        expected_aoa = direction[:2]
        assert np.allclose(aoa, expected_aoa, atol=1e-6)


def test_round_trip_with_enu_rotation() -> None:
    """Round-trip test with ENU frame rotation."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    rot_matrix = xyz_enu_rotation(collector_xyz)
    quat = Rotation.from_matrix(rot_matrix).as_quat()

    frequency_hz = C / 0.1
    element_positions = np.array(
        [
            [0.025, 0.0, 0.0],  # East-pointing
            [0.0, 0.025, 0.0],  # North-pointing
            [0.025, 0.025, 0.0],  # Diagonal
        ],
    )

    # Source to the East in ECEF (which is +X in ENU)
    source_xyz = np.array([6378137.0, 10000.0, 0.0])

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
    aoa = aoas[0]

    # Expected: direction is East = [1, 0, 0] in ENU, so aoa = [1, 0]
    assert np.allclose(aoa, [1.0, 0.0], atol=1e-6)


def test_round_trip_elevated_source_3d_array() -> None:
    """Round-trip test for source with elevation angle using 3D array.

    A planar array cannot recover the Z component of a 3D direction.
    This test uses a 3D array configuration to properly recover all
    direction components.
    """
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])

    frequency_hz = C / 0.3
    # 3D array (elements not all in same plane) for full direction recovery
    element_positions = np.array(
        [
            [0.1, 0.0, 0.0],
            [0.0, 0.1, 0.0],
            [0.0, 0.0, 0.1],  # Z-axis element enables elevation recovery
        ],
    )

    # Source at 30 deg elevation, 45 deg azimuth
    elevation = np.radians(30)
    azimuth = np.radians(45)

    # In array frame
    direction = np.array(
        [
            np.cos(elevation) * np.sin(azimuth),
            np.cos(elevation) * np.cos(azimuth),
            np.sin(elevation),
        ],
    )

    source_xyz = collector_xyz + 100000.0 * direction

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
    aoa = aoas[0]

    expected_aoa = direction[:2]
    assert np.allclose(aoa, expected_aoa, atol=1e-5)


def test_planar_array_recovers_azimuth_ratio() -> None:
    """Planar array recovers correct azimuth even with elevation.

    A planar array in XY cannot recover the Z component, but the
    ratio of X to Y (i.e., azimuth direction) should be correct.
    """
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])

    frequency_hz = C / 0.3
    # Planar array in XY
    element_positions = np.array(
        [
            [0.1, 0.0, 0.0],
            [0.0, 0.1, 0.0],
        ],
    )

    # Source at 30 deg elevation, 45 deg azimuth
    elevation = np.radians(30)
    azimuth = np.radians(45)

    direction = np.array(
        [
            np.cos(elevation) * np.sin(azimuth),
            np.cos(elevation) * np.cos(azimuth),
            np.sin(elevation),
        ],
    )

    source_xyz = collector_xyz + 100000.0 * direction

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
    aoa = aoas[0]

    # The ratio X/Y should match the azimuth (tan(azimuth) = sin/cos = X/Y)
    # At 45 degrees, X should equal Y
    assert np.allclose(aoa[0] / aoa[1], np.tan(azimuth), atol=1e-5)


# Overdetermined system tests


def test_overdetermined_system_4_elements() -> None:
    """Test with 4 elements (overdetermined for 3D direction)."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])

    frequency_hz = C / 0.3
    # 4 elements in non-degenerate configuration
    element_positions = np.array(
        [
            [0.1, 0.0, 0.0],
            [0.0, 0.1, 0.0],
            [0.0, 0.0, 0.1],
            [0.1, 0.1, 0.0],
        ],
    )

    direction = np.array([0.6, 0.7, np.sqrt(1 - 0.6**2 - 0.7**2)])
    source_xyz = collector_xyz + 100000.0 * direction

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
    aoa = aoas[0]

    expected_aoa = direction[:2]
    assert np.allclose(aoa, expected_aoa, atol=1e-5)


# Edge cases


def test_source_perpendicular_to_array_plane() -> None:
    """Source directly above array (perpendicular to XY plane)."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])

    frequency_hz = C / 0.3
    # Planar array in XY plane
    element_positions = np.array(
        [
            [0.1, 0.0, 0.0],
            [0.0, 0.1, 0.0],
        ],
    )

    # Source along +Z (perpendicular to array)
    direction = np.array([0.0, 0.0, 1.0])
    source_xyz = collector_xyz + 100000.0 * direction

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
    aoa = aoas[0]

    # All phases should be ~0 (source perpendicular to elements)
    # AOA should be [0, 0]
    assert np.allclose(aoa, [0.0, 0.0], atol=1e-6)


def test_minimum_elements_raises_error() -> None:
    """Test that single element raises ValueError."""
    pdoa = np.array([0.5])
    element_positions = np.array([[0.1, 0.0, 0.0]])
    frequency_hz = C / 0.3

    with pytest.raises(ValueError, match="at least 2 elements"):
        pdoa_to_aoa(pdoa, element_positions, frequency_hz)


# Input format flexibility


def test_accepts_list_inputs() -> None:
    """Verify list inputs are accepted."""
    pdoa = [0.5, 0.3]
    element_positions = [[0.1, 0.0, 0.0], [0.0, 0.1, 0.0]]
    frequency_hz = C / 0.3

    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)

    assert isinstance(aoas, np.ndarray)
    assert aoas.shape[1] == 2  # (K, 2) where K >= 1


# Wavelength scaling


def test_frequency_scaling_inverse() -> None:
    """Verify pdoa_to_aoa correctly accounts for frequency."""
    element_positions = np.array(
        [
            [0.1, 0.0, 0.0],
            [0.0, 0.1, 0.0],
        ],
    )

    # Same direction should give same AOA regardless of frequency
    # (as long as pdoa is scaled correctly)
    direction = np.array([0.6, 0.8, 0.0])

    # At frequency corresponding to 0.3m wavelength
    freq_1 = C / 0.3
    pdoa_1 = 2.0 * np.pi * freq_1 / C * (element_positions @ direction)
    aoas_1 = pdoa_to_aoa(pdoa_1, element_positions, freq_1)
    aoa_1 = aoas_1[0]

    # At frequency corresponding to 0.15m wavelength (double frequency)
    freq_2 = C / 0.15
    pdoa_2 = 2.0 * np.pi * freq_2 / C * (element_positions @ direction)
    aoas_2 = pdoa_to_aoa(pdoa_2, element_positions, freq_2)
    aoa_2 = aoas_2[0]

    # Both should give same direction
    assert np.allclose(aoa_1, aoa_2, atol=1e-10)
    assert np.allclose(aoa_1, direction[:2], atol=1e-10)


# Planar array geometry tests


@pytest.mark.parametrize(
    ("element_positions", "direction", "expected_indices"),
    [
        # X-Y planar array: returns X, Y direction cosines
        (
            np.array([[0.1, 0.0, 0.0], [0.0, 0.1, 0.0]]),
            np.array([0.6, 0.8, 0.0]),
            (0, 1),
        ),
        # X-Z planar array: returns X, Z direction cosines
        (
            np.array([[0.1, 0.0, 0.0], [0.0, 0.0, 0.1]]),
            np.array([0.6, 0.0, 0.8]),
            (0, 2),
        ),
        # Y-Z planar array: returns Y, Z direction cosines
        (
            np.array([[0.0, 0.1, 0.0], [0.0, 0.0, 0.1]]),
            np.array([0.0, 0.6, 0.8]),
            (1, 2),
        ),
    ],
    ids=["xy_plane", "xz_plane", "yz_plane"],
)
def test_planar_array_returns_correct_cosines(
    element_positions: np.ndarray,
    direction: np.ndarray,
    expected_indices: tuple[int, int],
) -> None:
    """Planar arrays return direction cosines for their two active axes."""
    frequency_hz = C / 0.3

    # Compute PDOA directly: phi_k = (2*pi*f/c) * dot(baseline_k, direction)
    k = 2 * np.pi * frequency_hz / C
    pdoa = k * (element_positions @ direction)

    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)
    aoa = aoas[0]

    # Verify we get the correct direction cosines for the active axes
    expected_aoa = direction[list(expected_indices)]
    assert np.allclose(aoa, expected_aoa, atol=1e-10)


# Ambiguity tests


def test_one_wrap_ambiguity_returns_two_candidates() -> None:
    """Array with d = lambda spacing returns exactly two valid AOA candidates.

    For element spacing d where lambda/2 <= d < 3*lambda/2, there is exactly
    one possible wrap per element, resulting in two valid AOA solutions:
    the true direction and one grating lobe ambiguity.
    """
    # Use wavelength = 0.3m, element spacing = 0.3m (d = lambda, one wrap)
    wavelength = 0.3
    frequency_hz = C / wavelength

    # Simple 2-element array along X axis with d = lambda spacing
    element_positions = np.array(
        [
            [wavelength, 0.0, 0.0],  # d = lambda along X
            [0.0, wavelength, 0.0],  # d = lambda along Y
        ],
    )

    # Source at 30 degrees off boresight in X direction
    # Direction cosines: [sin(30), 0, cos(30)] = [0.5, 0, 0.866]
    theta = np.radians(30)
    direction = np.array([np.sin(theta), 0.0, np.cos(theta)])

    # Compute PDOA
    k = 2 * np.pi * frequency_hz / C
    pdoa = k * (element_positions @ direction)

    # Get all candidate AOAs
    aoas = pdoa_to_aoa(pdoa, element_positions, frequency_hz)

    # Should have exactly 2 candidates for one-wrap ambiguity
    assert aoas.shape[0] == 2, f"Expected 2 candidates, got {aoas.shape[0]}"

    # The true direction should be among the candidates
    true_aoa = direction[:2]
    found_true = any(np.allclose(aoa, true_aoa, atol=1e-6) for aoa in aoas)
    assert found_true, f"True AOA {true_aoa} not found in candidates {aoas}"

    # Both candidates should be physically valid (|cx|, |cy| <= 1, cx^2+cy^2 <= 1)
    for aoa in aoas:
        assert np.all(np.abs(aoa) <= 1.0), f"Invalid direction cosine magnitude: {aoa}"
        assert np.sum(aoa**2) <= 1.0, f"Invalid direction vector magnitude: {aoa}"
