"""Time Difference of Arrival (TDOA) observable computation.

TDOA measures the difference in signal arrival time at two collectors. It is
related to range difference by the speed of light:

    TDOA = range_diff / c = (||emit - collector1|| - ||emit - collector2||) / c

A positive value means the signal arrives at collector2 first (emitter is closer
to collector2).
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils.constants import C

from ._range_utils import _range_and_unit

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import NDArray


def get_tdoa(
    emit_xyz: NDArray[np.floating] | Sequence,
    collector1_xyz: NDArray[np.floating] | Sequence,
    collector2_xyz: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Compute Time Difference of Arrival from emitter to two collectors.

    Calculates the difference in signal arrival time at two collectors:
        TDOA = (||emit - collector1|| - ||emit - collector2||) / c

    Args:
        emit_xyz: Emitter position(s) in ECEF meters. Shape (3,) or (N, 3).
        collector1_xyz: First collector position in ECEF meters. Shape (3,).
        collector2_xyz: Second collector position in ECEF meters. Shape (3,).

    Returns:
        TDOA in seconds. Shape () if emit is (3,), or (N,) if emit is (N, 3).
        Positive means signal arrives at collector2 first.
    """
    emit_xyz = np.asarray(emit_xyz)
    collector1_xyz = np.asarray(collector1_xyz)
    collector2_xyz = np.asarray(collector2_xyz)

    r1, _ = _range_and_unit(collector1_xyz, emit_xyz)
    r2, _ = _range_and_unit(collector2_xyz, emit_xyz)

    return np.asarray((r1 - r2) / C)


def get_tdoa_residual(
    collector1_xyz: NDArray[np.floating] | Sequence,
    collector2_xyz: NDArray[np.floating] | Sequence,
    tdoa_seconds: float,
) -> Callable[[NDArray[np.floating]], NDArray[np.floating]]:
    """Create a vectorized TDOA residual function.

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

    Args:
        collector1_xyz: First collector position in ECEF meters. Shape (3,).
        collector2_xyz: Second collector position in ECEF meters. Shape (3,).
        tdoa_seconds: Measured TDOA in seconds. Positive means signal arrives
            at collector2 first.

    Returns:
        Vectorized residual function: takes points (N, 3) or (3,), returns
        residual values (N,) or (). Zero where predicted TDOA equals measured.
    """
    c1 = np.asarray(collector1_xyz, dtype=np.float64)
    c2 = np.asarray(collector2_xyz, dtype=np.float64)
    range_diff_target = tdoa_seconds * C

    def residual(points: NDArray[np.floating]) -> NDArray[np.floating]:
        pts = np.asarray(points)
        r1 = np.linalg.norm(pts - c1, axis=-1)
        r2 = np.linalg.norm(pts - c2, axis=-1)
        return np.asarray(r1 - r2 - range_diff_target)

    return residual
