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

TEME (True Equator Mean Equinox) is an Earth-Centered Inertial frame used by
the SGP4/SDP4 satellite propagation models. Like J2000 ECI, velocity
transformations require accounting for Earth's rotation rate.

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

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

This module is useful for converting SGP4 output (which includes velocity)
to Earth-fixed coordinates.

Functions:
    ecef_to_teme_vel(xyz_ecef_m, vel_ecef_m_s, unix_s) -> vel_teme_m_s
    teme_to_ecef_vel(xyz_teme_m, vel_teme_m_s, unix_s) -> vel_ecef_m_s

References:
    - Vallado, D. "Fundamentals of Astrodynamics and Applications", 4th ed.
    - AIAA 2006-6753 "Revisiting Spacetrack Report #3"
    - python-sgp4 documentation: https://pypi.org/project/sgp4/
"""

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_teme_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 TEME velocity at a given time.

    Transforms velocity from the rotating Earth-fixed frame to the TEME
    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_teme = R(gmst) @ v_ecef + omega x r_teme

    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:
        TEME velocity in m/s. Shape matches input broadcasting rules.

    Examples:
        >>> import numpy as np
        >>> # Convert ground station velocity to TEME for comparison with SGP4
        >>> pos_ecef = np.array([6378137.0, 0.0, 0.0])
        >>> vel_ecef = np.array([0.0, 0.0, 0.0])  # Stationary on ground
        >>> vel_teme = ecef_to_teme_vel(pos_ecef, vel_ecef, 0.0)
    """
    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 TEME
    x_ecef = xyz_ecef_m[..., 0]
    y_ecef = xyz_ecef_m[..., 1]

    x_teme = cos_gmst * x_ecef - sin_gmst * y_ecef
    y_teme = sin_gmst * x_ecef + cos_gmst * y_ecef

    # Rotate velocity to TEME (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_teme where omega = [0, 0, omega_earth]
    # omega x r = [0, 0, omega] x [x, y, z] = [-omega*y, omega*x, 0]
    vx_teme = vx_rot - OMEGA_EARTH * y_teme
    vy_teme = vy_rot + OMEGA_EARTH * x_teme
    vz_teme = vz_rot

    return np.stack((vx_teme, vy_teme, vz_teme), axis=-1)


def teme_to_ecef_vel(
    xyz_teme_m: NDArray[np.floating] | Sequence,
    vel_teme_m_s: NDArray[np.floating] | Sequence,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> NDArray[np.floating]:
    """Convert TEME velocity to ECEF velocity at a given time.

    Transforms velocity from the TEME 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.

    This is particularly useful for converting satellite velocities from SGP4
    output to Earth-fixed coordinates.

    The transformation is:
        v_ecef = R(-gmst) @ (v_teme - omega x r_teme)

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

    Args:
        xyz_teme_m: TEME position in meters. Shape (3,) or (N, 3).
        vel_teme_m_s: TEME 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
        >>> # Convert SGP4 satellite state to ECEF
        >>> pos_teme = np.array([7000e3, 0.0, 0.0])  # 7000 km from center
        >>> vel_teme = np.array([0.0, 7500.0, 0.0])  # ~7.5 km/s orbital velocity
        >>> vel_ecef = teme_to_ecef_vel(pos_teme, vel_teme, 0.0)
    """
    xyz_teme_m = np.asarray(xyz_teme_m)
    vel_teme_m_s = np.asarray(vel_teme_m_s)
    gmst = unix_to_gmst(unix_s)

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

    # Get TEME position components
    x_teme = xyz_teme_m[..., 0]
    y_teme = xyz_teme_m[..., 1]

    # Get TEME velocity components
    vx_teme = vel_teme_m_s[..., 0]
    vy_teme = vel_teme_m_s[..., 1]
    vz_teme = vel_teme_m_s[..., 2]

    # Subtract transport velocity: omega x r_teme where omega = [0, 0, omega_earth]
    # omega x r = [0, 0, omega] x [x, y, z] = [-omega*y, omega*x, 0]
    vx_inertial = vx_teme - (-OMEGA_EARTH * y_teme)  # vx + omega*y
    vy_inertial = vy_teme - (OMEGA_EARTH * x_teme)  # vy - omega*x
    vz_inertial = vz_teme

    # 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)
