"""Angle of Arrival (AOA) observable computation with quaternion orientation.

AOA computes direction cosines from a collector (sensor/antenna) to a signal source,
expressed in the collector's local reference frame defined by a quaternion.

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

For ENU-based AOA (without explicit quaternion), see aoa_enu.py.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from scipy.spatial.transform import Rotation

from ._range_utils import _range_and_unit
from .rotation_utils import rotation_to_quat_array

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def get_aoa_quat(
    collector_xyz: NDArray[np.floating] | Sequence,
    source_xyz: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
) -> NDArray[np.floating]:
    """Get AOA direction cosines from collector to source using quaternion orientation.

    Computes the unit direction vector from collector to source, rotated into the
    collector's local frame using the provided rotation.

    Args:
        collector_xyz: Collector position(s) in ECEF. Shape (3,) or (n, 3).
        source_xyz: Source position(s) in ECEF. Shape (3,) or (n, 3).
        collector_rotation: Rotation defining collector orientation relative to ECEF.
            Either a scipy Rotation object, or quaternion array with shape (4,) or
            (n, 4) using scipy convention [x, y, z, w] (scalar-last).

    Returns:
        AOA direction cosines (X, Y components of unit vector in local frame).
        Shape (2,) if inputs are single points, or (n, 2) for multiple points.

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

    rotation: Rotation = (
        collector_rotation
        if isinstance(collector_rotation, Rotation)
        else Rotation.from_quat(np.asarray(collector_rotation))
    )

    # Compute unit direction vector from collector to source
    _, direction_unit = _range_and_unit(collector_xyz, source_xyz)

    # Rotate direction vector to the local frame using rotation
    aoa_3d = rotation.apply(direction_unit)

    # Return only X and Y components
    return aoa_3d[..., :2]


def get_aoa_quat_broadcast(
    collector_xyz: NDArray[np.floating] | Sequence,
    source_xyz: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
) -> NDArray[np.floating]:
    """Broadcast-capable wrapper for get_aoa_quat.

    Follows numpy broadcasting semantics for batch dimensions. Inputs with shape
    (..., M) are broadcast together on the (...) dimensions, then flattened for
    computation and reshaped back.

    Args:
        collector_xyz: Collector position(s) in ECEF. Shape (..., 3).
        source_xyz: Source position(s) in ECEF. Shape (..., 3).
        collector_rotation: Rotation defining collector orientation relative to ECEF.
            Either a scipy Rotation object, or quaternion array with shape (..., 4)
            using scipy convention [x, y, z, w] (scalar-last).

    Returns:
        AOA direction cosines (X, Y components of unit vector in local frame).
        Shape (..., 2) matching the broadcast batch shape of inputs.
        Single inputs (shape (M,)) return shape (2,).

    Example:
        >>> # Single collector, multiple sources
        >>> collector = np.array([6378000.0, 0.0, 0.0])  # (3,)
        >>> sources = np.random.randn(10, 5, 3) * 100 + collector  # (10, 5, 3)
        >>> quat = np.array([0.0, 0.0, 0.0, 1.0])  # (4,)
        >>> aoa = get_aoa_quat_broadcast(collector, sources, quat)  # (10, 5, 2)
    """
    collector_xyz = np.asarray(collector_xyz)
    source_xyz = np.asarray(source_xyz)
    collector_quat = rotation_to_quat_array(collector_rotation)

    # Track if all inputs are 1D (scalar case)
    all_1d = (
        collector_xyz.ndim == 1 and source_xyz.ndim == 1 and collector_quat.ndim == 1
    )

    # Ensure at least 2D for batch dimension handling
    collector_xyz = np.atleast_2d(collector_xyz)
    source_xyz = np.atleast_2d(source_xyz)
    collector_quat = np.atleast_2d(collector_quat)

    # Extract batch shapes (everything except last dim)
    batch_c = collector_xyz.shape[:-1]
    batch_s = source_xyz.shape[:-1]
    batch_q = collector_quat.shape[:-1]

    # Compute broadcast batch shape
    batch_shape = np.broadcast_shapes(batch_c, batch_s, batch_q)

    # Broadcast to common shape
    collector_xyz = np.broadcast_to(collector_xyz, (*batch_shape, 3))
    source_xyz = np.broadcast_to(source_xyz, (*batch_shape, 3))
    collector_quat = np.broadcast_to(collector_quat, (*batch_shape, 4))

    # Flatten, compute, reshape
    n = int(np.prod(batch_shape)) if batch_shape else 1
    result_flat = get_aoa_quat(
        collector_xyz.reshape(n, 3),
        source_xyz.reshape(n, 3),
        collector_quat.reshape(n, 4),
    )

    # Return scalar shape if all inputs were 1D
    if all_1d:
        return result_flat.reshape(2)
    return result_flat.reshape(*batch_shape, 2)
