"""Time Difference of Arrival (TDOA) with gradient computation.

Provides TDOA computation along with the gradient (Jacobian) for use in
optimization and fitting algorithms.

For basic TDOA computation without gradients, see tdoa.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

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def tdoa_with_gradient(
    emit_xyz: NDArray[np.floating] | Sequence,
    collector1_xyz: NDArray[np.floating] | Sequence,
    collector2_xyz: NDArray[np.floating] | Sequence,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Compute TDOA and gradient w.r.t. emitter position.

    Calculates the TDOA and its gradient:
        TDOA = (||emit - collector1|| - ||emit - collector2||) / c
        gradient = (unit_vec_1 - unit_vec_2) / c

    where unit_vec_i points from collector_i to the emitter.

    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:
        Tuple of (tdoa, gradient):
            - tdoa: TDOA in seconds. Shape () or (N,).
            - gradient: Gradient of TDOA w.r.t. emit_xyz in seconds/meter.
              Shape (3,) if emit is (3,), or (N, 3) if emit is (N, 3).
              gradient[i] = d(TDOA)/d(emit_xyz[i])
    """
    emit_xyz = np.asarray(emit_xyz)
    collector1_xyz = np.asarray(collector1_xyz)
    collector2_xyz = np.asarray(collector2_xyz)

    r1, unit1 = _range_and_unit(collector1_xyz, emit_xyz)
    r2, unit2 = _range_and_unit(collector2_xyz, emit_xyz)

    tdoa = np.asarray((r1 - r2) / C)
    gradient = np.asarray((unit1 - unit2) / C)

    return tdoa, gradient
