"""Inverse AOA computation: convert AOA direction cosines back to XYZ.

Reconstructs the 3D direction vector from AOA measurements, rotating from the
collector's local frame back to ECEF.

For forward AOA computation (XYZ to AOA), see aoa_quat.py.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

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

from gri_utils.conversion import cosines_to_unit_vector

from .rotation_utils import rotation_to_quat_array

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


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

    Reconstructs the 3D unit direction vector from the X, Y direction cosines
    and rotates it from the collector's local frame back to ECEF.

    Args:
        aoa: AOA direction cosines (X, Y components). Shape (2,) or (n, 2).
        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).
        z_positive: If True (default), assume source is in front of collector
            (positive Z in local frame). If False, source is behind (negative Z).

    Returns:
        Unit direction vector(s) in ECEF pointing from collector toward source.
        Shape (3,) if inputs are single points, or (n, 3) for multiple points.

    Note:
        To get the source position, multiply by range and add to collector position:
        ``source_xyz = collector_xyz + aoa_quat_to_xyz(aoa, rot) * range_m``
    """
    aoa = np.asarray(aoa)

    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 = aoa.ndim == 1

    # Ensure 2D for consistent processing
    if single_point:
        aoa = aoa[np.newaxis, :]

    # Reconstruct full 3D unit vector in local frame from direction cosines
    local_unit = cosines_to_unit_vector(aoa, z_positive=z_positive)

    # Rotate from local frame back to ECEF using inverse rotation
    xyz_unit = rotation.inv().apply(local_unit)

    # Return to original shape if single point
    if single_point:
        xyz_unit = xyz_unit[0]

    return xyz_unit


def aoa_quat_to_xyz_broadcast(
    aoa: NDArray[np.floating] | Sequence,
    collector_rotation: NDArray[np.floating] | Sequence | Rotation,
    *,
    z_positive: bool = True,
) -> NDArray[np.floating]:
    """Broadcast-capable wrapper for aoa_quat_to_xyz.

    Follows numpy broadcasting semantics for batch dimensions.

    Args:
        aoa: AOA direction cosines (X, Y components). Shape (..., 2).
        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).
        z_positive: If True (default), assume source is in front of collector
            (positive Z in local frame). If False, source is behind (negative Z).

    Returns:
        Unit direction vector(s) in ECEF pointing from collector toward source.
        Shape (..., 3) matching the broadcast batch shape of inputs.
        Single inputs (shape (M,)) return shape (3,).
    """
    aoa = np.asarray(aoa)
    collector_quat = rotation_to_quat_array(collector_rotation)

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

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

    # Extract batch shapes
    batch_a = aoa.shape[:-1]
    batch_q = collector_quat.shape[:-1]

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

    # Broadcast to common shape
    aoa = np.broadcast_to(aoa, (*batch_shape, 2))
    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 = aoa_quat_to_xyz(
        aoa.reshape(n, 2),
        collector_quat.reshape(n, 4),
        z_positive=z_positive,
    )

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