import numpy as np
import pytest

from gri_utils.conversion import aer_to_xyz, get_aer, translate_aer, xyz_to_aer

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


@pytest.mark.parametrize(
    ("xyz_pos", "xyz_pos2", "expected"),
    [
        ((6378000, 0, 0), (6378100, 0, 0), (0, 90, 100)),
        ((6378000, 0, 0), (6378000, 100, 0), (90, 0, 100)),
        ((6378000, 0, 0), (6378000, 0, 100), (0, 0, 100)),
    ],
)
def test_get_aer(xyz_pos: XYZ, xyz_pos2: XYZ, expected: AER) -> None:
    """Test get_aer with array and sequence inputs."""
    # Array input
    aer = get_aer(np.array(xyz_pos), np.array(xyz_pos2))
    assert np.allclose(np.array(expected), aer), f"{aer}"

    # Sequence input
    aer_seq = get_aer(xyz_pos, xyz_pos2)
    assert aer_seq == pytest.approx(expected), f"{aer_seq}"


@pytest.mark.parametrize(
    ("xyz_pos", "aer_vec", "expected"),
    [
        ((6378000, 0, 0), (0, 90, 100), (6378100, 0, 0)),
        ((6378000, 0, 0), (90, 0, 100), (6378000, 100, 0)),
        ((6378000, 0, 0), (0, 0, 100), (6378000, 0, 100)),
    ],
)
def test_translate_aer(xyz_pos: XYZ, aer_vec: AER, expected: XYZ) -> None:
    """Test translate_aer with array and sequence inputs."""
    # Array input
    xyz = translate_aer(np.array(xyz_pos), np.array(aer_vec))
    assert np.allclose(np.array(expected), xyz), f"{xyz}"

    # Sequence input
    xyz_seq = translate_aer(xyz_pos, aer_vec)
    assert xyz_seq == pytest.approx(expected), f"{xyz_seq}"


@pytest.mark.parametrize(
    ("xyz_pos", "xyz_vec", "expected"),
    [
        ((6378000, 0, 0), (100, 0, 0), (0, 90, 100)),
        ((6378000, 0, 0), (0, 100, 0), (90, 0, 100)),
        ((6378000, 0, 0), (0, 0, 100), (0, 0, 100)),
    ],
)
def test_xyz_to_aer(xyz_pos: XYZ, xyz_vec: XYZ, expected: AER) -> None:
    """Test xyz_to_aer with array and sequence inputs."""
    # Array input
    aer = xyz_to_aer(np.array(xyz_pos), np.array(xyz_vec))
    assert np.allclose(np.array(expected), aer), f"{aer}"

    # Sequence input
    aer_seq = xyz_to_aer(xyz_pos, xyz_vec)
    assert aer_seq == pytest.approx(expected), f"{aer_seq}"


@pytest.mark.parametrize(
    ("xyz_pos", "aer_vec", "expected"),
    [
        ((6378000, 0, 0), (0, 90, 100), (100, 0, 0)),
        ((6378000, 0, 0), (90, 0, 100), (0, 100, 0)),
        ((6378000, 0, 0), (0, 0, 100), (0, 0, 100)),
    ],
)
def test_aer_to_xyz(xyz_pos: XYZ, aer_vec: AER, expected: XYZ) -> None:
    """Test aer_to_xyz with array and sequence inputs."""
    # Array input
    xyz = aer_to_xyz(np.array(xyz_pos), np.array(aer_vec))
    assert np.allclose(np.array(expected), xyz), f"{xyz}"

    # Sequence input
    xyz_seq = aer_to_xyz(xyz_pos, aer_vec)
    assert xyz_seq == pytest.approx(expected), f"{xyz_seq}"


# Azimuth wrapping tests (0->360)
@pytest.mark.parametrize(
    ("xyz_pos2", "expected"),
    [
        ((6378000, 0, -100), (180, 0, 100)),
        ((6378000, -100, 0), (270, 0, 100)),
    ],
)
def test_azim_wrap(xyz_pos2: XYZ, expected: AER) -> None:
    """Test azimuth wrapping to 0-360 range."""
    xyz_pos = (6378000, 0, 0)
    aer = get_aer(xyz_pos, xyz_pos2)
    assert aer == pytest.approx(expected), f"{aer}"


def test_azim_near_360() -> None:
    """Test azimuth near 360 degrees doesn't wrap to 0."""
    xyz_pos = (6378000, 0, 0)
    aer = get_aer(xyz_pos, (6378000, -1, 100))
    assert 359 < aer[0] < 360


# Elevation wrapping tests (-90->90)
@pytest.mark.parametrize(
    ("xyz_pos2", "expected"),
    [
        ((6378100, 0, 0), (0, 90, 100)),
        ((6377900, 0, 0), (0, -90, 100)),
        ((6377900, 0, -100), (180, -45, (100**2 * 2) ** 0.5)),
        ((6378100, 0, -100), (180, 45, (100**2 * 2) ** 0.5)),
    ],
)
def test_elev_wrap(xyz_pos2: XYZ, expected: AER) -> None:
    """Test elevation wrapping to -90 to 90 range."""
    xyz_pos = (6378000, 0, 0)
    aer = get_aer(xyz_pos, xyz_pos2)
    assert aer == pytest.approx(expected), f"{aer}"


# Input azimuth handling tests
@pytest.mark.parametrize(
    ("aer_vec", "expected"),
    [
        ((360, 0, 100), (6378000, 0, 100)),
        ((450, 0, 100), (6378000, 100, 0)),
        ((-270, 0, 100), (6378000, 100, 0)),
        ((270, 0, 100), (6378000, -100, 0)),
    ],
)
def test_input_azim(aer_vec: AER, expected: XYZ) -> None:
    """Test input azimuth values outside 0-360 range."""
    xyz_pos = (6378000, 0, 0)
    xyz = translate_aer(xyz_pos, aer_vec)
    assert xyz == pytest.approx(expected)


# Input elevation handling tests
@pytest.mark.parametrize(
    ("aer_vec", "expected"),
    [
        ((0, 360, 100), (6378000, 0, 100)),
        ((0, 180, 100), (6378000, 0, -100)),
        ((0, -180, 100), (6378000, 0, -100)),
        ((0, -90, 100), (6377900, 0, 0)),
    ],
)
def test_input_elev(aer_vec: AER, expected: XYZ) -> None:
    """Test input elevation values outside -90 to 90 range."""
    xyz_pos = (6378000, 0, 0)
    xyz = translate_aer(xyz_pos, aer_vec)
    assert xyz == pytest.approx(expected)


# Vectorized tests


def test_get_aer_vectorized() -> None:
    """Test get_aer with multiple target positions."""
    # Reference position on equator at prime meridian
    xyz_ref: XYZ = (6378000, 0, 0)

    # Multiple target positions: +100m in each direction
    xyz_targets = np.array(
        [
            (6378100, 0, 0),  # +100m radially out (Up, elev=90)
            (6378000, 100, 0),  # +100m in Y (East, az=90)
            (6378000, 0, 100),  # +100m in Z (North, az=0)
        ],
    )
    expected = np.array(
        [
            (0, 90, 100),  # Up: az=0, elev=90
            (90, 0, 100),  # East: az=90, elev=0
            (0, 0, 100),  # North: az=0, elev=0
        ],
    )

    aer_vecs = get_aer(xyz_ref, xyz_targets)

    assert aer_vecs.shape == (3, 3)
    assert np.allclose(aer_vecs, expected)


def test_translate_aer_vectorized() -> None:
    """Test translate_aer with multiple AER vectors."""
    xyz_ref: XYZ = (6378000, 0, 0)

    # Multiple AER vectors
    aer_vecs = np.array(
        [
            (0, 90, 100),  # Up -> radially out
            (90, 0, 100),  # East
            (0, 0, 100),  # North
        ],
    )
    expected = np.array(
        [
            (6378100, 0, 0),
            (6378000, 100, 0),
            (6378000, 0, 100),
        ],
    )

    xyzs = translate_aer(xyz_ref, aer_vecs)

    assert xyzs.shape == (3, 3)
    assert np.allclose(xyzs, expected)


def test_xyz_to_aer_vectorized() -> None:
    """Test xyz_to_aer with multiple XYZ vectors."""
    xyz_ref: XYZ = (6378000, 0, 0)

    # Use 4 vectors to avoid (3,3) matrix interpretation in ENU layer
    xyz_vecs = np.array(
        [
            (100, 0, 0),  # Radial out -> Up (elev=90)
            (0, 100, 0),  # Y -> East (az=90)
            (0, 0, 100),  # Z -> North (az=0)
            (0, 0, -100),  # -Z -> South (az=180)
        ],
    )
    expected = np.array(
        [
            (0, 90, 100),
            (90, 0, 100),
            (0, 0, 100),
            (180, 0, 100),
        ],
    )

    aer_vecs = xyz_to_aer(xyz_ref, xyz_vecs)

    assert aer_vecs.shape == (4, 3)
    assert np.allclose(aer_vecs, expected)


def test_aer_to_xyz_vectorized() -> None:
    """Test aer_to_xyz with multiple AER vectors."""
    xyz_ref: XYZ = (6378000, 0, 0)

    # Use 4 vectors to avoid (3,3) matrix interpretation in ENU layer
    aer_vecs = np.array(
        [
            (0, 90, 100),  # Up -> X (radial)
            (90, 0, 100),  # East -> Y
            (0, 0, 100),  # North -> Z
            (180, 0, 100),  # South -> -Z
        ],
    )
    expected = np.array(
        [
            (100, 0, 0),
            (0, 100, 0),
            (0, 0, 100),
            (0, 0, -100),
        ],
    )

    xyz_vecs = aer_to_xyz(xyz_ref, aer_vecs)

    assert xyz_vecs.shape == (4, 3)
    assert np.allclose(xyz_vecs, expected)


def test_aer_xyz_vectorized_round_trip() -> None:
    """Test round-trip AER<->XYZ conversion with batched inputs."""
    xyz_ref: XYZ = (6378000, 0, 0)

    aer_vecs = np.array(
        [
            (45, 30, 100),
            (90, 0, 100),
            (0, 45, 100),
            (180, -30, 100),
            (270, 60, 100),
        ],
    )

    xyz_vecs = aer_to_xyz(xyz_ref, aer_vecs)
    aer_back = xyz_to_aer(xyz_ref, xyz_vecs)

    assert aer_back.shape == aer_vecs.shape
    assert np.allclose(aer_back, aer_vecs)
