"""Transform XYZ (ECEF) vectors <=> AER (Azimuth, Elevation, Range) vectors.

XYZ is in meters and AER is in degrees, degrees, meters. Takes and returns numpy arrays.

# Get the AER vector between the two positions
get_aer(xyz_pos1,xyz_pos2) -> aer

# Get the new XYZ location shifted by this AER vector
translate_aer(xyz_pos, aer) -> xyz_pos

# Convert the XYZ vector to an AER vector
xyz_to_aer(xyz_pos, xyz_vec) -> aer

# Convert the AER vector to an XYZ vector
aer_to_xyz(xyz_pos, aer) -> xyz_vec
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .xyz_enu import enu_to_xyz, get_enu, translate_enu, xyz_to_enu

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def get_aer(
    xyz_pos1_m: NDArray[np.floating] | Sequence,
    xyz_pos2_m: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Return the AER vector from pos1 to pos2 in the plane of pos1.

    Args:
        xyz_pos1_m(np.ndarray|Sequence): XYZ reference location in meters. Shape (3,).
        xyz_pos2_m(np.ndarray|Sequence): XYZ target location(s) in meters. Shape (3,)
            or (N, 3).

    Returns:
        np.ndarray: AER vector(s) in the plane of pos1, pointed to pos2. In degrees
            (0,360), degrees (-90,90), meters. Shape (3,) or (N, 3) matching
            xyz_pos2_m.
    """
    enu = get_enu(xyz_pos1_m, xyz_pos2_m)
    return _enu_to_aer(enu)


def translate_aer(
    xyz_pos_m: NDArray[np.floating] | Sequence,
    aer_ddm: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Return the XYZ location of pos shifted by this AER.

    Args:
        xyz_pos_m(np.ndarray|Sequence): XYZ reference location in meters. Shape (3,).
        aer_ddm(np.ndarray|Sequence): AER vector(s) from pos in degrees, degrees,
            meters. Shape (3,) or (N, 3).

    Returns:
        np.ndarray: XYZ location(s) pointed at by aer in meters. Shape (3,) or (N, 3)
            matching aer_ddm.
    """
    enu = _aer_to_enu(np.asarray(aer_ddm))
    return translate_enu(xyz_pos_m, enu)


def xyz_to_aer(
    xyz_pos_m: NDArray[np.floating] | Sequence,
    xyz_vec_m: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Convert the XYZ vector to AER in the plane of the XYZ position.

    Args:
        xyz_pos_m(np.ndarray|Sequence): XYZ reference location in meters. Shape (3,).
        xyz_vec_m(np.ndarray|Sequence): XYZ vector(s) from that location in meters.
            Shape (3,) or (N, 3).

    Returns:
        np.ndarray: AER vector(s) in the plane of pos in degrees(0,360),
            degrees(-90,90), meters. Shape (3,) or (N, 3) matching xyz_vec_m.
    """
    enu = xyz_to_enu(xyz_pos_m, xyz_vec_m)
    return _enu_to_aer(enu)


def aer_to_xyz(
    xyz_pos_m: NDArray[np.floating] | Sequence,
    aer_ddm: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Convert the AER vector to an XYZ vector.

    Args:
        xyz_pos_m(np.ndarray|Sequence): XYZ reference location in meters. Shape (3,).
        aer_ddm(np.ndarray|Sequence): AER vector(s) from pos in degrees, degrees,
            meters. Shape (3,) or (N, 3).

    Returns:
        np.ndarray: XYZ vector(s) from pos in meters. Shape (3,) or (N, 3) matching
            aer_ddm.
    """
    enu = _aer_to_enu(np.asarray(aer_ddm))
    return enu_to_xyz(xyz_pos_m, enu)


def _enu_to_aer(enu_m: NDArray[np.floating]) -> NDArray[np.floating]:
    """Convert an ENU to an AER in this plane.

    Args:
        enu_m: ENU vector(s). Shape (3,) or (N, 3).

    Returns:
        AER vector(s). Shape (3,) or (N, 3) matching input.
    """
    e = enu_m[..., 0]
    n = enu_m[..., 1]
    u = enu_m[..., 2]
    r = np.linalg.norm(enu_m, axis=-1)
    r_surf = (e**2 + n**2) ** 0.5
    elev = np.degrees(np.arctan2(u, r_surf))
    az = np.degrees(np.arctan2(e, n) % (2 * np.pi))
    return np.stack((az, elev, r), axis=-1)


def _aer_to_enu(aer_ddm: NDArray[np.floating]) -> NDArray[np.floating]:
    """Convert an AER to ENU in this plane.

    Args:
        aer_ddm: AER vector(s). Shape (3,) or (N, 3).

    Returns:
        ENU vector(s). Shape (3,) or (N, 3) matching input.
    """
    azim_rad = np.radians(aer_ddm[..., 0])
    elev_rad = np.radians(aer_ddm[..., 1])
    r = aer_ddm[..., 2]
    e = np.sin(azim_rad) * np.cos(elev_rad) * r
    n = np.cos(azim_rad) * np.cos(elev_rad) * r
    u = np.sin(elev_rad) * r
    return np.stack((e, n, u), axis=-1)
