"""Tests for inverse AOA computation (AOA to XYZ direction).

Tests verify that aoa_quat_to_xyz correctly reconstructs the ECEF direction
vector from AOA direction cosines.
"""

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

from gri_utils.conversion import xyz_enu_rotation
from gri_utils.observables import (
    aoa_quat_to_xyz,
    aoa_quat_to_xyz_broadcast,
    get_aoa_quat,
    get_aoa_quat_broadcast,
)

# Round-trip verification: the primary use case


def test_round_trip_reconstructs_source_position() -> None:
    """Forward AOA + inverse AOA + range reconstructs source position."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    target_xyz = np.array([6378137.0, 1000.0, 500.0])

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

    # Forward: get AOA from positions
    aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)

    # Inverse: get direction vector from AOA
    direction_unit = aoa_quat_to_xyz(aoa, sensor_quat)

    # Reconstruct source using range
    range_m = np.linalg.norm(target_xyz - sensor_xyz)
    reconstructed = sensor_xyz + direction_unit * range_m

    assert np.allclose(reconstructed, target_xyz, atol=1e-10)


def test_round_trip_various_target_positions() -> None:
    """Round-trip works for targets in various directions."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    # Targets in different directions
    test_targets = [
        sensor_xyz + np.array([10000.0, 0.0, 0.0]),  # Up
        sensor_xyz + np.array([0.0, 5000.0, 0.0]),  # East
        sensor_xyz + np.array([0.0, 0.0, 5000.0]),  # North
        sensor_xyz + np.array([1000.0, 3000.0, 4000.0]),  # Diagonal
        sensor_xyz + np.array([500.0, -2000.0, 1500.0]),  # Another diagonal
    ]

    for target_xyz in test_targets:
        aoa = get_aoa_quat(sensor_xyz, target_xyz, sensor_quat)
        direction_unit = aoa_quat_to_xyz(aoa, sensor_quat)
        range_m = np.linalg.norm(target_xyz - sensor_xyz)
        reconstructed = sensor_xyz + direction_unit * range_m

        assert np.allclose(reconstructed, target_xyz, atol=1e-8)


# Direction vector properties


def test_output_is_unit_vector() -> None:
    """Returned direction vector has unit magnitude."""
    aoas = np.array(
        [
            [0.0, 0.0],  # Straight ahead (Z only)
            [1.0, 0.0],  # East
            [0.0, 1.0],  # North
            [0.5, 0.5],  # Diagonal
            [0.6, 0.8],  # Another direction
        ],
    )

    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity
    quats = np.tile(quat, (5, 1))

    directions = aoa_quat_to_xyz(aoas, quats)

    norms = np.linalg.norm(directions, axis=-1)
    assert np.allclose(norms, 1.0, atol=1e-10)


# Z-sign handling for targets in front/behind


def test_z_positive_default_for_targets_ahead() -> None:
    """Default z_positive=True for targets in front of sensor."""
    aoa = np.array([0.0, 0.0])  # Pure Z direction
    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity

    direction = aoa_quat_to_xyz(aoa, quat)

    # With identity rotation and z_positive=True, should point +Z
    assert np.allclose(direction, [0.0, 0.0, 1.0], atol=1e-10)


def test_z_negative_for_targets_behind() -> None:
    """z_positive=False for targets behind the sensor."""
    aoa = np.array([0.0, 0.0])  # Pure Z direction
    quat = np.array([0.0, 0.0, 0.0, 1.0])  # Identity

    direction_pos = aoa_quat_to_xyz(aoa, quat, z_positive=True)
    direction_neg = aoa_quat_to_xyz(aoa, quat, z_positive=False)

    assert np.allclose(direction_pos, [0.0, 0.0, 1.0], atol=1e-10)
    assert np.allclose(direction_neg, [0.0, 0.0, -1.0], atol=1e-10)


# Batch processing


def test_batch_round_trip_multiple_targets() -> None:
    """Round-trip works for batch of targets."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    # Multiple targets
    targets_xyz = np.array(
        [
            [6388137.0, 0.0, 0.0],
            [6378137.0, 5000.0, 0.0],
            [6378137.0, 0.0, 5000.0],
        ],
    )

    sensors_xyz = np.tile(sensor_xyz, (3, 1))
    quats = np.tile(sensor_quat, (3, 1))

    # Forward
    aoas = get_aoa_quat(sensors_xyz, targets_xyz, quats)

    # Inverse
    directions = aoa_quat_to_xyz(aoas, quats)

    # Reconstruct
    ranges = np.linalg.norm(targets_xyz - sensors_xyz, axis=-1, keepdims=True)
    reconstructed = sensors_xyz + directions * ranges

    assert np.allclose(reconstructed, targets_xyz, atol=1e-8)


# Broadcast functionality


def test_broadcast_single_inputs() -> None:
    """Broadcast wrapper matches non-broadcast for single inputs."""
    aoa = np.array([0.5, 0.3])
    quat = np.array([0.0, 0.0, 0.0, 1.0])

    direction_ref = aoa_quat_to_xyz(aoa, quat)
    direction_bc = aoa_quat_to_xyz_broadcast(aoa, quat)

    assert direction_bc.shape == (3,)
    assert np.allclose(direction_bc, direction_ref, atol=1e-10)


def test_broadcast_multidimensional_aoas() -> None:
    """Broadcast over multidimensional AOA array."""
    # 2x3 grid of AOAs
    aoas = np.zeros((2, 3, 2))
    aoas[0, 0] = [1.0, 0.0]  # East
    aoas[0, 1] = [0.0, 1.0]  # North
    aoas[0, 2] = [0.5, 0.5]
    aoas[1, 0] = [-1.0, 0.0]  # West
    aoas[1, 1] = [0.0, -1.0]  # South
    aoas[1, 2] = [-0.5, -0.5]

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

    directions = aoa_quat_to_xyz_broadcast(aoas, quat)

    assert directions.shape == (2, 3, 3)

    # All should be unit vectors
    norms = np.linalg.norm(directions, axis=-1)
    assert np.allclose(norms, 1.0, atol=1e-10)


def test_broadcast_round_trip() -> None:
    """Full round-trip using broadcast functions."""
    sensor_xyz = np.array([6378137.0, 0.0, 0.0])
    rot_matrix = xyz_enu_rotation(sensor_xyz)
    sensor_quat = Rotation.from_matrix(rot_matrix).as_quat()

    # Random targets (all in front of sensor for z_positive=True)
    rng = np.random.default_rng(42)
    targets_xyz = np.zeros((3, 4, 3))
    targets_xyz[..., 0] = sensor_xyz[0] + rng.uniform(100, 1000, (3, 4))
    targets_xyz[..., 1] = sensor_xyz[1] + rng.uniform(-500, 500, (3, 4))
    targets_xyz[..., 2] = sensor_xyz[2] + rng.uniform(-500, 500, (3, 4))

    # Forward
    aoas = get_aoa_quat_broadcast(sensor_xyz, targets_xyz, sensor_quat)

    # Inverse
    directions = aoa_quat_to_xyz_broadcast(aoas, sensor_quat)

    # Reconstruct
    ranges = np.linalg.norm(targets_xyz - sensor_xyz, axis=-1, keepdims=True)
    reconstructed = sensor_xyz + directions * ranges

    assert np.allclose(reconstructed, targets_xyz, atol=1e-8)


# Input format flexibility


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

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

    aoa = get_aoa_quat(sensor_xyz, target_xyz, rotation)

    direction_from_rotation = aoa_quat_to_xyz(aoa, rotation)
    direction_from_quat = aoa_quat_to_xyz(aoa, quat)

    assert np.allclose(direction_from_rotation, direction_from_quat, atol=1e-10)
