import numpy as np
import pytest

from gri_utils import constants
from gri_utils.conversion import lla_to_xyz, xyz_to_lla

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

# Test data: Munich, Germany
_LLA_MUNICH: LLA = (48.13743, 11.57549, 550)
_XYZ_MUNICH: XYZ = (4177977.2298, 855753.1594, 4727497.3435)

# North pole
_LLA_NORTH_POLE: LLA = (90, 0, 0)
_XYZ_NORTH_POLE: XYZ = (0, 0, 6356752.3142)

# South pole
_LLA_SOUTH_POLE: LLA = (-90, 0, 0)
_XYZ_SOUTH_POLE: XYZ = (0, 0, -6356752.3142)

# Near north pole (89.9999 degrees) - avoids singularity
_LLA_NEAR_NORTH_POLE: LLA = (89.9999, 0, 0)

# Equator at prime meridian
_LLA_EQUATOR_PRIME: LLA = (0, 0, 0)
_XYZ_EQUATOR_PRIME: XYZ = (6378137.0, 0, 0)

# Equator at international date line (lon=180)
_LLA_EQUATOR_DATE_LINE: LLA = (0, 180, 0)
_XYZ_EQUATOR_DATE_LINE: XYZ = (-6378137.0, 0, 0)


@pytest.mark.parametrize(
    ("lla", "expected"),
    [
        # Standard location (Munich)
        (_LLA_MUNICH, _XYZ_MUNICH),
        # Equator at prime meridian
        (_LLA_EQUATOR_PRIME, _XYZ_EQUATOR_PRIME),
        # Equator at date line
        (_LLA_EQUATOR_DATE_LINE, _XYZ_EQUATOR_DATE_LINE),
        # North pole
        (_LLA_NORTH_POLE, _XYZ_NORTH_POLE),
        # South pole
        (_LLA_SOUTH_POLE, _XYZ_SOUTH_POLE),
    ],
)
def test_lla_to_xyz(lla: LLA, expected: XYZ) -> None:
    """Test LLA to XYZ conversion with array and sequence inputs."""
    # Array input
    xyz = lla_to_xyz(np.array(lla))
    assert np.allclose(xyz, expected, rtol=1e-6), f"{xyz}"

    # Sequence input
    xyz_seq = lla_to_xyz(lla)
    assert np.allclose(xyz_seq, expected, rtol=1e-6), f"{xyz_seq}"


@pytest.mark.parametrize(
    ("xyz", "expected"),
    [
        # Standard location (Munich)
        (_XYZ_MUNICH, _LLA_MUNICH),
        # Equator at prime meridian
        (_XYZ_EQUATOR_PRIME, _LLA_EQUATOR_PRIME),
        # Date line - atan2 returns pi, so lon=180
        (_XYZ_EQUATOR_DATE_LINE, _LLA_EQUATOR_DATE_LINE),
        # Poles (singularity now handled with simplified polar geometry)
        (_XYZ_NORTH_POLE, _LLA_NORTH_POLE),
        (_XYZ_SOUTH_POLE, _LLA_SOUTH_POLE),
    ],
)
def test_xyz_to_lla(xyz: XYZ, expected: LLA) -> None:
    """Test XYZ to LLA conversion with array and sequence inputs."""
    # Array input
    lla = xyz_to_lla(np.array(xyz))
    # Use atol for altitude since rtol fails when expected is 0
    assert np.allclose(lla[:2], expected[:2], rtol=1e-6), f"{lla}"
    assert np.isclose(lla[2], expected[2], atol=1e-4), f"alt: {lla[2]} vs {expected[2]}"

    # Sequence input
    lla_seq = xyz_to_lla(xyz)
    assert np.allclose(lla_seq[:2], expected[:2], rtol=1e-6), f"{lla_seq}"
    assert np.isclose(lla_seq[2], expected[2], atol=1e-4), f"alt: {lla_seq[2]}"


@pytest.mark.parametrize(
    ("lla", "xyz"),
    [
        (_LLA_MUNICH, _XYZ_MUNICH),
        (_LLA_EQUATOR_PRIME, _XYZ_EQUATOR_PRIME),
        (_LLA_NORTH_POLE, _XYZ_NORTH_POLE),
        (_LLA_SOUTH_POLE, _XYZ_SOUTH_POLE),
    ],
)
def test_round_trip(lla: LLA, xyz: XYZ) -> None:
    """Test round-trip conversion LLA->XYZ->LLA and XYZ->LLA->XYZ."""
    assert np.allclose(xyz_to_lla(lla_to_xyz(lla)), lla, rtol=1e-6)
    assert np.allclose(lla_to_xyz(xyz_to_lla(xyz)), xyz, rtol=1e-6)


def test_near_pole_round_trip() -> None:
    """Test round-trip near poles."""
    # Near north pole
    xyz = lla_to_xyz(_LLA_NEAR_NORTH_POLE)
    lla_back = xyz_to_lla(xyz)
    assert np.allclose(lla_back[:2], _LLA_NEAR_NORTH_POLE[:2], rtol=1e-4)
    assert abs(lla_back[2] - _LLA_NEAR_NORTH_POLE[2]) < 1  # within 1 meter


def test_xyz_to_lla_exact_pole_with_altitude() -> None:
    """Test xyz_to_lla at exact poles with x=y=0 and non-zero altitude.

    This tests the polar singularity fix. When x=y=0 exactly, the standard
    altitude formula (p/cos(lat) - n) fails because cos(lat)=0 at poles.
    The fix uses simplified polar geometry (alt = |z| - polar_radius) when
    the horizontal distance p < 1m.
    """
    # North pole at 1000m altitude with exact x=y=0
    xyz_north = np.array([0.0, 0.0, constants.PR_M + 1000.0])
    lla_north = xyz_to_lla(xyz_north)
    assert lla_north[0] == pytest.approx(90.0, abs=1e-6)
    assert lla_north[2] == pytest.approx(1000.0, abs=1e-6)

    # South pole at 500m altitude with exact x=y=0
    xyz_south = np.array([0.0, 0.0, -(constants.PR_M + 500.0)])
    lla_south = xyz_to_lla(xyz_south)
    assert lla_south[0] == pytest.approx(-90.0, abs=1e-6)
    assert lla_south[2] == pytest.approx(500.0, abs=1e-6)

    # Test vectorized with exact poles
    xyzs = np.array([xyz_north, xyz_south])
    llas = xyz_to_lla(xyzs)
    assert llas[0, 0] == pytest.approx(90.0, abs=1e-6)
    assert llas[0, 2] == pytest.approx(1000.0, abs=1e-6)
    assert llas[1, 0] == pytest.approx(-90.0, abs=1e-6)
    assert llas[1, 2] == pytest.approx(500.0, abs=1e-6)


# Vectorized tests


def test_lla_to_xyz_vectorized() -> None:
    """Test lla_to_xyz with batched (N, 3) input."""
    # Stack multiple known LLA points
    llas = np.array(
        [
            _LLA_MUNICH,
            _LLA_EQUATOR_PRIME,
            _LLA_EQUATOR_DATE_LINE,
        ],
    )
    expected = np.array(
        [
            _XYZ_MUNICH,
            _XYZ_EQUATOR_PRIME,
            _XYZ_EQUATOR_DATE_LINE,
        ],
    )

    xyzs = lla_to_xyz(llas)

    assert xyzs.shape == (3, 3), f"Expected shape (3, 3), got {xyzs.shape}"
    assert np.allclose(xyzs, expected, rtol=1e-6)


def test_xyz_to_lla_vectorized() -> None:
    """Test xyz_to_lla with batched (N, 3) input."""
    # Stack multiple known XYZ points
    xyzs = np.array(
        [
            _XYZ_MUNICH,
            _XYZ_EQUATOR_PRIME,
            _XYZ_EQUATOR_DATE_LINE,
        ],
    )
    expected = np.array(
        [
            _LLA_MUNICH,
            _LLA_EQUATOR_PRIME,
            _LLA_EQUATOR_DATE_LINE,
        ],
    )

    llas = xyz_to_lla(xyzs)

    assert llas.shape == (3, 3), f"Expected shape (3, 3), got {llas.shape}"
    assert np.allclose(llas, expected, rtol=1e-6)


def test_lla_xyz_vectorized_round_trip() -> None:
    """Test round-trip conversion with batched inputs."""
    llas = np.array(
        [
            [40.0, -75.0, 100.0],
            [0.0, 0.0, 0.0],
            [-33.9, 151.2, 50.0],  # Sydney
            [51.5, -0.1, 10.0],  # London
        ],
    )

    xyzs = lla_to_xyz(llas)
    llas_back = xyz_to_lla(xyzs)

    assert llas_back.shape == llas.shape
    assert np.allclose(llas_back, llas, rtol=1e-10)
