"""Angle of Arrival (AOA) observable computation using ENU frame.

AOA computes direction cosines from a collector (sensor/antenna) to a signal source,
expressed in the ENU (East-North-Up) frame at the collector's position.

The output is the X and Y components (East and North) of a unit vector pointing
from collector to source. The Z component (Up) is implicit: +/-sqrt(1 - X**2 - Y**2).

For arbitrary collector orientations using quaternions, see aoa_quat.py.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils.conversion import cosines_to_unit_vector, enu_to_xyz, xyz_to_enu

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import NDArray


def get_aoa_enu(
    collector_xyz: NDArray[np.floating] | Sequence,
    source_xyz: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Get AOA direction cosines from collector to source in ENU frame.

    Computes the unit direction vector from collector to source, expressed
    in the ENU frame at the collector's position.

    Args:
        collector_xyz: Collector position in ECEF. Shape (3,).
        source_xyz: Source position(s) in ECEF. Shape (3,) or (n, 3).

    Returns:
        AOA direction cosines (East, North components of unit vector).
        Shape (2,) if source is single point, or (n, 2) for multiple sources.

    Note:
        Does not account for additional effects (e.g., atmospheric refraction).
    """
    collector_xyz = np.asarray(collector_xyz)
    source_xyz = np.asarray(source_xyz)

    # Get ENU vector from collector to source
    enu_vec = source_xyz - collector_xyz
    enu_vec = xyz_to_enu(collector_xyz, enu_vec)

    # Normalize to get unit direction
    if enu_vec.ndim == 1:
        norm = np.linalg.norm(enu_vec)
        enu_unit = enu_vec / norm
    else:
        norm = np.linalg.norm(enu_vec, axis=-1, keepdims=True)
        enu_unit = enu_vec / norm

    # Return only East and North components
    return enu_unit[..., :2]


def translate_aoa(
    collector_xyz: NDArray[np.floating] | Sequence,
    aoa: NDArray[np.floating] | Sequence,
    range_m: NDArray[np.floating] | Sequence[float] | np.floating | float,
    *,
    z_positive: bool = True,
) -> NDArray[np.floating]:
    """Translate from collector position using AOA and range to get source position.

    Args:
        collector_xyz: Collector position in ECEF meters. Shape (3,).
        aoa: AOA direction cosines (East, North components). Shape (2,) or (n, 2).
        range_m: Range from collector to source in meters. Scalar or shape (n,).
        z_positive: If True (default), source is above horizon (positive Up).
            If False, source is below horizon (negative Up).

    Returns:
        Source position(s) in ECEF meters. Shape (3,) or (n, 3).
    """
    collector_xyz = np.asarray(collector_xyz)
    aoa = np.asarray(aoa)
    range_m = np.asarray(range_m)

    # Get the ENU unit vector
    enu_unit = aoa_to_enu(aoa, z_positive=z_positive)

    # Scale by range
    if enu_unit.ndim == 1 or range_m.ndim == 0:
        enu_vec = enu_unit * range_m
    else:
        enu_vec = enu_unit * range_m[:, np.newaxis]

    # Convert ENU vector to ECEF and add to collector position
    xyz_vec = enu_to_xyz(collector_xyz, enu_vec)

    if xyz_vec.ndim == 1:
        return collector_xyz + xyz_vec
    return collector_xyz + xyz_vec


def xyz_to_aoa(
    collector_xyz: NDArray[np.floating] | Sequence,
    xyz_vec: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Convert XYZ vector to AOA direction cosines in ENU frame at collector.

    Args:
        collector_xyz: Collector position in ECEF. Shape (3,).
        xyz_vec: XYZ vector(s) in ECEF. Shape (3,) or (n, 3).

    Returns:
        AOA direction cosines (East, North components of normalized vector).
        Shape (2,) or (n, 2).
    """
    collector_xyz = np.asarray(collector_xyz)
    xyz_vec = np.asarray(xyz_vec)

    # Convert to ENU
    enu_vec = xyz_to_enu(collector_xyz, xyz_vec)

    # Normalize
    if enu_vec.ndim == 1:
        norm = np.linalg.norm(enu_vec)
        enu_unit = enu_vec / norm
    else:
        norm = np.linalg.norm(enu_vec, axis=-1, keepdims=True)
        enu_unit = enu_vec / norm

    # Return only East and North components
    return enu_unit[..., :2]


def aoa_to_xyz(
    collector_xyz: NDArray[np.floating] | Sequence,
    aoa: NDArray[np.floating] | Sequence,
    *,
    z_positive: bool = True,
) -> NDArray[np.floating]:
    """Convert AOA direction cosines to XYZ unit vector in ECEF.

    Args:
        collector_xyz: Collector position in ECEF. Shape (3,).
        aoa: AOA direction cosines (East, North components). Shape (2,) or (n, 2).
        z_positive: If True (default), source is above horizon (positive Up).
            If False, source is below horizon (negative Up).

    Returns:
        Unit direction vector(s) in ECEF. Shape (3,) or (n, 3).
    """
    collector_xyz = np.asarray(collector_xyz)
    aoa = np.asarray(aoa)

    # Get full ENU unit vector
    enu_unit = aoa_to_enu(aoa, z_positive=z_positive)

    # Convert to XYZ
    return enu_to_xyz(collector_xyz, enu_unit)


def aoa_to_enu(
    aoa: NDArray[np.floating] | Sequence,
    *,
    z_positive: bool = True,
) -> NDArray[np.floating]:
    """Convert AOA direction cosines to full ENU unit vector.

    Reconstructs the 3D unit direction vector from the East, North direction
    cosines by computing the Up component.

    Args:
        aoa: AOA direction cosines (East, North components). Shape (2,) or (n, 2).
        z_positive: If True (default), source is above horizon (positive Up).
            If False, source is below horizon (negative Up).

    Returns:
        ENU unit direction vector(s). Shape (3,) or (n, 3).
    """
    return cosines_to_unit_vector(aoa, z_positive=z_positive)


def get_aoa_residual(
    collector_xyz: NDArray[np.floating] | Sequence,
    aoa_cosines: NDArray[np.floating] | Sequence,
    *,
    z_positive: bool = True,
) -> Callable[[NDArray[np.floating]], NDArray[np.floating]]:
    """Create a vectorized AOA angular residual function.

    Returns a function that computes the angular deviation between the
    measured AOA bearing direction and the direction to candidate emitter
    positions. The zero-crossing defines the AOA bearing line.

    Args:
        collector_xyz: Collector position in ECEF meters. Shape (3,).
        aoa_cosines: Measured direction cosines (East, North) in the ENU
            frame at the collector. Shape (2,).
        z_positive: If True (default), source is above horizon.

    Returns:
        Vectorized residual function: takes points (N, 3) or (3,), returns
        angular deviation in radians (N,) or (). Zero along the bearing line.
    """
    col = np.asarray(collector_xyz, dtype=np.float64)
    aoa = np.asarray(aoa_cosines, dtype=np.float64)

    # Convert AOA direction cosines to ECEF bearing direction
    bearing_ecef = aoa_to_xyz(col, aoa, z_positive=z_positive)
    bearing_ecef = bearing_ecef / np.linalg.norm(bearing_ecef)

    def residual(points: NDArray[np.floating]) -> NDArray[np.floating]:
        pts = np.asarray(points)
        vec = pts - col
        dist = np.linalg.norm(vec, axis=-1, keepdims=True)
        unit_vec = vec / np.maximum(dist, 1e-12)

        cos_angle = np.sum(unit_vec * bearing_ecef, axis=-1)
        cos_angle = np.clip(cos_angle, -1.0, 1.0)
        return np.asarray(np.arccos(cos_angle))

    return residual
