"""Range observable computation.

Range measures the Euclidean distance from a collector to an emitter:

    range = ||emit_xyz - collector_xyz||

This is the simplest geolocation observable. The iso-surface is a sphere
centered on the collector.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import NDArray


def get_range(
    collector_xyz: NDArray[np.floating] | Sequence,
    emit_xyz: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Compute range from collector to emitter.

    Args:
        collector_xyz: Collector position in ECEF meters. Shape (3,).
        emit_xyz: Emitter position(s) in ECEF meters. Shape (3,) or (N, 3).

    Returns:
        Range in meters. Shape () if emit is (3,), or (N,) if emit is (N, 3).
    """
    collector_xyz = np.asarray(collector_xyz)
    emit_xyz = np.asarray(emit_xyz)

    return np.asarray(np.linalg.norm(emit_xyz - collector_xyz, axis=-1))


def get_range_residual(
    collector_xyz: NDArray[np.floating] | Sequence,
    range_m: float,
) -> Callable[[NDArray[np.floating]], NDArray[np.floating]]:
    """Create a vectorized range residual function.

    Returns a function that computes the difference between observed range
    and the range predicted for candidate emitter positions. The zero-crossing
    of this residual defines the range iso-surface (sphere).

    Args:
        collector_xyz: Collector position in ECEF meters. Shape (3,).
        range_m: Measured range in meters.

    Returns:
        Vectorized residual function: takes points (N, 3) or (3,), returns
        residual values (N,) or (). Zero where predicted range equals measured.
    """
    center = np.asarray(collector_xyz, dtype=np.float64)

    def residual(points: NDArray[np.floating]) -> NDArray[np.floating]:
        pts = np.asarray(points)
        return np.asarray(np.linalg.norm(pts - center, axis=-1) - range_m)

    return residual
