"""Frequency Difference of Arrival (FDOA) observable computation.

FDOA measures the difference in Doppler shift observed at two collectors due to
relative motion between emitter and collectors. The full bistatic formula is:

    FDOA = f0/c * ((v_emit - v_coll1) . unit1 - (v_emit - v_coll2) . unit2)

where:
    - f0 is the carrier frequency in Hz
    - c is the speed of light
    - v_emit is the emitter velocity vector
    - v_coll1, v_coll2 are the collector velocity vectors
    - unit1, unit2 are unit vectors from emitter to each collector

A positive value means higher Doppler at collector1 (emitter approaching
collector1 faster than collector2, or collector1 approaching emitter faster).

For stationary collectors (v_coll = 0), this simplifies to:
    FDOA = f0/c * (v_emit . unit1 - v_emit . unit2)
"""

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_fdoa(  # noqa: PLR0913
    emit_xyz: NDArray[np.floating] | Sequence,
    emit_vel: NDArray[np.floating] | Sequence,
    collector1_xyz: NDArray[np.floating] | Sequence,
    collector2_xyz: NDArray[np.floating] | Sequence,
    freq_hz: float,
    collector1_vel: NDArray[np.floating] | Sequence | None = None,
    collector2_vel: NDArray[np.floating] | Sequence | None = None,
) -> NDArray[np.floating]:
    """Compute Frequency Difference of Arrival from emitter to two collectors.

    Calculates the difference in Doppler shift at two collectors using the
    full bistatic formula that accounts for both emitter and collector motion:

        FDOA = f0/c * ((v_emit - v_coll1) . unit1 - (v_emit - v_coll2) . unit2)

    where unit vectors point from emitter to each collector.

    Args:
        emit_xyz: Emitter position(s) in ECEF meters. Shape (3,) or (N, 3).
        emit_vel: Emitter velocity in ECEF meters/second. 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,).
        freq_hz: Carrier frequency in Hz.
        collector1_vel: First collector velocity in ECEF m/s. Shape (3,).
            If None, collector is assumed stationary.
        collector2_vel: Second collector velocity in ECEF m/s. Shape (3,).
            If None, collector is assumed stationary.

    Returns:
        FDOA in Hz. Shape () if emit is (3,), or (N,) if emit is (N, 3).
        Positive means higher Doppler shift at collector1.
    """
    emit_xyz = np.asarray(emit_xyz)
    emit_vel = np.asarray(emit_vel)
    collector1_xyz = np.asarray(collector1_xyz)
    collector2_xyz = np.asarray(collector2_xyz)

    # Handle collector velocities (default to zero if not provided)
    coll1_vel = np.zeros(3) if collector1_vel is None else np.asarray(collector1_vel)
    coll2_vel = np.zeros(3) if collector2_vel is None else np.asarray(collector2_vel)

    # Get unit vectors from emitter to each collector
    _, unit1 = _range_and_unit(emit_xyz, collector1_xyz)
    _, unit2 = _range_and_unit(emit_xyz, collector2_xyz)

    # Compute relative velocities (emitter velocity relative to each collector)
    # v_rel = v_emit - v_coll
    v_rel1 = emit_vel - coll1_vel
    v_rel2 = emit_vel - coll2_vel

    # Compute radial velocity components (relative velocity projected onto unit vectors)
    if emit_vel.ndim == 1:
        v_radial1 = np.dot(v_rel1, unit1)
        v_radial2 = np.dot(v_rel2, unit2)
    else:
        v_radial1 = np.sum(v_rel1 * unit1, axis=-1)
        v_radial2 = np.sum(v_rel2 * unit2, axis=-1)

    # FDOA = f0/c * (v_radial1 - v_radial2)
    fdoa = (freq_hz / C) * (v_radial1 - v_radial2)

    return np.asarray(fdoa)


def get_fdoa_residual(  # noqa: PLR0913
    collector1_xyz: NDArray[np.floating] | Sequence,
    collector1_vel: NDArray[np.floating] | Sequence,
    collector2_xyz: NDArray[np.floating] | Sequence,
    collector2_vel: NDArray[np.floating] | Sequence,
    fdoa_hz: float,
    freq_hz: float,
    *,
    emitter_vel: NDArray[np.floating] | Sequence | None = None,
) -> Callable[[NDArray[np.floating]], NDArray[np.floating]]:
    """Create a vectorized FDOA residual function.

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

    Args:
        collector1_xyz: First collector position in ECEF meters. Shape (3,).
        collector1_vel: First collector velocity in ECEF m/s. Shape (3,).
        collector2_xyz: Second collector position in ECEF meters. Shape (3,).
        collector2_vel: Second collector velocity in ECEF m/s. Shape (3,).
        fdoa_hz: Measured FDOA in Hz.
        freq_hz: Carrier frequency in Hz.
        emitter_vel: Emitter velocity in ECEF m/s. Shape (3,).
            If None, emitter is assumed stationary.

    Returns:
        Vectorized residual function: takes points (N, 3) or (3,), returns
        residual values (N,) or (). Zero where predicted FDOA equals measured.
    """
    c1 = np.asarray(collector1_xyz, dtype=np.float64)
    c2 = np.asarray(collector2_xyz, dtype=np.float64)
    v1 = np.asarray(collector1_vel, dtype=np.float64)
    v2 = np.asarray(collector2_vel, dtype=np.float64)
    ev = (
        np.zeros(3, dtype=np.float64)
        if emitter_vel is None
        else np.asarray(
            emitter_vel,
            dtype=np.float64,
        )
    )

    # Relative velocities (collector - emitter)
    rel_vel1 = v1 - ev
    rel_vel2 = v2 - ev

    def residual(points: NDArray[np.floating]) -> NDArray[np.floating]:
        pts = np.asarray(points)
        vec1 = c1 - pts
        vec2 = c2 - pts

        r1 = np.linalg.norm(vec1, axis=-1, keepdims=True)
        r2 = np.linalg.norm(vec2, axis=-1, keepdims=True)

        u1 = vec1 / np.maximum(r1, 1e-12)
        u2 = vec2 / np.maximum(r2, 1e-12)

        doppler1 = np.sum(rel_vel1 * u1, axis=-1)
        doppler2 = np.sum(rel_vel2 * u2, axis=-1)

        computed = (freq_hz / C) * (doppler1 - doppler2)
        return np.asarray(computed - fdoa_hz)

    return residual
