"""Adaptive terrain-observable intersection algorithm.

Finds where implicit observable surfaces (TDOA hyperboloids, range spheres, etc.)
intersect terrain sheets using adaptive subsampling for efficiency.

The algorithm:
1. Loads full-resolution terrain covering the region of interest
2. Subsamples for a coarse initial pass to find cells with sign changes
3. Runs marching squares at full resolution only where intersections exist
4. Stitches contour segments into polylines

This approach is much faster than evaluating the full grid when intersections
are sparse (the typical case for geolocation observables).

Example::

    from gri_utils import terrain
    from gri_utils.observables import get_tdoa_residual

    residual = get_tdoa_residual(c1, c2, tdoa_seconds=0.001)
    contours = terrain.intersect_adaptive(source, residual, lats, lons)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .marching_squares import contours_to_xyz, find_contours
from .stitch import stitch_sheets

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import NDArray

    from .source import TerrainSource


def intersect_adaptive(
    source: TerrainSource,
    residual_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]],
    lats: NDArray[np.floating] | Sequence[float] | float,
    lons: NDArray[np.floating] | Sequence[float] | float,
    subsample_factor: int = 10,
) -> list[NDArray[np.floating]]:
    """Find observable-terrain intersection with adaptive resolution.

    Loads full-resolution terrain, subsamples for coarse initial pass to
    find intersection regions, then runs marching squares at full resolution
    only where intersections exist.

    Args:
        source: Terrain data source implementing TerrainSource ABC.
        residual_fn: Vectorized function taking XYZ points with shape (N, 3)
            and returning residual values with shape (N,). The intersection
            surface is where residual = 0.
        lats: Latitude(s) defining region of interest in decimal degrees.
        lons: Longitude(s) defining region of interest in decimal degrees.
        subsample_factor: Factor for coarse pass. 10 means evaluate every
            10th point initially. Must be >= 1. Default 10.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.

    Raises:
        ValueError: If no terrain data is available for the region.
    """
    if subsample_factor < 1:
        msg = "subsample_factor must be >= 1"
        raise ValueError(msg)

    # Get terrain sheets covering the region
    sheets_with_bounds = source.get_covering_sheets(lats, lons)
    if not sheets_with_bounds:
        msg = "No terrain data available for the specified region"
        raise ValueError(msg)

    # Stitch into single sheet if needed
    full_sheet, _bounds = stitch_sheets(sheets_with_bounds)
    m, n, _ = full_sheet.shape

    # If subsample_factor is 1 or sheet is small, just run full resolution
    is_small = m <= subsample_factor * 2 and n <= subsample_factor * 2
    if subsample_factor == 1 or is_small:
        return _intersect_full_resolution(full_sheet, residual_fn)

    # Subsample for coarse pass
    coarse_rows = np.arange(0, m, subsample_factor)
    coarse_cols = np.arange(0, n, subsample_factor)

    # Ensure we include the last row/col
    if coarse_rows[-1] != m - 1:
        coarse_rows = np.append(coarse_rows, m - 1)
    if coarse_cols[-1] != n - 1:
        coarse_cols = np.append(coarse_cols, n - 1)

    # Extract coarse grid
    coarse_sheet = full_sheet[np.ix_(coarse_rows, coarse_cols)]
    cm, cn, _ = coarse_sheet.shape

    # Evaluate residual on coarse grid
    coarse_flat = coarse_sheet.reshape(-1, 3)
    coarse_residuals = residual_fn(coarse_flat).reshape(cm, cn)

    # Find cells with sign changes (potential intersections)
    sign_change_cells = _find_sign_change_cells(coarse_residuals)

    if not sign_change_cells:
        # No intersections found
        return []

    # Process each sign-change cell at full resolution
    all_contours: list[NDArray[np.floating]] = []

    for ci, cj in sign_change_cells:
        # Map coarse cell indices back to full resolution indices
        row_start = coarse_rows[ci]
        row_end = coarse_rows[ci + 1] + 1 if ci + 1 < len(coarse_rows) else m
        col_start = coarse_cols[cj]
        col_end = coarse_cols[cj + 1] + 1 if cj + 1 < len(coarse_cols) else n

        # Extract full-resolution sub-sheet for this cell
        sub_sheet = full_sheet[row_start:row_end, col_start:col_end]

        # Run marching squares on this sub-sheet
        cell_contours = _intersect_full_resolution(sub_sheet, residual_fn)

        all_contours.extend(cell_contours)

    # TODO: Stitch contours that cross cell boundaries
    # For now, return all contour segments - they may have small gaps at boundaries

    return all_contours


def _intersect_full_resolution(
    sheet: NDArray[np.floating],
    residual_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]],
) -> list[NDArray[np.floating]]:
    """Run intersection on a sheet at full resolution.

    Args:
        sheet: XYZ terrain sheet, shape (M, N, 3).
        residual_fn: Vectorized residual function.

    Returns:
        List of contours in XYZ coordinates.
    """
    m, n, _ = sheet.shape

    # Evaluate residual at all grid points
    flat_xyz = sheet.reshape(-1, 3)

    # Handle NaN values (gaps in stitched sheets)
    valid_mask = ~np.isnan(flat_xyz[:, 0])
    residuals = np.full(flat_xyz.shape[0], np.nan, dtype=np.float64)
    if np.any(valid_mask):
        residuals[valid_mask] = residual_fn(flat_xyz[valid_mask])

    residuals = residuals.reshape(m, n)

    # Find contours using marching squares
    contours = find_contours(residuals, level=0.0)

    # Convert to XYZ coordinates
    return contours_to_xyz(contours, sheet)


def _find_sign_change_cells(
    residuals: NDArray[np.floating],
) -> list[tuple[int, int]]:
    """Find cells where residual changes sign (potential intersections).

    Args:
        residuals: 2D array of residual values, shape (M, N).

    Returns:
        List of (row, col) indices for cells with sign changes.
        Each cell is defined by its top-left corner index.
    """
    m, n = residuals.shape
    sign_change_cells = []

    for i in range(m - 1):
        for j in range(n - 1):
            # Get corner values
            corners = [
                residuals[i, j],
                residuals[i, j + 1],
                residuals[i + 1, j],
                residuals[i + 1, j + 1],
            ]

            # Skip if any corner is NaN
            if any(np.isnan(c) for c in corners):
                continue

            # Check for sign change
            signs = [c > 0 for c in corners]
            if not all(signs) and any(signs):
                sign_change_cells.append((i, j))

    return sign_change_cells


# -----------------------------------------------------------------------------
# Convenience wrappers for common observables
# -----------------------------------------------------------------------------


def intersect_source_with_tdoa(  # noqa: PLR0913
    source: TerrainSource,
    lats: NDArray[np.floating] | Sequence[float] | float,
    lons: NDArray[np.floating] | Sequence[float] | float,
    collector1_xyz: NDArray[np.floating] | Sequence[float],
    collector2_xyz: NDArray[np.floating] | Sequence[float],
    tdoa_seconds: float,
    subsample_factor: int = 10,
) -> list[NDArray[np.floating]]:
    """Find TDOA-terrain intersection with adaptive resolution.

    Args:
        source: Terrain data source implementing TerrainSource ABC.
        lats: Latitude(s) defining region of interest in decimal degrees.
        lons: Longitude(s) defining region of interest in decimal degrees.
        collector1_xyz: First collector position in ECEF meters, shape (3,).
        collector2_xyz: Second collector position in ECEF meters, shape (3,).
        tdoa_seconds: TDOA value in seconds. Positive means signal arrives
            at collector2 first.
        subsample_factor: Factor for coarse pass. Default 10.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.
    """
    from gri_utils.observables import get_tdoa_residual  # noqa: PLC0415

    residual_fn = get_tdoa_residual(collector1_xyz, collector2_xyz, tdoa_seconds)
    return intersect_adaptive(source, residual_fn, lats, lons, subsample_factor)


def intersect_source_with_range(  # noqa: PLR0913
    source: TerrainSource,
    lats: NDArray[np.floating] | Sequence[float] | float,
    lons: NDArray[np.floating] | Sequence[float] | float,
    center_xyz: NDArray[np.floating] | Sequence[float],
    range_m: float,
    subsample_factor: int = 10,
) -> list[NDArray[np.floating]]:
    """Find range sphere-terrain intersection with adaptive resolution.

    Args:
        source: Terrain data source implementing TerrainSource ABC.
        lats: Latitude(s) defining region of interest in decimal degrees.
        lons: Longitude(s) defining region of interest in decimal degrees.
        center_xyz: Sphere center in ECEF meters, shape (3,).
        range_m: Range value (sphere radius) in meters.
        subsample_factor: Factor for coarse pass. Default 10.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.
    """
    from gri_utils.observables import get_range_residual  # noqa: PLC0415

    residual_fn = get_range_residual(center_xyz, range_m)
    return intersect_adaptive(source, residual_fn, lats, lons, subsample_factor)


def intersect_source_with_fdoa(  # noqa: PLR0913
    source: TerrainSource,
    lats: NDArray[np.floating] | Sequence[float] | float,
    lons: NDArray[np.floating] | Sequence[float] | float,
    collector1_xyz: NDArray[np.floating] | Sequence[float],
    collector1_vel: NDArray[np.floating] | Sequence[float],
    collector2_xyz: NDArray[np.floating] | Sequence[float],
    collector2_vel: NDArray[np.floating] | Sequence[float],
    fdoa_hz: float,
    freq_hz: float,
    subsample_factor: int = 10,
    *,
    emitter_vel: NDArray[np.floating] | Sequence[float] | None = None,
) -> list[NDArray[np.floating]]:
    """Find FDOA iso-Doppler surface-terrain intersection with adaptive resolution.

    Args:
        source: Terrain data source implementing TerrainSource ABC.
        lats: Latitude(s) defining region of interest in decimal degrees.
        lons: Longitude(s) defining region of interest in decimal degrees.
        collector1_xyz: First collector position in ECEF meters, shape (3,).
        collector1_vel: First collector velocity in m/s, shape (3,).
        collector2_xyz: Second collector position in ECEF meters, shape (3,).
        collector2_vel: Second collector velocity in m/s, shape (3,).
        fdoa_hz: FDOA value in Hz.
        freq_hz: Carrier frequency in Hz.
        subsample_factor: Factor for coarse pass. Default 10.
        emitter_vel: Emitter velocity in m/s, shape (3,). Default is zero.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.
    """
    from gri_utils.observables import get_fdoa_residual  # noqa: PLC0415

    residual_fn = get_fdoa_residual(
        collector1_xyz,
        collector1_vel,
        collector2_xyz,
        collector2_vel,
        fdoa_hz,
        freq_hz,
        emitter_vel=emitter_vel,
    )
    return intersect_adaptive(source, residual_fn, lats, lons, subsample_factor)


def intersect_source_with_aoa(  # noqa: PLR0913
    source: TerrainSource,
    lats: NDArray[np.floating] | Sequence[float] | float,
    lons: NDArray[np.floating] | Sequence[float] | float,
    collector_xyz: NDArray[np.floating] | Sequence[float],
    aoa_cosines: NDArray[np.floating] | Sequence[float],
    subsample_factor: int = 10,
) -> list[NDArray[np.floating]]:
    """Find AOA bearing line-terrain intersection with adaptive resolution.

    Args:
        source: Terrain data source implementing TerrainSource ABC.
        lats: Latitude(s) defining region of interest in decimal degrees.
        lons: Longitude(s) defining region of interest in decimal degrees.
        collector_xyz: Collector position in ECEF meters, shape (3,).
        aoa_cosines: Direction cosines (cos_east, cos_north) defining the
            bearing direction in ENU frame, shape (2,).
        subsample_factor: Factor for coarse pass. Default 10.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.
    """
    from gri_utils.observables import get_aoa_residual  # noqa: PLC0415

    residual_fn = get_aoa_residual(collector_xyz, aoa_cosines)
    return intersect_adaptive(source, residual_fn, lats, lons, subsample_factor)
