"""Transform XYZ (ECEF) locations <=> ECI (Earth-Centered Inertial) locations.

ECEF (Earth-Centered Earth-Fixed) rotates with the Earth, while ECI
(Earth-Centered Inertial) is fixed relative to the stars. The transformation
between them depends on time via the Greenwich Mean Sidereal Time (GMST).

These functions translate between ECEF and ECI coordinates in meters. They
take and return numpy arrays.

# Convert ECEF to ECI at a given time (Unix timestamp)
ecef_to_eci(xyz_ecef_m, unix_s)

# Convert ECI to ECEF at a given time (Unix timestamp)
eci_to_ecef(xyz_eci_m, unix_s)

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

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_eci(
    xyz_ecef_m: NDArray[np.floating] | Sequence,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> NDArray[np.floating]:
    """Convert ECEF coordinates to ECI 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 inertial frame.

    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:
        ECI 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
        >>> ecef = np.array([6378137.0, 0.0, 0.0])
        >>> eci = ecef_to_eci(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 ECI
    # ECI = R_z(gmst) @ ECEF
    # For positive gmst, ECEF X-axis has rotated eastward from ECI X-axis
    x_ecef = xyz_ecef_m[..., 0]
    y_ecef = xyz_ecef_m[..., 1]
    z_ecef = xyz_ecef_m[..., 2]

    x_eci = cos_gmst * x_ecef - sin_gmst * y_ecef
    y_eci = sin_gmst * x_ecef + cos_gmst * y_ecef
    z_eci = z_ecef  # Z-axis is the same for both frames

    return np.stack((x_eci, y_eci, z_eci), axis=-1)


def eci_to_ecef(
    xyz_eci_m: NDArray[np.floating] | Sequence,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> NDArray[np.floating]:
    """Convert ECI 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 the inertial frame to the
    Earth-fixed frame.

    Args:
        xyz_eci_m: ECI 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 ECI frame
        >>> eci = np.array([6378137.0, 0.0, 0.0])
        >>> ecef = eci_to_ecef(eci, 0.0)  # At Unix epoch
    """
    xyz_eci_m = np.asarray(xyz_eci_m)
    gmst = unix_to_gmst(unix_s)

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

    # Inverse rotation R_z(-gmst) applied to ECI gives ECEF
    # ECEF = R_z(-gmst) @ ECI = R_z(gmst)^T @ ECI
    x_eci = xyz_eci_m[..., 0]
    y_eci = xyz_eci_m[..., 1]
    z_eci = xyz_eci_m[..., 2]

    x_ecef = cos_gmst * x_eci + sin_gmst * y_eci
    y_ecef = -sin_gmst * x_eci + cos_gmst * y_eci
    z_ecef = z_eci  # Z-axis is the same for both frames

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