"""Transform XYZ (ECEF) vectors <=> ENU (East North Up) vectors.

Each of the following functions is unitless. EG: Pass in kilometers and the result is
in kilometers. Takes and returns numpy arrays.

# Get the ENU vector between the two positions
get_enu(xyz_pos1,xyz_pos2) -> enu

# Get the new XYZ location shifted by this ENU vector
translate_enu(xyz_pos, enu) -> xyz_pos

# Convert the XYZ vector to an ENU vector
xyz_to_enu(xyz_pos, xyz_vec) -> enu

# Convert the ENU vector to an XYZ vector
enu_to_xyz(xyz_pos, enu) -> xyz_vec
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .xyz_enu_rotation import enu_xyz_rotation, xyz_enu_rotation

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def get_enu(
    xyz_pos1: NDArray[np.floating] | Sequence,
    xyz_pos2: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Return the ENU vector from pos1 to pos2 in the plane of pos1.

    Args:
        xyz_pos1(np.ndarray|Sequence): XYZ reference location. Shape (3,).
        xyz_pos2(np.ndarray|Sequence): XYZ target location(s). Shape (3,) or (N, 3).

    Returns:
        np.ndarray: ENU vector(s) in the plane of pos1, pointed to pos2. Units are the
            same as pos1/pos2. Shape (3,) or (N, 3) matching xyz_pos2.
    """
    xyz_pos1 = np.asarray(xyz_pos1)
    xyz_pos2 = np.asarray(xyz_pos2)
    t_enu = xyz_enu_rotation(xyz_pos1)
    diff = xyz_pos2 - xyz_pos1
    if diff.ndim == 1:
        return t_enu @ diff
    return (t_enu @ diff.T).T


def translate_enu(
    xyz_pos: NDArray[np.floating] | Sequence,
    enu: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Return the XYZ location of pos shifted by this ENU.

    Args:
        xyz_pos(np.ndarray|Sequence): XYZ reference location. Shape (3,).
        enu(np.ndarray|Sequence): ENU vector(s) from pos. Shape (3,) or (N, 3).

    Returns:
        np.ndarray: XYZ location(s) pointed at by enu. Shape (3,) or (N, 3) matching
            enu.
    """
    xyz_pos = np.asarray(xyz_pos)
    enu = np.asarray(enu)
    t_xyz = enu_xyz_rotation(xyz_pos)
    if enu.ndim == 1:
        return t_xyz @ enu + xyz_pos
    return (t_xyz @ enu.T).T + xyz_pos


def xyz_to_enu(
    xyz_pos: NDArray[np.floating] | Sequence,
    xyz: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Convert the XYZ vector(s) to ENU in the plane of the XYZ position.

    Args:
        xyz_pos(np.ndarray|Sequence): XYZ reference location. Shape (3,).
        xyz(np.ndarray|Sequence): XYZ vector(s). Shape (3,) or (N, 3).

    Returns:
        np.ndarray: ENU vector(s). Shape (3,) or (N, 3) matching input xyz.

    Note:
        For covariance matrix transformation, use xyz_to_enu_cov() instead.
    """
    xyz_pos = np.asarray(xyz_pos)
    xyz = np.asarray(xyz)
    t_enu = xyz_enu_rotation(xyz_pos)
    if xyz.ndim == 1:
        return t_enu @ xyz
    return (t_enu @ xyz.T).T


def enu_to_xyz(
    xyz_pos: NDArray[np.floating] | Sequence,
    enu: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Convert the ENU vector(s) to an XYZ vector.

    Args:
        xyz_pos(np.ndarray|Sequence): XYZ reference location. Shape (3,).
        enu(np.ndarray|Sequence): ENU vector(s). Shape (3,) or (N, 3).

    Returns:
        np.ndarray: XYZ vector(s). Shape (3,) or (N, 3) matching input enu.

    Note:
        For covariance matrix transformation, use enu_to_xyz_cov() instead.
    """
    xyz_pos = np.asarray(xyz_pos)
    enu = np.asarray(enu)
    t_xyz = enu_xyz_rotation(xyz_pos)
    if enu.ndim == 1:
        return t_xyz @ enu
    return (t_xyz @ enu.T).T
