"""Generic terrain sheet intersection with implicit surfaces.

Find where an implicit surface (defined by a residual function) intersects a
terrain sheet. The residual function can represent any observable iso-surface:
TDOA hyperboloids, range spheres, FDOA iso-Doppler surfaces, AOA cones, etc.

The approach:
1. Evaluate the residual function at each sheet grid point
2. Use marching squares to find zero-crossing contours
3. Convert contours to XYZ coordinates

Use residual factory functions from gri_utils.observables to create residual
functions for specific observables::

    from gri_utils.observables import get_tdoa_residual
    from gri_utils.terrain import intersect_sheet

    residual = get_tdoa_residual(collector1_xyz, collector2_xyz, tdoa_seconds)
    contours = intersect_sheet(sheet, residual)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .marching_squares import contours_to_xyz, find_contours

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import NDArray


def intersect_sheet(
    sheet_xyz: NDArray[np.floating] | Sequence,
    residual_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]],
    level: float = 0.0,
) -> list[NDArray[np.floating]]:
    """Find where an implicit surface intersects a terrain sheet.

    Evaluates a vectorized residual function at each grid point, finds
    zero-crossings using marching squares, and converts to XYZ coordinates.

    Args:
        sheet_xyz: Terrain sheet in ECEF meters. Shape (M, N, 3).
        residual_fn: Vectorized function taking XYZ points with shape (K, 3)
            and returning residual values with shape (K,). The intersection
            surface is where residual = level.
        level: Iso-surface level to find. Default 0.0.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    m, n, _ = sheet_xyz.shape

    flat_xyz = sheet_xyz.reshape(-1, 3)
    flat_residuals = residual_fn(flat_xyz)
    residuals = flat_residuals.reshape(m, n)

    contours = find_contours(residuals, level=level)

    return contours_to_xyz(contours, sheet_xyz)


def intersect_sheet_scalar(
    sheet_xyz: NDArray[np.floating] | Sequence,
    residual_fn: Callable[[NDArray[np.floating]], float],
    level: float = 0.0,
) -> list[NDArray[np.floating]]:
    """Find where an implicit surface intersects a terrain sheet (scalar version).

    Like intersect_sheet but for residual functions that process one point at
    a time. Prefer intersect_sheet with vectorized functions for performance.

    Args:
        sheet_xyz: Terrain sheet in ECEF meters. Shape (M, N, 3).
        residual_fn: Function taking an XYZ point (shape (3,)) and returning
            a scalar residual value. The surface is where residual = level.
        level: Iso-surface level to find. Default 0.0.

    Returns:
        List of intersection contours, each (K, 3) in ECEF meters.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    m, n, _ = sheet_xyz.shape

    residuals = np.zeros((m, n), dtype=np.float64)
    for i in range(m):
        for j in range(n):
            residuals[i, j] = residual_fn(sheet_xyz[i, j])

    contours = find_contours(residuals, level=level)

    return contours_to_xyz(contours, sheet_xyz)
