"""Tests for AOA computation with quaternion orientation.

Tests verify that get_aoa_quat correctly computes direction cosines from collector
to source in the collector's local reference frame.
"""

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

from gri_utils.conversion import xyz_enu_rotation
from gri_utils.observables import get_aoa_quat, get_aoa_quat_broadcast

# Realistic scenario: ground-based sensor observing targets


def test_sensor_observing_target_directly_above() -> None:
    """Sensor on equator observes target directly overhead."""
    # Sensor at sea level on equator, prime meridian
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    # Aircraft 10km directly above
    target_xyz = np.array([6388137.0, 0.0, 0.0])

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

    aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)

    # Target is straight up: East=0, North=0, Up=1
    # AOA returns only East, North components
    assert np.allclose(aoa, [0.0, 0.0], atol=1e-10)


def test_sensor_observing_target_to_east() -> None:
    """Sensor observes target due East at same altitude."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    # Target 1km to the East (positive Y in ECEF at this location)
    target_xyz = np.array([6378137.0, 1000.0, 0.0])

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

    aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)

    # Target is due East: direction cosines should be [1, 0]
    assert np.allclose(aoa, [1.0, 0.0], atol=1e-10)


def test_sensor_observing_target_to_north() -> None:
    """Sensor observes target due North at same altitude."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    # Target 1km to the North (positive Z in ECEF at this location)
    target_xyz = np.array([6378137.0, 0.0, 1000.0])

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

    aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)

    # Target is due North: direction cosines should be [0, 1]
    assert np.allclose(aoa, [0.0, 1.0], atol=1e-10)


def test_sensor_observing_target_northeast_elevated() -> None:
    """Sensor observes target at 45° azimuth with 30° elevation."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])

    # Target at 45° azimuth (NE), 30° elevation, 10km range
    range_m = 10000.0
    elevation = np.radians(30)
    azimuth = np.radians(45)

    # In ENU: East = range * cos(el) * sin(az)
    #         North = range * cos(el) * cos(az)
    #         Up = range * sin(el)
    east = range_m * np.cos(elevation) * np.sin(azimuth)
    north = range_m * np.cos(elevation) * np.cos(azimuth)
    up = range_m * np.sin(elevation)

    # Convert ENU offset to ECEF (at equator/prime meridian: E=+Y, N=+Z, U=+X)
    target_xyz = sensor_xyz + np.array([up, east, north])

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

    aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)

    # Expected direction cosines (unit vector components)
    expected_east = np.cos(elevation) * np.sin(azimuth)
    expected_north = np.cos(elevation) * np.cos(azimuth)

    assert np.allclose(aoa, [expected_east, expected_north], atol=1e-10)


@pytest.mark.parametrize(
    ("azimuth_deg", "expected_aoa"),
    [
        (0, (0.0, 1.0)),  # North
        (90, (1.0, 0.0)),  # East
        (180, (0.0, -1.0)),  # South
        (270, (-1.0, 0.0)),  # West
        (45, (np.sqrt(2) / 2, np.sqrt(2) / 2)),  # Northeast
    ],
)
def test_cardinal_and_intercardinal_directions(
    azimuth_deg: float,
    expected_aoa: tuple[float, float],
) -> None:
    """Verify AOA for targets at cardinal and intercardinal directions."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])

    # Target at given azimuth, 0° elevation (on horizon), 1km range
    azimuth = np.radians(azimuth_deg)
    east = 1000.0 * np.sin(azimuth)
    north = 1000.0 * np.cos(azimuth)

    target_xyz = sensor_xyz + np.array([0.0, east, north])

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

    aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)

    assert np.allclose(aoa, expected_aoa, atol=1e-10)


# Multiple target tracking scenario


def test_tracking_multiple_targets_simultaneously() -> None:
    """Sensor tracks 3 targets at different positions."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])

    # Three targets: overhead, to the east, to the north
    targets_xyz = np.array(
        [
            [6388137.0, 0.0, 0.0],  # 10km up
            [6378137.0, 5000.0, 0.0],  # 5km east
            [6378137.0, 0.0, 5000.0],  # 5km north
        ],
    )

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

    # Replicate sensor position and quaternion for each target
    sensors_xyz = np.tile(sensor_xyz, (3, 1))
    quats = np.tile(sensor_quat, (3, 1))

    aoas = get_aoa_quat(sensors_xyz, targets_xyz, quats)

    expected = np.array(
        [
            [0.0, 0.0],  # Overhead: E=0, N=0
            [1.0, 0.0],  # East: E=1, N=0
            [0.0, 1.0],  # North: E=0, N=1
        ],
    )

    assert aoas.shape == (3, 2)
    assert np.allclose(aoas, expected, atol=1e-10)


# Broadcast functionality for batch processing


def test_broadcast_single_sensor_multiple_targets() -> None:
    """Single sensor observing multiple targets (common surveillance scenario)."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])

    # Grid of targets at different azimuths
    n_targets = 8
    azimuths = np.linspace(0, 2 * np.pi, n_targets, endpoint=False)
    range_m = 10000.0

    targets_xyz = np.zeros((n_targets, 3))
    targets_xyz[:, 0] = sensor_xyz[0]  # Same altitude
    targets_xyz[:, 1] = sensor_xyz[1] + range_m * np.sin(azimuths)  # East offset
    targets_xyz[:, 2] = sensor_xyz[2] + range_m * np.cos(azimuths)  # North offset

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

    # Broadcast: single sensor to all targets
    aoas = get_aoa_quat_broadcast(sensor_xyz, targets_xyz, sensor_quat)

    assert aoas.shape == (n_targets, 2)
    # Each AOA should be a unit vector projection (E^2 + N^2 <= 1)
    magnitudes_squared = aoas[:, 0] ** 2 + aoas[:, 1] ** 2
    assert np.all(magnitudes_squared <= 1.0 + 1e-10)


def test_broadcast_multidimensional_target_grid() -> None:
    """Sensor observing a 2D grid of targets (e.g., radar scan)."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])

    # 4x5 grid of targets
    east_offsets = np.linspace(-2000, 2000, 4)
    north_offsets = np.linspace(-2500, 2500, 5)
    ee, nn = np.meshgrid(east_offsets, north_offsets, indexing="ij")

    targets_xyz = np.zeros((4, 5, 3))
    targets_xyz[..., 0] = sensor_xyz[0] + 1000.0  # 1km above
    targets_xyz[..., 1] = sensor_xyz[1] + ee
    targets_xyz[..., 2] = sensor_xyz[2] + nn

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

    aoas = get_aoa_quat_broadcast(sensor_xyz, targets_xyz, sensor_quat)

    assert aoas.shape == (4, 5, 2)


# Input format flexibility


def test_accepts_scipy_rotation_object() -> None:
    """Verify scipy Rotation objects work directly (common user pattern)."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    target_xyz = np.array([6378137.0, 1000.0, 0.0])

    rot_matrix = xyz_enu_rotation(sensor_xyz)
    rotation = Rotation.from_matrix(rot_matrix)

    # Should work with Rotation object directly
    aoa = get_aoa_quat(sensor_xyz, target_xyz, rotation)

    assert aoa.shape == (2,)
    assert np.allclose(aoa, [1.0, 0.0], atol=1e-10)


def test_accepts_list_inputs() -> None:
    """Verify list inputs are accepted (user convenience)."""
    sensor_xyz = [6378137.0, 0.0, 0.0]
    target_xyz = [6378137.0, 1000.0, 0.0]
    quat = [0.0, 0.0, 0.0, 1.0]  # Identity

    aoa = get_aoa_quat(sensor_xyz, target_xyz, quat)

    assert isinstance(aoa, np.ndarray)
    assert aoa.shape == (2,)
