"""Frequency Difference of Arrival (FDOA) with gradient computation.

Provides FDOA computation along with gradients (Jacobians) with respect to both
emitter position and velocity, for use in optimization and fitting algorithms.

The full bistatic FDOA formula with moving collectors is:
    FDOA = f0/c * ((v_emit - v_coll1) . unit1 - (v_emit - v_coll2) . unit2)

The gradients are computed with respect to emitter position and velocity only
(collector positions/velocities are treated as fixed measurements).

For basic FDOA computation without gradients, see fdoa.py.
"""

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, _unit_vector_jacobian

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def fdoa_with_gradient(  # 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,
) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]:
    """Compute FDOA and gradients w.r.t. emitter position and velocity.

    Calculates the FDOA using the full bistatic formula and its gradients:
        FDOA = f0/c * ((v_emit - v_coll1) . unit1 - (v_emit - v_coll2) . unit2)

        d(FDOA)/d(emit_vel) = f0/c * (unit1 - unit2)

        d(FDOA)/d(emit_pos) = f0/c * (v_rel1.d(unit1)/d(pos) - v_rel2.d(unit2)/d(pos))

    where v_rel = v_emit - v_coll is the relative velocity.

    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:
        Tuple of (fdoa, grad_pos, grad_vel):
            - fdoa: FDOA in Hz. Shape () or (N,).
            - grad_pos: Gradient of FDOA w.r.t. emit_xyz in Hz/meter.
              Shape (3,) if emit is (3,), or (N, 3) if emit is (N, 3).
            - grad_vel: Gradient of FDOA w.r.t. emit_vel in Hz/(m/s).
              Shape (3,) if emit is (3,), or (N, 3) if emit is (N, 3).
    """
    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)

    single_point = emit_xyz.ndim == 1

    # Get range and unit vectors from emitter to each collector
    range1, unit1 = _range_and_unit(emit_xyz, collector1_xyz)
    range2, unit2 = _range_and_unit(emit_xyz, collector2_xyz)

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

    # Compute radial velocity components using relative velocities
    if single_point:
        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 value
    scale = freq_hz / C
    fdoa = scale * (v_radial1 - v_radial2)

    # Gradient w.r.t. emitter velocity: d(FDOA)/d(emit_vel) = f0/c * (unit1 - unit2)
    # Note: collector velocities don't affect this gradient
    grad_vel = scale * (unit1 - unit2)

    # Gradient w.r.t. emitter position requires unit vector Jacobians
    # d(FDOA)/d(pos) = f0/c * (v_rel1 . d(unit1)/d(pos) - v_rel2 . d(unit2)/d(pos))
    # Note: d(unit)/d(emit_pos) = -d(unit)/d(diff) since unit points emit->collector
    jac1 = _unit_vector_jacobian(unit1, range1)  # (3,3) or (N,3,3)
    jac2 = _unit_vector_jacobian(unit2, range2)

    if single_point:
        # v_rel @ jac gives d(v_rel.unit)/d(diff), negate for d/d(emit_pos)
        grad_pos = -scale * (v_rel1 @ jac1 - v_rel2 @ jac2)
    else:
        # Batch: need einsum for (N,3) @ (N,3,3) -> (N,3)
        grad_pos = -scale * (
            np.einsum("ni,nij->nj", v_rel1, jac1)
            - np.einsum("ni,nij->nj", v_rel2, jac2)
        )

    return np.asarray(fdoa), np.asarray(grad_pos), np.asarray(grad_vel)
