"""The rotation matrix for XYZ to NED.

The rotation matrix (and its transpose) can be used to transform positions or ellipses
(coordinate spaces) from one to the other.

The transpose in numpy is a view of the matrix, with no time complexity to calculate,
so we don't include it, but it can be obtained with ".T" from the rotation matrix.

To transform points:

XYZ vector to an NED vector:

t_ned.T @ xyz

NED vector to XYZ vector:

t_ned @ ned

To simplify, we have two functions: ned_xyz_rotation and xyz_ned_rotation.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils import constants

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def xyz_ned_rotation(
    xyz_pos: NDArray[np.floating] | Sequence,
    *,
    geocentric: bool = False,
) -> NDArray[np.floating]:
    """Rotation matrix for XYZ to NED.

    EG:
    t_ned = xyz_ned_rotation(location)
    ned = t_ned @ xyz_delta

    To reverse, just the the transpose ".T" of this matrix. t_xyz.T == t_ned

    Args:
        xyz_pos(ndarray|Sequence): Position for transformation
        geocentric(bool): Normally using local vertical (geodedic), but can use radius
            (geocentric)

    Returns:
        XYZ->NED 3x3 numpy array
    """
    # Copy to modify
    u = np.array(xyz_pos, copy=True)
    if not geocentric:
        r3fc = constants.ER_M / constants.PR_M
        u[2] = u[2] * r3fc**2
    # normalize, don't use /= with numpy because of internal representation mappings
    u = u / np.linalg.norm(u)
    # e = z x u normalized
    e = np.array([-u[1], u[0], 0])
    e = e / np.linalg.norm(e)
    # n = u x e
    n = np.cross(u, e)
    # ned = [
    #     [ n[0], e[0], -u[0]],
    #     [ n[1], e[1], -u[1]],
    #     [ n[2], e[2], -u[2]],
    # ]
    return np.vstack((n, e, -u))


def ned_xyz_rotation(
    xyz_pos: NDArray[np.floating] | Sequence,
    *,
    geocentric: bool = False,
) -> NDArray[np.floating]:
    """Rotation matrix for NED to XYZ.

    EG:
    t_xyz = ned_xyz_rotation(location)
    xyz_delta = t_xyz @ ned

    To reverse, just the the transpose ".T" of this matrix. t_ned.T == t_xyz

    Args:
        xyz_pos(ndarray|Sequence): Position for transformation
        geocentric(bool): Normally using local vertical (geodedic), but can use radius
            (geocentric)

    Returns:
        NED->XYZ 3x3 numpy array
    """
    return xyz_ned_rotation(xyz_pos, geocentric=geocentric).T
