"""Shared utilities for range-based observable computations.

Private module providing helper functions used by range_diff, tdoa, and aoa modules.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from numpy.typing import NDArray


def _range_and_unit(
    from_xyz: NDArray[np.floating],
    to_xyz: NDArray[np.floating],
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Compute range and unit vector between two points.

    Args:
        from_xyz: Origin point(s) in ECEF. Shape (3,) or (N, 3).
        to_xyz: Target point(s) in ECEF. Shape (3,) or (N, 3).

    Returns:
        Tuple of (range_val, unit):
            - range_val: Euclidean distance. Shape () or (N,).
            - unit: Unit direction vector from origin to target. Shape (3,) or (N, 3).
    """
    diff = to_xyz - from_xyz

    if diff.ndim == 1:
        range_val: NDArray[np.floating] = np.asarray(np.linalg.norm(diff))
        unit = diff / range_val
    else:
        range_val = np.linalg.norm(diff, axis=-1)
        unit = diff / range_val[..., np.newaxis]

    return range_val, unit


def _unit_vector_jacobian(
    unit: NDArray[np.floating],
    range_val: NDArray[np.floating],
) -> NDArray[np.floating]:
    """Compute Jacobian of unit vector with respect to the un-normalized vector.

    The Jacobian of normalizing a vector v to v/||v|| is:
        d(v/||v||)/dv = (I - u ⊗ u) / ||v||

    where u is the unit vector and ⊗ denotes outer product.

    Args:
        unit: Unit vector(s). Shape (3,) or (N, 3).
        range_val: Norm of original vector(s). Shape () or (N,).

    Returns:
        Jacobian matrix. Shape (3, 3) for single input, or (N, 3, 3) for batch.
        jacobian[i, j] = d(unit_i)/d(vec_j)
    """
    single_point = unit.ndim == 1

    if single_point:
        unit = unit[np.newaxis, :]
        range_val = np.atleast_1d(range_val)

    n = unit.shape[0]

    # Extract components
    ux = unit[:, 0]
    uy = unit[:, 1]
    uz = unit[:, 2]

    # Build Jacobian: (I - u ⊗ u) / r
    # J[i,j] = (delta_ij - u_i * u_j) / r
    jacobian = np.zeros((n, 3, 3))
    jacobian[:, 0, 0] = (1 - ux * ux) / range_val
    jacobian[:, 0, 1] = (-ux * uy) / range_val
    jacobian[:, 0, 2] = (-ux * uz) / range_val
    jacobian[:, 1, 0] = (-uy * ux) / range_val
    jacobian[:, 1, 1] = (1 - uy * uy) / range_val
    jacobian[:, 1, 2] = (-uy * uz) / range_val
    jacobian[:, 2, 0] = (-uz * ux) / range_val
    jacobian[:, 2, 1] = (-uz * uy) / range_val
    jacobian[:, 2, 2] = (1 - uz * uz) / range_val

    if single_point:
        return jacobian[0]

    return jacobian
