"""Transform velocity vectors between ECEF and ECI reference frames.

ECEF (Earth-Centered Earth-Fixed) rotates with the Earth, while ECI
(Earth-Centered Inertial) is fixed relative to the stars. Velocity
transformations require accounting for Earth's rotation rate.

The velocity transformation includes the transport term from Earth's rotation:
    v_eci = R(t) @ v_ecef + omega x r_eci

where omega is the Earth's rotation vector [0, 0, omega_earth] and R(t) is the
time-dependent rotation matrix.

Functions:
    ecef_to_eci_vel(xyz_ecef_m, vel_ecef_m_s, unix_s) -> vel_eci_m_s
    eci_to_ecef_vel(xyz_eci_m, vel_eci_m_s, unix_s) -> vel_ecef_m_s

References:
    - Vallado, D. "Fundamentals of Astrodynamics and Applications", 4th ed.
    - IERS Conventions (2010), Chapter 5
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils.constants import OMEGA_EARTH

from .gmst import unix_to_gmst

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def ecef_to_eci_vel(
    xyz_ecef_m: NDArray[np.floating] | Sequence,
    vel_ecef_m_s: NDArray[np.floating] | Sequence,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> NDArray[np.floating]:
    """Convert ECEF velocity to ECI velocity at a given time.

    Transforms velocity from the rotating Earth-fixed frame to the inertial
    frame. This accounts for both the rotation of the velocity vector and
    the transport velocity due to Earth's rotation.

    The transformation is:
        v_eci = R(gmst) @ v_ecef + omega x r_eci

    where omega = [0, 0, omega_earth] is Earth's rotation vector.

    Args:
        xyz_ecef_m: ECEF position in meters. Shape (3,) or (N, 3).
        vel_ecef_m_s: ECEF velocity in m/s. Shape (3,) or (N, 3).
        unix_s: Unix timestamp(s) in seconds. Scalar or array of shape (N,).

    Returns:
        ECI velocity in m/s. Shape matches input broadcasting rules.

    Examples:
        >>> import numpy as np
        >>> # Stationary point on equator in ECEF has velocity in ECI
        >>> pos_ecef = np.array([6378137.0, 0.0, 0.0])
        >>> vel_ecef = np.array([0.0, 0.0, 0.0])
        >>> vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, 0.0)
        >>> # vel_eci will be ~465 m/s in Y direction (Earth surface velocity)
    """
    xyz_ecef_m = np.asarray(xyz_ecef_m)
    vel_ecef_m_s = np.asarray(vel_ecef_m_s)
    gmst = unix_to_gmst(unix_s)

    cos_gmst = np.cos(gmst)
    sin_gmst = np.sin(gmst)

    # Rotate position to ECI (only x,y needed for cross product with omega along z-axis)
    x_ecef = xyz_ecef_m[..., 0]
    y_ecef = xyz_ecef_m[..., 1]

    x_eci = cos_gmst * x_ecef - sin_gmst * y_ecef
    y_eci = sin_gmst * x_ecef + cos_gmst * y_ecef

    # Rotate velocity to ECI (without transport term)
    vx_ecef = vel_ecef_m_s[..., 0]
    vy_ecef = vel_ecef_m_s[..., 1]
    vz_ecef = vel_ecef_m_s[..., 2]

    vx_rot = cos_gmst * vx_ecef - sin_gmst * vy_ecef
    vy_rot = sin_gmst * vx_ecef + cos_gmst * vy_ecef
    vz_rot = vz_ecef

    # Add transport velocity: omega x r_eci where omega = [0, 0, omega_earth]
    # omega x r = [0, 0, omega] x [x, y, z] = [-omega*y, omega*x, 0]
    vx_eci = vx_rot - OMEGA_EARTH * y_eci
    vy_eci = vy_rot + OMEGA_EARTH * x_eci
    vz_eci = vz_rot

    return np.stack((vx_eci, vy_eci, vz_eci), axis=-1)


def eci_to_ecef_vel(
    xyz_eci_m: NDArray[np.floating] | Sequence,
    vel_eci_m_s: NDArray[np.floating] | Sequence,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> NDArray[np.floating]:
    """Convert ECI velocity to ECEF velocity at a given time.

    Transforms velocity from the inertial frame to the rotating Earth-fixed
    frame. This accounts for both the rotation of the velocity vector and
    the transport velocity due to Earth's rotation.

    The transformation is:
        v_ecef = R(-gmst) @ (v_eci - omega x r_eci)

    where omega = [0, 0, omega_earth] is Earth's rotation vector.

    Args:
        xyz_eci_m: ECI position in meters. Shape (3,) or (N, 3).
        vel_eci_m_s: ECI velocity in m/s. Shape (3,) or (N, 3).
        unix_s: Unix timestamp(s) in seconds. Scalar or array of shape (N,).

    Returns:
        ECEF velocity in m/s. Shape matches input broadcasting rules.

    Examples:
        >>> import numpy as np
        >>> # Point with velocity in ECI
        >>> pos_eci = np.array([6378137.0, 0.0, 0.0])
        >>> vel_eci = np.array([0.0, 465.0, 0.0])
        >>> vel_ecef = eci_to_ecef_vel(pos_eci, vel_eci, 0.0)
    """
    xyz_eci_m = np.asarray(xyz_eci_m)
    vel_eci_m_s = np.asarray(vel_eci_m_s)
    gmst = unix_to_gmst(unix_s)

    cos_gmst = np.cos(gmst)
    sin_gmst = np.sin(gmst)

    # Get ECI position components
    x_eci = xyz_eci_m[..., 0]
    y_eci = xyz_eci_m[..., 1]

    # Get ECI velocity components
    vx_eci = vel_eci_m_s[..., 0]
    vy_eci = vel_eci_m_s[..., 1]
    vz_eci = vel_eci_m_s[..., 2]

    # Subtract transport velocity: omega x r_eci where omega = [0, 0, omega_earth]
    # omega x r = [0, 0, omega] x [x, y, z] = [-omega*y, omega*x, 0]
    vx_inertial = vx_eci - (-OMEGA_EARTH * y_eci)  # vx + omega*y
    vy_inertial = vy_eci - (OMEGA_EARTH * x_eci)  # vy - omega*x
    vz_inertial = vz_eci

    # Rotate to ECEF: R(-gmst) = R(gmst)^T
    vx_ecef = cos_gmst * vx_inertial + sin_gmst * vy_inertial
    vy_ecef = -sin_gmst * vx_inertial + cos_gmst * vy_inertial
    vz_ecef = vz_inertial

    return np.stack((vx_ecef, vy_ecef, vz_ecef), axis=-1)
