"""Angle of Arrival (AOA) with gradient computation.

Provides AOA computation along with Jacobian matrices for use in optimization
and fitting algorithms.

For basic AOA computation without gradients, see aoa_quat.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, _unit_vector_jacobian
from .rotation_utils import rotation_to_quat_array

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def aoa_quat_with_gradient(
    collector_xyz: NDArray[np.floating] | Sequence,
    source_xyz: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Compute Angle of Arrival direction cosines and gradient w.r.t. source position.

    Computes the unit direction vector from collector to source, rotated into the
    collector's local frame, along with the Jacobian matrix of the AOA with respect
    to the source position.

    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:
        Tuple of (aoa, gradient):
            - aoa: AOA direction cosines (X, Y components). Shape (2,) or (n, 2).
            - gradient: Jacobian of AOA w.r.t. source position.
              Shape (2, 3) or (n, 2, 3). gradient[i, j] = dAOA_i/dsource_j.

    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))
    )

    # Determine if single point or batch
    single_point = collector_xyz.ndim == 1

    # Ensure 2D for consistent processing
    if single_point:
        collector_xyz = collector_xyz[np.newaxis, :]
        source_xyz = source_xyz[np.newaxis, :]
        # Wrap single rotation in array-like for consistent indexing
        if rotation.single:
            rotation = Rotation.concatenate([rotation])

    # Compute range and unit vector using shared utility
    range_val, direction_unit = _range_and_unit(collector_xyz, source_xyz)

    # Get rotation matrix
    rot_matrix = rotation.as_matrix()  # (n, 3, 3)

    # Rotate direction vector to the local frame
    aoa_3d = rotation.apply(direction_unit)  # (n, 3)
    aoa = aoa_3d[:, :2]  # (n, 2)

    # Compute gradient (Jacobian of AOA w.r.t. source position)
    # The Jacobian of a normalized vector d_n = d/||d|| is:
    #   dd_n/dd = (I - d_n @ d_n.T) / ||d||
    # Then we rotate this Jacobian by the rotation matrix R:
    #   dAOA/dsource = R @ dd_n/dd

    # Get the Jacobian of normalized direction in ECEF frame using shared utility
    j_ecef = _unit_vector_jacobian(direction_unit, range_val)  # (n, 3, 3)

    # Rotate Jacobian to local frame: J_local = R @ J_ecef
    # rot_matrix is (n, 3, 3), j_ecef is (n, 3, 3)
    j_local = np.einsum("nij,njk->nik", rot_matrix, j_ecef)  # (n, 3, 3)

    # Extract only X and Y rows for the gradient
    gradient = j_local[:, :2, :]  # (n, 2, 3)

    # Return to original shape if single point
    if single_point:
        aoa = aoa[0]  # (2,)
        gradient = gradient[0]  # (2, 3)

    return aoa, gradient


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

    Follows numpy broadcasting semantics for batch dimensions.

    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:
        Tuple of (aoa, gradient):
            - aoa: AOA direction cosines. Shape (..., 2).
            - gradient: Jacobian of AOA w.r.t. source position. Shape (..., 2, 3).
        Single inputs (shape (M,)) return shapes (2,) and (2, 3).
    """
    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
    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
    aoa_flat, gradient_flat = aoa_quat_with_gradient(
        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 aoa_flat.reshape(2), gradient_flat.reshape(2, 3)
    return aoa_flat.reshape(*batch_shape, 2), gradient_flat.reshape(*batch_shape, 2, 3)
