"""Tests for PDOA (Phase Difference of Arrival) observable computation.

Tests verify that get_pdoa correctly computes phase differences for
interferometer array elements based on signal direction.
"""

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, get_pdoa_broadcast, get_pdoa_residual

# Basic physics tests with simple geometry


def test_source_along_positive_x_axis_of_array() -> None:
    """Source along array +X axis produces expected phase differences."""
    # Sensor at equator, identity rotation (array frame = ECEF)
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    # Source 10km along +X in ECEF (which is +X in array frame with identity)
    source_xyz = np.array([6388137.0, 0.0, 0.0])

    # Identity quaternion: array frame aligned with ECEF
    quat = np.array([0.0, 0.0, 0.0, 1.0])

    # Single element at (1, 0, 0) in array frame
    wavelength = 1.0  # 1 meter wavelength for easy math
    frequency_hz = C / wavelength
    element_positions = np.array([[1.0, 0.0, 0.0]])

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)

    # Direction to source is +X, so dot(element_pos, direction) = 1.0
    # Phase = 2*pi*f/c * 1.0 = 2*pi (when wavelength = 1m)
    expected = 2.0 * np.pi
    assert np.allclose(pdoa, [expected], atol=1e-10)


def test_source_perpendicular_to_element_gives_zero_phase() -> None:
    """Source perpendicular to element position vector gives zero phase."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    # Source along +X direction
    source_xyz = np.array([6388137.0, 0.0, 0.0])

    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity

    # Element along Y axis - perpendicular to source direction
    wavelength = 1.0
    frequency_hz = C / wavelength
    element_positions = np.array([[0.0, 1.0, 0.0]])

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)

    # dot(Y_element, X_direction) = 0
    assert np.allclose(pdoa, [0.0], atol=1e-10)


def test_multiple_elements_linear_array() -> None:
    """Linear array along X produces linearly increasing phase."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6388137.0, 0.0, 0.0])  # +X direction

    quat = np.array([0.0, 0.0, 0.0, 1.0])

    wavelength = 1.0
    frequency_hz = C / wavelength
    # Elements at 0.25, 0.5, 0.75 meters along X
    element_positions = np.array(
        [
            [0.25, 0.0, 0.0],
            [0.5, 0.0, 0.0],
            [0.75, 0.0, 0.0],
        ],
    )

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)

    # Expected phases: 2*pi * [0.25, 0.5, 0.75]
    expected = 2.0 * np.pi * np.array([0.25, 0.5, 0.75])
    assert np.allclose(pdoa, expected, atol=1e-10)


# Test with realistic ENU frame


def test_source_to_east_with_enu_frame() -> None:
    """Source to East produces phase on East-pointing elements."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    # Target 10km to the East (positive Y in ECEF at equator)
    source_xyz = np.array([6378137.0, 10000.0, 0.0])

    # Sensor oriented in local ENU frame
    rot_matrix = xyz_enu_rotation(collector_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    wavelength = 0.1  # 10 cm = 3 GHz signal
    frequency_hz = C / wavelength
    # Element 0.05m (half wavelength) along East direction in array frame
    # In ENU frame: East = X, North = Y, Up = Z
    element_positions = np.array([[0.05, 0.0, 0.0]])  # East-pointing

    pdoa = get_pdoa(
        collector_xyz,
        source_xyz,
        sensor_quat,
        element_positions,
        frequency_hz,
    )

    # Direction is pure East in ENU, element is along East
    # dot product = 0.05, phase = 2*pi * 0.05 / 0.1 = pi
    expected = np.pi
    assert np.allclose(pdoa, [expected], atol=1e-10)


def test_source_to_north_with_enu_frame() -> None:
    """Source to North produces phase on North-pointing elements."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    # Target 10km to the North (positive Z in ECEF at equator)
    source_xyz = np.array([6378137.0, 0.0, 10000.0])

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

    wavelength = 0.1
    frequency_hz = C / wavelength
    # Element along North (Y in ENU)
    element_positions = np.array([[0.0, 0.05, 0.0]])  # North-pointing

    pdoa = get_pdoa(
        collector_xyz,
        source_xyz,
        sensor_quat,
        element_positions,
        frequency_hz,
    )

    expected = np.pi  # Same calculation as East case
    assert np.allclose(pdoa, [expected], atol=1e-10)


# Consistency with ULA steering formula


def test_consistency_with_ula_formula() -> None:
    """Verify PDOA matches ULA steering vector formula for 1D array.

    ULA formula: phase[k] = 2*pi * k * d / lambda * sin(theta)
    where theta is angle from broadside.
    """
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    # Array along X axis in array frame, source at 30 degrees from broadside
    # Broadside is perpendicular to array axis (Y direction)
    # 30 degrees toward +X means direction = [sin(30), cos(30), 0]
    theta_deg = 30.0
    theta = np.radians(theta_deg)
    direction = np.array([np.sin(theta), np.cos(theta), 0.0])

    # Put source far away in that direction (using identity rotation)
    source_xyz = collector_xyz + 100000.0 * direction

    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity

    d = 0.15  # Element spacing
    wavelength = 0.3  # lambda = 0.3m (1 GHz)
    frequency_hz = C / wavelength

    # ULA with 3 elements at positions d, 2d, 3d along X
    element_positions = np.array(
        [
            [d, 0.0, 0.0],
            [2 * d, 0.0, 0.0],
            [3 * d, 0.0, 0.0],
        ],
    )

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)

    # ULA formula: phase[k] = 2*pi * (k+1) * d / lambda * sin(theta)
    # where k=0,1,2 for our elements at d, 2d, 3d
    k = np.array([1, 2, 3])
    expected = 2.0 * np.pi * k * d / wavelength * np.sin(theta)

    assert np.allclose(pdoa, expected, atol=1e-10)


# Batch processing tests


def test_multiple_sources() -> None:
    """Compute PDOA for multiple source positions."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    # Three sources along different directions
    sources_xyz = np.array(
        [
            [6388137.0, 0.0, 0.0],  # +X
            [6378137.0, 10000.0, 0.0],  # +Y
            [6378137.0, 0.0, 10000.0],  # +Z
        ],
    )

    quat = np.array([0.0, 0.0, 0.0, 1.0])
    wavelength = 1.0
    frequency_hz = C / wavelength

    # Single element along X
    element_positions = np.array([[0.5, 0.0, 0.0]])

    pdoa = get_pdoa(collector_xyz, sources_xyz, quat, element_positions, frequency_hz)

    assert pdoa.shape == (3, 1)
    # Source +X: phase = 2*pi*0.5 = pi
    # Source +Y: phase = 0 (perpendicular)
    # Source +Z: phase = 0 (perpendicular)
    expected = np.array([[np.pi], [0.0], [0.0]])
    assert np.allclose(pdoa, expected, atol=1e-10)


def test_broadcast_single_collector_multiple_sources() -> None:
    """Broadcast with single collector observing multiple sources."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    # Grid of sources
    n_sources = 5
    sources_xyz = np.zeros((n_sources, 3))
    sources_xyz[:, 0] = collector_xyz[0] + 10000.0  # All 10km along +X

    quat = np.array([0.0, 0.0, 0.0, 1.0])
    wavelength = 1.0
    frequency_hz = C / wavelength
    element_positions = np.array([[0.25, 0.0, 0.0], [0.0, 0.25, 0.0]])

    pdoa = get_pdoa_broadcast(
        collector_xyz,
        sources_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    assert pdoa.shape == (n_sources, 2)


def test_broadcast_multidimensional() -> None:
    """Broadcast with 2D grid of sources."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])

    # 3x4 grid of sources
    sources_xyz = np.zeros((3, 4, 3))
    sources_xyz[..., 0] = collector_xyz[0] + 10000.0

    quat = np.array([0.0, 0.0, 0.0, 1.0])
    wavelength = 1.0
    frequency_hz = C / wavelength
    element_positions = np.array([[0.25, 0.0, 0.0]])

    pdoa = get_pdoa_broadcast(
        collector_xyz,
        sources_xyz,
        quat,
        element_positions,
        frequency_hz,
    )

    assert pdoa.shape == (3, 4, 1)


# Input format flexibility


def test_accepts_list_inputs() -> None:
    """Verify list inputs are accepted."""
    collector_xyz = [6378137.0, 0.0, 0.0]
    source_xyz = [6388137.0, 0.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]
    element_positions = [[0.5, 0.0, 0.0]]
    wavelength = 1.0
    frequency_hz = C / wavelength

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)

    assert isinstance(pdoa, np.ndarray)
    assert pdoa.shape == (1,)


def test_accepts_scipy_rotation() -> None:
    """Verify scipy Rotation objects work directly."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6388137.0, 0.0, 0.0])

    rotation = Rotation.identity()
    element_positions = np.array([[0.5, 0.0, 0.0]])
    wavelength = 1.0
    frequency_hz = C / wavelength

    pdoa = get_pdoa(
        collector_xyz,
        source_xyz,
        rotation,
        element_positions,
        frequency_hz,
    )

    assert pdoa.shape == (1,)


# Edge cases


def test_frequency_scaling() -> None:
    """Phase scales proportionally with frequency."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6388137.0, 0.0, 0.0])
    quat = np.array([0.0, 0.0, 0.0, 1.0])
    element_positions = np.array([[1.0, 0.0, 0.0]])

    freq_1 = C / 1.0  # 1m wavelength
    freq_2 = C / 2.0  # 2m wavelength

    pdoa_1m = get_pdoa(collector_xyz, source_xyz, quat, element_positions, freq_1)
    pdoa_2m = get_pdoa(collector_xyz, source_xyz, quat, element_positions, freq_2)

    # Doubling frequency (halving wavelength) doubles phase
    assert np.allclose(pdoa_1m, 2.0 * pdoa_2m, atol=1e-10)


@pytest.mark.parametrize(
    ("direction", "element", "expected_sign"),
    [
        ([1, 0, 0], [1, 0, 0], 1),  # Same direction: positive phase
        ([1, 0, 0], [-1, 0, 0], -1),  # Opposite direction: negative phase
        ([0, 1, 0], [1, 0, 0], 0),  # Perpendicular: zero phase
    ],
)
def test_phase_sign_convention(
    direction: list[float],
    element: list[float],
    expected_sign: int,
) -> None:
    """Verify phase sign matches dot product of direction and element position."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = collector_xyz + 10000.0 * np.array(direction)
    quat = np.array([0.0, 0.0, 0.0, 1.0])
    element_positions = np.array([element])
    wavelength = 1.0
    frequency_hz = C / wavelength

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)

    if expected_sign == 0:
        assert np.allclose(pdoa, 0.0, atol=1e-10)
    else:
        assert np.sign(pdoa[0]) == expected_sign


# Residual tests


def test_residual_zero_at_true_source() -> None:
    """Residual is near zero at the true source position."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6388137.0, 0.0, 0.0])

    quat = np.array([0.0, 0.0, 0.0, 1.0])
    wavelength = 0.3
    frequency_hz = C / wavelength
    element_positions = np.array(
        [[0.1, 0.0, 0.0], [0.0, 0.1, 0.0]],
    )

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    residual_fn = get_pdoa_residual(
        collector_xyz,
        quat,
        element_positions,
        frequency_hz,
        pdoa,
    )

    result = residual_fn(source_xyz)
    assert np.abs(result) < 1e-6


def test_residual_vectorized() -> None:
    """Residual accepts batch of points and returns correct shapes."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6388137.0, 0.0, 0.0])

    quat = np.array([0.0, 0.0, 0.0, 1.0])
    wavelength = 0.3
    frequency_hz = C / wavelength
    element_positions = np.array(
        [[0.1, 0.0, 0.0], [0.0, 0.1, 0.0]],
    )

    pdoa = get_pdoa(collector_xyz, source_xyz, quat, element_positions, frequency_hz)
    residual_fn = get_pdoa_residual(
        collector_xyz,
        quat,
        element_positions,
        frequency_hz,
        pdoa,
    )

    # Batch of points: true source and an offset point
    points = np.array(
        [source_xyz, source_xyz + np.array([0, 1000, 0])],
    )
    result = residual_fn(points)

    assert result.shape == (2,)
    assert result[0] < 1e-6
    assert result[1] > result[0]


def test_residual_with_enu_frame() -> None:
    """Residual works with ENU-oriented collector."""
    collector_xyz = np.array([6378137.0, 0.0, 0.0])
    source_xyz = np.array([6378137.0, 10000.0, 0.0])

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

    wavelength = 0.1
    frequency_hz = C / wavelength
    element_positions = np.array(
        [[0.04, 0.0, 0.0], [0.0, 0.04, 0.0]],
    )

    pdoa = get_pdoa(
        collector_xyz,
        source_xyz,
        sensor_quat,
        element_positions,
        frequency_hz,
    )
    residual_fn = get_pdoa_residual(
        collector_xyz,
        sensor_quat,
        element_positions,
        frequency_hz,
        pdoa,
    )

    result = residual_fn(source_xyz)
    assert np.abs(result) < 1e-6
