"""The rotation matrix for XYZ to ENU.

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 ENU vector:

t_enu.T @ xyz

ENU vector to XYZ vector:

t_enu @ enu

To simplify, we have two functions: enu_xyz_rotation and xyz_enu_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_enu_rotation(
    xyz_pos: NDArray[np.floating] | Sequence,
    *,
    geocentric: bool = False,
) -> NDArray[np.floating]:
    """Rotation matrix for XYZ to ENU.

    EG:
    t_enu = xyz_enu_rotation(xyz_pos)
    enu_vec = t_enu @ xyz_vec
    enu_mat = t_enu @ xyz_mat @ t_enu.T

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

    Note: In numpy the transpose is simply a view of this matrix, and is effectively
    O(1), so there is no need to stash the result somewhere for efficiency. Just call .T

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

    Returns:
        XYZ->ENU 3x3 numpy array
    """
    # Copy to modify
    u = np.array(xyz_pos)
    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])
    norm = np.linalg.norm(e)
    # When both x and y are zero, answer is unstable, prefer positive
    e = np.array([1, 0, 0]) if norm == 0 else e / norm
    # n = u x e
    n = np.cross(u, e)
    # enu = [
    #     [ e[0], n[0], u[0]],
    #     [ e[1], n[1], u[1]],
    #     [ e[2], n[2], u[2]],
    # ]
    return np.vstack((e, n, u))


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

    EG:
    t_xyz = enu_xyz_rotation(xyz_pos)
    xyz_vec = t_xyz @ enu_vec
    xyz_mat = t_xyz @ enu_mat @ t_xyz.T

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

    Note: In numpy the transpose is simply a view of this matrix, and is effectively
    O(1), so there is no need to stash the result somewhere for efficiency. Just call .T

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

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