"""Transform acceleration 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. Acceleration
transformations require accounting for Earth's rotation rate, including
Coriolis and centrifugal terms.

The acceleration transformation is:
    a_eci = R(t) @ a_ecef + 2*omega x v_eci + omega x (omega x r_eci)

where:
    - R(t) is the time-dependent rotation matrix
    - omega is Earth's rotation vector [0, 0, omega_earth]
    - 2*omega x v_eci is the Coriolis term
    - omega x (omega x r_eci) is the centrifugal term

Functions:
    ecef_to_eci_acc(xyz_ecef_m, vel_ecef_m_s, acc_ecef_m_s2, unix_s) -> acc_eci_m_s2
    eci_to_ecef_acc(xyz_eci_m, vel_eci_m_s, acc_eci_m_s2, unix_s) -> acc_ecef_m_s2

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
from .xyz_eci_vel import ecef_to_eci_vel

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


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

    Transforms acceleration from the rotating Earth-fixed frame to the inertial
    frame. This accounts for the rotation of the acceleration vector, the
    Coriolis acceleration (2*omega x v), and the centrifugal acceleration
    (omega x (omega x r)).

    The transformation is:
        a_eci = R(gmst) @ a_ecef + 2*omega x v_eci + omega x (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).
        acc_ecef_m_s2: ECEF acceleration in m/s**2. Shape (3,) or (N, 3).
        unix_s: Unix timestamp(s) in seconds. Scalar or array of shape (N,).

    Returns:
        ECI acceleration in m/s**2. Shape matches input broadcasting rules.

    Examples:
        >>> import numpy as np
        >>> # Object at rest in ECEF experiences centrifugal acceleration in ECI
        >>> pos_ecef = np.array([6378137.0, 0.0, 0.0])
        >>> vel_ecef = np.array([0.0, 0.0, 0.0])
        >>> acc_ecef = np.array([0.0, 0.0, 0.0])
        >>> acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, 0.0)
        >>> # acc_eci will have centrifugal component pointing outward
    """
    xyz_ecef_m = np.asarray(xyz_ecef_m)
    vel_ecef_m_s = np.asarray(vel_ecef_m_s)
    acc_ecef_m_s2 = np.asarray(acc_ecef_m_s2)
    gmst = unix_to_gmst(unix_s)

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

    # Rotate position to ECI
    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
    # z_eci = z_ecef (not needed, centrifugal is radial in X-Y)

    # Get velocity in ECI for Coriolis term
    vel_eci = ecef_to_eci_vel(xyz_ecef_m, vel_ecef_m_s, unix_s)
    vx_eci = vel_eci[..., 0]
    vy_eci = vel_eci[..., 1]

    # Rotate acceleration to ECI (without fictitious forces)
    ax_ecef = acc_ecef_m_s2[..., 0]
    ay_ecef = acc_ecef_m_s2[..., 1]
    az_ecef = acc_ecef_m_s2[..., 2]

    ax_rot = cos_gmst * ax_ecef - sin_gmst * ay_ecef
    ay_rot = sin_gmst * ax_ecef + cos_gmst * ay_ecef
    az_rot = az_ecef

    # Coriolis term: 2*omega x v_eci where omega = [0, 0, omega]
    # 2*omega x v = 2 * [0, 0, omega] x [vx, vy, vz] = 2 * [-omega*vy, omega*vx, 0]
    coriolis_x = -2 * OMEGA_EARTH * vy_eci
    coriolis_y = 2 * OMEGA_EARTH * vx_eci

    # Centrifugal term: omega x (omega x r_eci)
    # omega x r
    # = [0, 0, omega] x [x, y, z]
    # = [-omega*y, omega*x, 0]
    # omega x (omega x r)
    # = [0, 0, omega] x [-omega*y, omega*x, 0]
    # = [-omega**2*x, -omega**2*y, 0]
    omega_sq = OMEGA_EARTH * OMEGA_EARTH
    centrifugal_x = -omega_sq * x_eci
    centrifugal_y = -omega_sq * y_eci

    # Combine all terms
    ax_eci = ax_rot + coriolis_x + centrifugal_x
    ay_eci = ay_rot + coriolis_y + centrifugal_y
    az_eci = az_rot

    return np.stack((ax_eci, ay_eci, az_eci), axis=-1)


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

    Transforms acceleration from the inertial frame to the rotating Earth-fixed
    frame. This accounts for the rotation of the acceleration vector, and
    removes the Coriolis and centrifugal terms.

    The transformation is:
        a_ecef = R(-gmst) @ (a_eci - 2*omega x v_eci - omega x (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).
        acc_eci_m_s2: ECI acceleration in m/s**2. Shape (3,) or (N, 3).
        unix_s: Unix timestamp(s) in seconds. Scalar or array of shape (N,).

    Returns:
        ECEF acceleration in m/s**2. Shape matches input broadcasting rules.

    Examples:
        >>> import numpy as np
        >>> pos_eci = np.array([6378137.0, 0.0, 0.0])
        >>> vel_eci = np.array([0.0, 465.0, 0.0])
        >>> acc_eci = np.array([0.0, 0.0, 0.0])
        >>> acc_ecef = eci_to_ecef_acc(pos_eci, vel_eci, acc_eci, 0.0)
    """
    xyz_eci_m = np.asarray(xyz_eci_m)
    vel_eci_m_s = np.asarray(vel_eci_m_s)
    acc_eci_m_s2 = np.asarray(acc_eci_m_s2)
    gmst = unix_to_gmst(unix_s)

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

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

    vx_eci = vel_eci_m_s[..., 0]
    vy_eci = vel_eci_m_s[..., 1]

    ax_eci = acc_eci_m_s2[..., 0]
    ay_eci = acc_eci_m_s2[..., 1]
    az_eci = acc_eci_m_s2[..., 2]

    # Coriolis term: 2*omega x v_eci
    coriolis_x = -2 * OMEGA_EARTH * vy_eci
    coriolis_y = 2 * OMEGA_EARTH * vx_eci

    # Centrifugal term: omega x (omega x r_eci) = [-omega**2*x, -omega**2*y, 0]
    omega_sq = OMEGA_EARTH * OMEGA_EARTH
    centrifugal_x = -omega_sq * x_eci
    centrifugal_y = -omega_sq * y_eci

    # Remove fictitious forces
    ax_inertial = ax_eci - coriolis_x - centrifugal_x
    ay_inertial = ay_eci - coriolis_y - centrifugal_y
    az_inertial = az_eci

    # Rotate to ECEF: R(-gmst) = R(gmst)^T
    ax_ecef = cos_gmst * ax_inertial + sin_gmst * ay_inertial
    ay_ecef = -sin_gmst * ax_inertial + cos_gmst * ay_inertial
    az_ecef = az_inertial

    return np.stack((ax_ecef, ay_ecef, az_ecef), axis=-1)
