"""Transform XYZ (ECEF) locations <=> TEME (True Equator Mean Equinox) locations.

TEME is an Earth-Centered Inertial (ECI) reference frame used specifically by the
SGP4/SDP4 satellite propagation models. It differs from standard J2000 ECI in that
it uses the "true equator" (accounting for nutation) and "mean equinox of date"
(not accounting for nutation in the equinox direction).

This frame is primarily useful for interoperability with SGP4 propagator output
from libraries like python-sgp4. For most other applications, use the J2000 ECI
frame via xyz_eci.py instead.

The transformation between ECEF and TEME requires the Greenwich Mean Sidereal Time
(GMST), which accounts for Earth's rotation relative to the stars.

Note: For high-precision applications (sub-arcsecond), additional corrections for
polar motion and LOD (length of day) variations may be needed. This implementation
provides accuracy suitable for typical satellite tracking applications.

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 .gmst import unix_to_gmst

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def ecef_to_teme(
    xyz_ecef_m: NDArray[np.floating] | Sequence,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> NDArray[np.floating]:
    """Convert ECEF coordinates to TEME coordinates at a given time.

    Applies a rotation about the Z-axis by the Greenwich Mean Sidereal Time
    (GMST) angle to transform from the Earth-fixed frame to the TEME frame.

    TEME (True Equator Mean Equinox) is the coordinate frame used by SGP4/SDP4
    satellite propagators. Use this function when comparing positions from
    python-sgp4 with Earth-fixed coordinates.

    Args:
        xyz_ecef_m: ECEF coordinates in meters. Shape (3,) for single point
            or (N, 3) for N points.
        unix_s: Unix timestamp(s) in seconds. Scalar for single time, or
            array of shape (N,) for N times (one per point).

    Returns:
        TEME coordinates in meters as a numpy array. Shape (3,) or (N, 3)
            matching input.

    Examples:
        >>> import numpy as np
        >>> # Point on prime meridian at equator (at equatorial radius)
        >>> ecef = np.array([6378137.0, 0.0, 0.0])
        >>> teme = ecef_to_teme(ecef, 0.0)  # At Unix epoch
    """
    xyz_ecef_m = np.asarray(xyz_ecef_m)
    gmst = unix_to_gmst(unix_s)

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

    # Rotation matrix R_z(gmst) applied to ECEF gives TEME
    # TEME = R_z(gmst) @ ECEF
    # For positive gmst, ECEF X-axis has rotated eastward from TEME X-axis
    x_ecef = xyz_ecef_m[..., 0]
    y_ecef = xyz_ecef_m[..., 1]
    z_ecef = xyz_ecef_m[..., 2]

    x_teme = cos_gmst * x_ecef - sin_gmst * y_ecef
    y_teme = sin_gmst * x_ecef + cos_gmst * y_ecef
    z_teme = z_ecef  # Z-axis is the same for both frames

    return np.stack((x_teme, y_teme, z_teme), axis=-1)


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

    Applies the inverse rotation about the Z-axis by the Greenwich Mean
    Sidereal Time (GMST) angle to transform from TEME to ECEF.

    TEME (True Equator Mean Equinox) is the coordinate frame used by SGP4/SDP4
    satellite propagators. Use this function to convert satellite positions
    from python-sgp4 output to Earth-fixed coordinates.

    Args:
        xyz_teme_m: TEME coordinates in meters. Shape (3,) for single point
            or (N, 3) for N points.
        unix_s: Unix timestamp(s) in seconds. Scalar for single time, or
            array of shape (N,) for N times (one per point).

    Returns:
        ECEF coordinates in meters as a numpy array. Shape (3,) or (N, 3)
            matching input.

    Examples:
        >>> import numpy as np
        >>> # Point in TEME frame
        >>> teme = np.array([6378137.0, 0.0, 0.0])
        >>> ecef = teme_to_ecef(teme, 0.0)  # At Unix epoch
    """
    xyz_teme_m = np.asarray(xyz_teme_m)
    gmst = unix_to_gmst(unix_s)

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

    # Inverse rotation R_z(-gmst) applied to TEME gives ECEF
    # ECEF = R_z(-gmst) @ TEME = R_z(gmst)^T @ TEME
    x_teme = xyz_teme_m[..., 0]
    y_teme = xyz_teme_m[..., 1]
    z_teme = xyz_teme_m[..., 2]

    x_ecef = cos_gmst * x_teme + sin_gmst * y_teme
    y_ecef = -sin_gmst * x_teme + cos_gmst * y_teme
    z_ecef = z_teme  # Z-axis is the same for both frames

    return np.stack((x_ecef, y_ecef, z_ecef), axis=-1)
