"""Convert between 2D direction cosines and 3D unit vectors.

Direction cosines represent the X and Y components of a unit direction vector,
with the Z component implicit: Z = +/-sqrt(1 - X**2 - Y**2).

This representation is commonly used for:
- Angle of Arrival (AOA) measurements in geolocation
- Antenna pointing directions
- Any application where the third axis is assumed positive or negative

The sign ambiguity in Z is resolved by the z_positive parameter.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def cosines_to_unit_vector(
    cosines: NDArray[np.floating] | Sequence,
    *,
    z_positive: bool = True,
) -> NDArray[np.floating]:
    """Convert 2D direction cosines to a 3D unit vector.

    Reconstructs the third component of a unit vector from its first two
    components (direction cosines).

    Args:
        cosines: Direction cosines (X, Y components). Shape (2,) or (n, 2).
        z_positive: If True (default), Z component is positive.
            If False, Z component is negative.

    Returns:
        Unit vector(s) with shape (3,) or (n, 3).

    Example:
        >>> cosines_to_unit_vector([1.0, 0.0])  # Pointing along +X
        array([1., 0., 0.])
        >>> cosines_to_unit_vector([0.0, 0.0])  # Pointing along +Z
        array([0., 0., 1.])
        >>> cosines_to_unit_vector([0.0, 0.0], z_positive=False)  # Along -Z
        array([ 0.,  0., -1.])
    """
    cosines = np.asarray(cosines)

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

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

    # Reconstruct Z component from X, Y
    x = cosines[:, 0]
    y = cosines[:, 1]
    z_squared = 1.0 - x * x - y * y
    # Clamp to avoid numerical issues with sqrt of negative numbers
    z_squared = np.clip(z_squared, 0.0, 1.0)
    z = np.sqrt(z_squared)
    if not z_positive:
        z = -z

    # Build full 3D unit vector
    unit_vector = np.column_stack([x, y, z])  # (n, 3)

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

    return unit_vector


def unit_vector_to_cosines(
    unit_vector: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Extract 2D direction cosines from a 3D unit vector.

    Returns the X and Y components of the unit vector. The vector is normalized
    before extracting components, so non-unit vectors are also accepted.

    Args:
        unit_vector: 3D vector(s). Shape (3,) or (n, 3). Need not be normalized.

    Returns:
        Direction cosines (X, Y components) with shape (2,) or (n, 2).

    Example:
        >>> unit_vector_to_cosines([1.0, 0.0, 0.0])  # Along +X
        array([1., 0.])
        >>> unit_vector_to_cosines([0.0, 0.0, 1.0])  # Along +Z
        array([0., 0.])
        >>> unit_vector_to_cosines([3.0, 4.0, 0.0])  # Non-unit, normalized
        array([0.6, 0.8])
    """
    unit_vector = np.asarray(unit_vector)

    # Normalize (handles both single and batch)
    if unit_vector.ndim == 1:
        norm = np.linalg.norm(unit_vector)
        normalized = unit_vector / norm
    else:
        norm = np.linalg.norm(unit_vector, axis=-1, keepdims=True)
        normalized = unit_vector / norm

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