"""Tests for adaptive terrain-observable intersection."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import pytest
from numpy.testing import assert_allclose

from gri_utils.conversion import lla_to_xyz
from gri_utils.terrain import lla_to_xyz_sheet
from gri_utils.terrain.adaptive import (
    _find_sign_change_cells,
    intersect_adaptive,
    intersect_source_with_range,
    intersect_source_with_tdoa,
)
from gri_utils.terrain.source import TerrainBounds, TerrainSource

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


class FlatTerrainSource(TerrainSource):
    """Simple flat terrain source for testing."""

    def __init__(self, altitude_m: float = 0.0, resolution_deg: float = 0.01) -> None:
        """Initialize with altitude and resolution."""
        self.altitude = altitude_m
        self.resolution = resolution_deg

    def get_covering_sheets(
        self,
        lats: NDArray[np.floating] | Sequence[float] | float,
        lons: NDArray[np.floating] | Sequence[float] | float,
    ) -> list[tuple[NDArray[np.floating], TerrainBounds]]:
        """Return a single flat terrain tile covering the requested region."""
        lats_arr = np.atleast_1d(np.asarray(lats, dtype=np.float64))
        lons_arr = np.atleast_1d(np.asarray(lons, dtype=np.float64))

        lat_min = np.floor(np.min(lats_arr))
        lat_max = np.ceil(np.max(lats_arr))
        lon_min = np.floor(np.min(lons_arr))
        lon_max = np.ceil(np.max(lons_arr))

        n_lat = int((lat_max - lat_min) / self.resolution) + 1
        n_lon = int((lon_max - lon_min) / self.resolution) + 1

        tile_lats = np.linspace(lat_min, lat_max, n_lat)
        tile_lons = np.linspace(lon_min, lon_max, n_lon)
        tile_alts = np.full((n_lat, n_lon), self.altitude)

        sheet = lla_to_xyz_sheet(tile_lats, tile_lons, tile_alts)
        bounds = TerrainBounds(lat_min, lat_max, lon_min, lon_max)

        return [(sheet, bounds)]

    def covers(self, lat: float, lon: float) -> bool:  # noqa: ARG002
        """Return True for all locations."""
        return True


# -----------------------------------------------------------------------------
# Helper function tests
# -----------------------------------------------------------------------------


def test_find_sign_change_cells_no_intersection() -> None:
    """Test sign change detection with no intersection."""
    # All positive values
    residuals = np.ones((5, 5))
    cells = _find_sign_change_cells(residuals)
    assert cells == []


def test_find_sign_change_cells_with_intersection() -> None:
    """Test sign change detection with intersection."""
    residuals = np.array(
        [
            [1.0, 1.0, 1.0],
            [1.0, -1.0, 1.0],
            [1.0, 1.0, 1.0],
        ],
    )
    cells = _find_sign_change_cells(residuals)

    # Should find 4 cells around the negative center
    assert len(cells) == 4
    assert (0, 0) in cells
    assert (0, 1) in cells
    assert (1, 0) in cells
    assert (1, 1) in cells


def test_find_sign_change_cells_skips_nan() -> None:
    """Test that cells with NaN are skipped."""
    residuals = np.array(
        [
            [1.0, np.nan, 1.0],
            [1.0, -1.0, 1.0],
            [1.0, 1.0, 1.0],
        ],
    )
    cells = _find_sign_change_cells(residuals)

    # Cell (0,0) has NaN corner, should be skipped
    # Cell (0,1) has NaN corner, should be skipped
    assert (0, 0) not in cells
    assert (0, 1) not in cells
    # Cell (1,0) and (1,1) should still be found
    assert (1, 0) in cells
    assert (1, 1) in cells


# -----------------------------------------------------------------------------
# Adaptive intersection tests
# -----------------------------------------------------------------------------


def test_intersect_adaptive_sphere() -> None:
    """Test adaptive intersection with a sphere (range)."""
    source = FlatTerrainSource(altitude_m=0.0, resolution_deg=0.01)

    # Sphere centered above the terrain at 39, -105
    center_lla = np.array([39.0, -105.0, 100000.0])  # 100 km altitude
    center_xyz = lla_to_xyz(center_lla)

    # Get a terrain point to compute distance to terrain surface
    terrain_lla = np.array([39.0, -105.0, 0.0])
    terrain_xyz = lla_to_xyz(terrain_lla)

    # Radius equals distance to terrain surface - so sphere touches terrain at center
    # and intersects in a ring around it
    radius = np.linalg.norm(center_xyz - terrain_xyz) + 50000.0  # Extend beyond center

    def residual_fn(points: NDArray[np.floating]) -> NDArray[np.floating]:
        return np.linalg.norm(points - center_xyz, axis=-1) - radius

    contours = intersect_adaptive(
        source,
        residual_fn,
        lats=[38.5, 39.5],
        lons=[-105.5, -104.5],
        subsample_factor=5,
    )

    # Should find at least one contour
    assert len(contours) >= 1

    # All contour points should be approximately at the sphere surface
    for contour in contours:
        distances = np.linalg.norm(contour - center_xyz, axis=-1)
        assert_allclose(distances, radius, rtol=0.01)


def test_intersect_adaptive_no_intersection() -> None:
    """Test adaptive intersection when there's no intersection."""
    source = FlatTerrainSource(altitude_m=0.0, resolution_deg=0.01)

    # Sphere way above terrain that doesn't intersect
    center_xyz = np.array([0.0, 0.0, 1e8])  # Very far from Earth
    radius = 1000.0  # Small radius

    def residual_fn(points: NDArray[np.floating]) -> NDArray[np.floating]:
        return np.linalg.norm(points - center_xyz, axis=-1) - radius

    contours = intersect_adaptive(
        source,
        residual_fn,
        lats=[38.5, 39.5],
        lons=[-105.5, -104.5],
        subsample_factor=5,
    )

    assert contours == []


def test_intersect_adaptive_subsample_1() -> None:
    """Test adaptive intersection with subsample_factor=1 (full resolution)."""
    source = FlatTerrainSource(altitude_m=0.0, resolution_deg=0.1)  # Coarser for speed

    center_lla = np.array([39.0, -105.0, 100000.0])
    center_xyz = lla_to_xyz(center_lla)

    # Compute radius to intersect terrain
    terrain_lla = np.array([39.0, -105.0, 0.0])
    terrain_xyz = lla_to_xyz(terrain_lla)
    radius = np.linalg.norm(center_xyz - terrain_xyz) + 20000.0

    def residual_fn(points: NDArray[np.floating]) -> NDArray[np.floating]:
        return np.linalg.norm(points - center_xyz, axis=-1) - radius

    contours = intersect_adaptive(
        source,
        residual_fn,
        lats=[38.8, 39.2],
        lons=[-105.2, -104.8],
        subsample_factor=1,  # Full resolution
    )

    assert len(contours) >= 1


def test_intersect_adaptive_invalid_subsample_raises() -> None:
    """Test that subsample_factor < 1 raises ValueError."""
    source = FlatTerrainSource()

    def residual_fn(points: NDArray[np.floating]) -> NDArray[np.floating]:
        return np.zeros(points.shape[0])

    with pytest.raises(ValueError, match="subsample_factor"):
        intersect_adaptive(source, residual_fn, [39.0], [-105.0], subsample_factor=0)


def test_intersect_adaptive_empty_source_raises() -> None:
    """Test that empty source result raises ValueError."""

    class EmptySource(TerrainSource):
        def get_covering_sheets(
            self,
            lats: NDArray[np.floating] | Sequence[float] | float,  # noqa: ARG002
            lons: NDArray[np.floating] | Sequence[float] | float,  # noqa: ARG002
        ) -> list[tuple[NDArray[np.floating], TerrainBounds]]:
            return []

        def covers(self, lat: float, lon: float) -> bool:  # noqa: ARG002
            return False

    source = EmptySource()

    def residual_fn(points: NDArray[np.floating]) -> NDArray[np.floating]:
        return np.zeros(points.shape[0])

    with pytest.raises(ValueError, match="No terrain data"):
        intersect_adaptive(source, residual_fn, [39.0], [-105.0])


# -----------------------------------------------------------------------------
# Convenience wrapper tests
# -----------------------------------------------------------------------------


def test_intersect_source_with_range() -> None:
    """Test range sphere intersection wrapper."""
    source = FlatTerrainSource(altitude_m=0.0, resolution_deg=0.05)

    center_lla = np.array([39.0, -105.0, 50000.0])
    center_xyz = lla_to_xyz(center_lla)

    # Compute radius to intersect terrain
    terrain_lla = np.array([39.0, -105.0, 0.0])
    terrain_xyz = lla_to_xyz(terrain_lla)
    radius = np.linalg.norm(center_xyz - terrain_xyz) + 20000.0

    contours = intersect_source_with_range(
        source,
        lats=[38.8, 39.2],
        lons=[-105.2, -104.8],
        center_xyz=center_xyz,
        range_m=float(radius),
        subsample_factor=3,
    )

    assert len(contours) >= 1


def test_intersect_source_with_tdoa() -> None:
    """Test TDOA hyperboloid intersection wrapper."""
    source = FlatTerrainSource(altitude_m=0.0, resolution_deg=0.05)

    # Two collectors above the terrain
    c1_lla = np.array([39.0, -105.5, 500000.0])
    c2_lla = np.array([39.0, -104.5, 500000.0])
    c1_xyz = lla_to_xyz(c1_lla)
    c2_xyz = lla_to_xyz(c2_lla)

    # Zero TDOA - equidistant points
    contours = intersect_source_with_tdoa(
        source,
        lats=[38.8, 39.2],
        lons=[-105.2, -104.8],
        collector1_xyz=c1_xyz,
        collector2_xyz=c2_xyz,
        tdoa_seconds=0.0,  # Zero TDOA = perpendicular bisector plane
        subsample_factor=3,
    )

    assert len(contours) >= 1

    # For zero TDOA, contour should be equidistant from both collectors
    for contour in contours:
        d1 = np.linalg.norm(contour - c1_xyz, axis=-1)
        d2 = np.linalg.norm(contour - c2_xyz, axis=-1)
        assert_allclose(d1, d2, rtol=0.01)
