"""Convert orbital elements to state vector at a given time.

This module implements Keplerian orbit propagation: given orbital elements
and a target time, compute the position, velocity, and acceleration vectors
in the ECI frame.

The key algorithm is solving Kepler's equation to find the eccentric anomaly,
then converting through true anomaly to position and velocity in the perifocal
frame, and finally rotating to ECI coordinates.

References:
    - Vallado, D. "Fundamentals of Astrodynamics and Applications", 4th ed.
    - Bate, Mueller, White "Fundamentals of Astrodynamics"
"""

from __future__ import annotations

import math
from typing import TYPE_CHECKING

import numpy as np

from gri_utils.constants import GM_EARTH

from .orbit_state import OrbitState

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray

    from .orbital_elements import OrbitalElements


def elements_to_state(
    elements: OrbitalElements,
    unix_s: float | NDArray[np.floating] | Sequence[float],
) -> OrbitState | list[OrbitState]:
    """Propagate orbital elements to state vector(s) at given time(s).

    Performs Keplerian (two-body) orbit propagation by:
    1. Propagating mean anomaly from epoch to target time
    2. Solving Kepler's equation for eccentric anomaly
    3. Computing true anomaly
    4. Computing position and velocity in perifocal frame
    5. Rotating to ECI frame
    6. Computing gravitational acceleration

    Args:
        elements: OrbitalElements defining the orbit.
        unix_s: Target time(s) as Unix timestamp(s) in seconds. Can be a scalar
            for a single time, or a sequence/array for multiple times.

    Returns:
        If unix_s is scalar: Single OrbitState at that time.
        If unix_s is array-like: List of OrbitState objects, one per time.

    Example:
        >>> from gri_utils.orbit import OrbitalElements, elements_to_state
        >>> elements = OrbitalElements(
        ...     semi_major_axis_m=7_000_000,
        ...     eccentricity=0.001,
        ...     inclination_rad=0.0,
        ...     raan_rad=0.0,
        ...     arg_periapsis_rad=0.0,
        ...     mean_anomaly_rad=0.0,
        ...     epoch_unix_s=0.0,
        ... )
        >>> state = elements_to_state(elements, 0.0)
        >>> print(f"Position: {state.position_eci_m}")
    """
    # Check if input is scalar or array
    is_scalar = np.isscalar(unix_s)
    times = np.atleast_1d(np.asarray(unix_s, dtype=np.float64))

    # Extract elements
    a = elements.semi_major_axis_m
    e = elements.eccentricity
    i = elements.inclination_rad
    raan = elements.raan_rad
    omega = elements.arg_periapsis_rad
    m0 = elements.mean_anomaly_rad
    t0 = elements.epoch_unix_s

    # Mean motion: n = sqrt(mu / a^3)
    n = math.sqrt(GM_EARTH / (a * a * a))

    # Precompute rotation matrix (same for all times)
    rot_matrix = _rotation_matrix_313(raan, i, omega)

    # Semi-latus rectum: p = a * (1 - e^2)
    p = a * (1.0 - e * e)

    results = []
    for t in times:
        # Propagate mean anomaly
        dt = t - t0
        mean_anomaly = (m0 + n * dt) % (2.0 * math.pi)

        # Solve Kepler's equation for eccentric anomaly
        eccentric_anomaly = _solve_kepler(mean_anomaly, e)

        # Convert to true anomaly
        true_anomaly = _eccentric_to_true_anomaly(eccentric_anomaly, e)

        # Position magnitude: r = p / (1 + e*cos(nu))
        cos_nu = math.cos(true_anomaly)
        sin_nu = math.sin(true_anomaly)
        r_mag = p / (1.0 + e * cos_nu)

        # Position in perifocal frame (x toward periapsis, z normal to orbit)
        r_perifocal = np.array([r_mag * cos_nu, r_mag * sin_nu, 0.0])

        # Velocity in perifocal frame
        # v = sqrt(mu/p) * [-sin(nu), e + cos(nu), 0]
        sqrt_mu_p = math.sqrt(GM_EARTH / p)
        v_perifocal = np.array(
            [
                -sqrt_mu_p * sin_nu,
                sqrt_mu_p * (e + cos_nu),
                0.0,
            ],
        )

        # Rotate to ECI frame
        r_eci = rot_matrix @ r_perifocal
        v_eci = rot_matrix @ v_perifocal

        # Gravitational acceleration: a = -mu/r^3 * r
        r_mag_cubed = r_mag * r_mag * r_mag
        a_eci = -GM_EARTH / r_mag_cubed * r_eci

        results.append(
            OrbitState(
                position_eci_m=r_eci,
                velocity_eci_ms=v_eci,
                acceleration_eci_ms2=a_eci,
            ),
        )

    if is_scalar:
        return results[0]
    return results


# -----------------------------------------------------------------------------
# Private helper functions
# -----------------------------------------------------------------------------


def _solve_kepler(
    mean_anomaly: float,
    eccentricity: float,
    tolerance: float = 1e-12,
    max_iterations: int = 50,
) -> float:
    """Solve Kepler's equation M = E - e*sin(E) for eccentric anomaly E.

    Uses Newton-Raphson iteration starting from E = M.

    Args:
        mean_anomaly: Mean anomaly M in radians.
        eccentricity: Orbital eccentricity e (0 <= e < 1 for elliptical).
        tolerance: Convergence tolerance for Newton-Raphson.
        max_iterations: Maximum number of iterations.

    Returns:
        Eccentric anomaly E in radians.

    Raises:
        RuntimeError: If iteration fails to converge.
    """
    # Initial guess
    e_anomaly = mean_anomaly

    for _ in range(max_iterations):
        # f(E) = E - e*sin(E) - M
        # f'(E) = 1 - e*cos(E)
        f_val = e_anomaly - eccentricity * math.sin(e_anomaly) - mean_anomaly
        f_prime = 1.0 - eccentricity * math.cos(e_anomaly)

        delta = f_val / f_prime
        e_anomaly -= delta

        if abs(delta) < tolerance:
            return e_anomaly

    msg = f"Kepler equation failed to converge after {max_iterations} iterations"
    raise RuntimeError(msg)


def _eccentric_to_true_anomaly(eccentric_anomaly: float, eccentricity: float) -> float:
    """Convert eccentric anomaly to true anomaly.

    Args:
        eccentric_anomaly: Eccentric anomaly E in radians.
        eccentricity: Orbital eccentricity e.

    Returns:
        True anomaly nu in radians.
    """
    # nu = 2 * atan2(sqrt(1+e) * sin(E/2), sqrt(1-e) * cos(E/2))
    half_e = eccentric_anomaly / 2.0
    sqrt_1_plus_e = math.sqrt(1.0 + eccentricity)
    sqrt_1_minus_e = math.sqrt(1.0 - eccentricity)

    return 2.0 * math.atan2(
        sqrt_1_plus_e * math.sin(half_e),
        sqrt_1_minus_e * math.cos(half_e),
    )


def _rotation_matrix_313(
    raan: float,
    inclination: float,
    arg_periapsis: float,
) -> NDArray[np.floating]:
    """Compute rotation matrix from perifocal to ECI frame (3-1-3 Euler angles).

    The perifocal frame has x pointing to periapsis and z normal to orbit plane.
    The transformation uses the sequence: R_z(-RAAN) @ R_x(-i) @ R_z(-omega)
    applied to perifocal coordinates to get ECI coordinates.

    Args:
        raan: Right ascension of ascending node (Omega) in radians.
        inclination: Inclination (i) in radians.
        arg_periapsis: Argument of periapsis (omega) in radians.

    Returns:
        3x3 rotation matrix from perifocal to ECI frame.
    """
    cos_raan = math.cos(raan)
    sin_raan = math.sin(raan)
    cos_inc = math.cos(inclination)
    sin_inc = math.sin(inclination)
    cos_argp = math.cos(arg_periapsis)
    sin_argp = math.sin(arg_periapsis)

    # Combined rotation matrix (perifocal to ECI)
    # This is R_z(Omega) @ R_x(i) @ R_z(omega) for the standard convention
    r11 = cos_raan * cos_argp - sin_raan * sin_argp * cos_inc
    r12 = -cos_raan * sin_argp - sin_raan * cos_argp * cos_inc
    r13 = sin_raan * sin_inc

    r21 = sin_raan * cos_argp + cos_raan * sin_argp * cos_inc
    r22 = -sin_raan * sin_argp + cos_raan * cos_argp * cos_inc
    r23 = -cos_raan * sin_inc

    r31 = sin_argp * sin_inc
    r32 = cos_argp * sin_inc
    r33 = cos_inc

    return np.array(
        [
            [r11, r12, r13],
            [r21, r22, r23],
            [r31, r32, r33],
        ],
        dtype=np.float64,
    )
