import math

import numpy as np
import pytest

from gri_utils.observables import (
    aoa_to_enu,
    aoa_to_xyz,
    get_aoa_enu,
    translate_aoa,
    xyz_to_aoa,
)

# Type aliases for readability
XYZ = tuple[float, float, float]
AOA = tuple[float, float]
ENU = tuple[float, float, float]

SQRT2_INV = 1.0 / math.sqrt(2)

# Reference collector position: on equator at prime meridian
# At this location: +X is Up, +Y is East, +Z is North
COLLECTOR: XYZ = (6378000.0, 0.0, 0.0)


# Test get_aoa_enu


@pytest.mark.parametrize(
    ("source_offset", "expected"),
    [
        ((0.0, 100.0, 0.0), (1.0, 0.0)),  # East
        ((0.0, -100.0, 0.0), (-1.0, 0.0)),  # West
        ((0.0, 0.0, 100.0), (0.0, 1.0)),  # North
        ((0.0, 0.0, -100.0), (0.0, -1.0)),  # South
        ((100.0, 0.0, 0.0), (0.0, 0.0)),  # Up
        ((0.0, 100.0, 100.0), (SQRT2_INV, SQRT2_INV)),  # 45° E-N
    ],
)
def test_get_aoa_enu(source_offset: XYZ, expected: AOA) -> None:
    """Test get_aoa_enu with array and sequence inputs."""
    collector = np.array(COLLECTOR)
    source = collector + np.array(source_offset)

    # Array input
    aoa = get_aoa_enu(collector, source)
    assert aoa.shape == (2,)
    assert np.allclose(aoa, expected, atol=1e-10)

    # Sequence input
    aoa_seq = get_aoa_enu(COLLECTOR, tuple(source))
    assert np.allclose(aoa_seq, expected, atol=1e-10)


# Test xyz_to_aoa


@pytest.mark.parametrize(
    ("xyz_vec", "expected"),
    [
        ((0.0, 100.0, 0.0), (1.0, 0.0)),  # +Y -> East
        ((0.0, 0.0, 100.0), (0.0, 1.0)),  # +Z -> North
        ((100.0, 0.0, 0.0), (0.0, 0.0)),  # +X -> Up
        ((0.0, 100.0, 100.0), (SQRT2_INV, SQRT2_INV)),  # Diagonal
    ],
)
def test_xyz_to_aoa(xyz_vec: XYZ, expected: AOA) -> None:
    """Test xyz_to_aoa with array and sequence inputs."""
    # Array input
    aoa = xyz_to_aoa(np.array(COLLECTOR), np.array(xyz_vec))
    assert aoa.shape == (2,)
    assert np.allclose(aoa, expected, atol=1e-10)

    # Sequence input
    aoa_seq = xyz_to_aoa(COLLECTOR, xyz_vec)
    assert np.allclose(aoa_seq, expected, atol=1e-10)


# Test aoa_to_enu


@pytest.mark.parametrize(
    ("aoa", "expected"),
    [
        ((1.0, 0.0), (1.0, 0.0, 0.0)),  # East
        ((0.0, 1.0), (0.0, 1.0, 0.0)),  # North
        ((0.0, 0.0), (0.0, 0.0, 1.0)),  # Up
        ((SQRT2_INV, SQRT2_INV), (SQRT2_INV, SQRT2_INV, 0.0)),  # 45° E-N
    ],
)
def test_aoa_to_enu(aoa: AOA, expected: ENU) -> None:
    """Test aoa_to_enu with array and sequence inputs."""
    # Array input
    enu = aoa_to_enu(np.array(aoa))
    assert enu.shape == (3,)
    assert np.allclose(enu, expected, atol=1e-7)

    # Sequence input
    enu_seq = aoa_to_enu(aoa)
    assert np.allclose(enu_seq, expected, atol=1e-7)


def test_aoa_to_enu_z_negative() -> None:
    """Test aoa_to_enu with z_positive=False."""
    aoa = (0.0, 0.0)
    enu = aoa_to_enu(aoa, z_positive=False)
    assert np.allclose(enu, [0.0, 0.0, -1.0], atol=1e-10)


# Test aoa_to_xyz


@pytest.mark.parametrize(
    ("aoa", "expected"),
    [
        ((1.0, 0.0), (0.0, 1.0, 0.0)),  # East -> +Y
        ((0.0, 1.0), (0.0, 0.0, 1.0)),  # North -> +Z
        ((0.0, 0.0), (1.0, 0.0, 0.0)),  # Up -> +X
    ],
)
def test_aoa_to_xyz(aoa: AOA, expected: XYZ) -> None:
    """Test aoa_to_xyz with array and sequence inputs."""
    # Array input
    xyz_unit = aoa_to_xyz(np.array(COLLECTOR), np.array(aoa))
    assert xyz_unit.shape == (3,)
    assert np.allclose(xyz_unit, expected, atol=1e-10)

    # Sequence input
    xyz_seq = aoa_to_xyz(COLLECTOR, aoa)
    assert np.allclose(xyz_seq, expected, atol=1e-10)


def test_aoa_to_xyz_z_negative() -> None:
    """Test aoa_to_xyz with z_positive=False."""
    aoa = (0.0, 0.0)
    xyz_unit = aoa_to_xyz(COLLECTOR, aoa, z_positive=False)
    # Down at this location is -X in ECEF
    assert np.allclose(xyz_unit, [-1.0, 0.0, 0.0], atol=1e-10)


# Test translate_aoa


@pytest.mark.parametrize(
    ("aoa", "range_m", "expected"),
    [
        ((1.0, 0.0), 100.0, (6378000.0, 100.0, 0.0)),  # East
        ((0.0, 1.0), 200.0, (6378000.0, 0.0, 200.0)),  # North
        ((0.0, 0.0), 100.0, (6378100.0, 0.0, 0.0)),  # Up
    ],
)
def test_translate_aoa(aoa: AOA, range_m: float, expected: XYZ) -> None:
    """Test translate_aoa with array and sequence inputs."""
    # Array input
    source = translate_aoa(np.array(COLLECTOR), np.array(aoa), range_m)
    assert source.shape == (3,)
    assert np.allclose(source, expected, atol=1e-10)

    # Sequence input
    source_seq = translate_aoa(COLLECTOR, aoa, range_m)
    assert np.allclose(source_seq, expected, atol=1e-10)


def test_translate_aoa_z_negative() -> None:
    """Test translate_aoa with z_positive=False."""
    aoa = (0.0, 0.0)
    range_m = 100.0
    source = translate_aoa(COLLECTOR, aoa, range_m, z_positive=False)
    # Down at this location is -X in ECEF
    assert np.allclose(source, [6377900.0, 0.0, 0.0], atol=1e-10)


# Vectorized tests


def test_get_aoa_enu_vectorized() -> None:
    """Test get_aoa_enu with multiple sources."""
    collector = np.array(COLLECTOR)
    sources = np.array(
        [
            [6378100.0, 0.0, 0.0],  # Up
            [6378000.0, 100.0, 0.0],  # East
            [6378000.0, 0.0, 100.0],  # North
        ],
    )
    expected = np.array(
        [
            [0.0, 0.0],  # Up
            [1.0, 0.0],  # East
            [0.0, 1.0],  # North
        ],
    )

    aoa = get_aoa_enu(collector, sources)

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


def test_xyz_to_aoa_vectorized() -> None:
    """Test xyz_to_aoa with multiple XYZ vectors."""
    xyz_vecs = np.array(
        [
            [100.0, 0.0, 0.0],  # +X -> Up
            [0.0, 100.0, 0.0],  # +Y -> East
            [0.0, 0.0, 100.0],  # +Z -> North
            [0.0, 100.0, 100.0],  # Diagonal
        ],
    )
    expected = np.array(
        [
            [0.0, 0.0],
            [1.0, 0.0],
            [0.0, 1.0],
            [SQRT2_INV, SQRT2_INV],
        ],
    )

    aoa = xyz_to_aoa(COLLECTOR, xyz_vecs)

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


def test_aoa_to_enu_vectorized() -> None:
    """Test aoa_to_enu with multiple AOAs."""
    aoas = np.array(
        [
            [1.0, 0.0],  # East
            [0.0, 1.0],  # North
            [0.0, 0.0],  # Up
        ],
    )
    expected = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
        ],
    )

    enus = aoa_to_enu(aoas)

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


def test_aoa_to_xyz_vectorized() -> None:
    """Test aoa_to_xyz with multiple AOAs."""
    aoas = np.array(
        [
            [1.0, 0.0],  # East -> +Y
            [0.0, 1.0],  # North -> +Z
            [0.0, 0.0],  # Up -> +X
        ],
    )
    expected = np.array(
        [
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
            [1.0, 0.0, 0.0],
        ],
    )

    xyz_units = aoa_to_xyz(COLLECTOR, aoas)

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


def test_translate_aoa_vectorized() -> None:
    """Test translate_aoa with multiple AOAs and ranges."""
    aoas = np.array(
        [
            [1.0, 0.0],  # East
            [0.0, 1.0],  # North
        ],
    )
    ranges = np.array([100.0, 200.0])
    expected = np.array(
        [
            [6378000.0, 100.0, 0.0],
            [6378000.0, 0.0, 200.0],
        ],
    )

    sources = translate_aoa(COLLECTOR, aoas, ranges)

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


def test_translate_aoa_scalar_range() -> None:
    """Test translate_aoa with scalar range for multiple AOAs."""
    aoas = np.array(
        [
            [1.0, 0.0],  # East
            [0.0, 1.0],  # North
        ],
    )
    range_m = 100.0
    expected = np.array(
        [
            [6378000.0, 100.0, 0.0],
            [6378000.0, 0.0, 100.0],
        ],
    )

    sources = translate_aoa(COLLECTOR, aoas, range_m)

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


# Round-trip consistency tests


@pytest.mark.parametrize(
    "source_offset",
    [
        (0.0, 100.0, 0.0),  # East
        (0.0, -100.0, 0.0),  # West
        (0.0, 0.0, 100.0),  # North
        (0.0, 0.0, -100.0),  # South
        (100.0, 0.0, 0.0),  # Up
        (50.0, 75.0, 100.0),  # Arbitrary
    ],
)
def test_get_aoa_enu_translate_aoa_round_trip(source_offset: XYZ) -> None:
    """Test round-trip: source -> get_aoa_enu -> translate_aoa -> source."""
    collector = np.array(COLLECTOR)
    source = collector + np.array(source_offset)

    aoa = get_aoa_enu(collector, source)
    range_m = np.linalg.norm(source - collector)
    source_recovered = translate_aoa(collector, aoa, range_m)

    assert np.allclose(source_recovered, source, atol=1e-10)


@pytest.mark.parametrize(
    "xyz_vec",
    [
        (0.0, 100.0, 0.0),  # +Y
        (0.0, 0.0, 100.0),  # +Z
        (100.0, 0.0, 0.0),  # +X
        (50.0, 75.0, 100.0),  # Arbitrary
    ],
)
def test_xyz_to_aoa_aoa_to_xyz_round_trip(xyz_vec: XYZ) -> None:
    """Test round-trip: xyz_to_aoa -> aoa_to_xyz recovers normalized vector."""
    xyz_vec_arr = np.array(xyz_vec)

    aoa = xyz_to_aoa(COLLECTOR, xyz_vec_arr)
    xyz_unit = aoa_to_xyz(COLLECTOR, aoa)

    expected = xyz_vec_arr / np.linalg.norm(xyz_vec_arr)
    assert np.allclose(xyz_unit, expected, atol=1e-10)


def test_vectorized_round_trip() -> None:
    """Test round-trip with batched inputs.

    Note: All source offsets must have positive Up component (positive X at
    this collector location) since translate_aoa defaults to z_positive=True.
    """
    collector = np.array(COLLECTOR)
    source_offsets = np.array(
        [
            [50.0, 30.0, 20.0],
            [100.0, 0.0, 0.0],
            [0.0, 100.0, 0.0],
            [0.0, 0.0, 100.0],
            [50.0, -30.0, -20.0],  # Still above horizon (positive X = Up)
        ],
    )
    sources = collector + source_offsets
    ranges = np.linalg.norm(source_offsets, axis=1)

    aoas = get_aoa_enu(collector, sources)
    sources_back = translate_aoa(collector, aoas, ranges)

    assert sources_back.shape == sources.shape
    assert np.allclose(sources_back, sources, atol=1e-10)
