"""Convert state vector to orbital elements.

This module implements the inverse of orbit propagation: given position and
velocity vectors in the ECI frame at a specific time, compute the classical
Keplerian orbital elements.

The algorithm uses vector algebra to compute the orbital angular momentum,
eccentricity vector, and node vector, from which the six classical elements
are derived.

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
from .orbital_elements import OrbitalElements

if TYPE_CHECKING:
    from numpy.typing import NDArray

TOLERANCE = 1e-10
"""Numerical tolerance for near-zero checks to avoid division by zero."""


def state_to_elements(  # noqa: C901, PLR0912, PLR0915
    state: OrbitState,
    unix_s: float,
) -> OrbitalElements:
    """Convert state vector to orbital elements.

    Computes classical Keplerian orbital elements from position and velocity
    vectors using standard orbital mechanics vector algebra.

    Args:
        state: OrbitState containing position and velocity in ECI frame.
            The acceleration component is not used.
        unix_s: Time of the state as Unix timestamp in seconds. This becomes
            the epoch of the resulting orbital elements.

    Returns:
        OrbitalElements representing the orbit, with mean_anomaly_rad set to
        the value at the specified epoch time.

    Raises:
        ValueError: If the orbit is parabolic/hyperbolic (e >= 1) or if the
            state represents a degenerate case (zero angular momentum).

    Example:
        >>> import numpy as np
        >>> from gri_utils.orbit import OrbitState, state_to_elements
        >>> # Circular equatorial orbit
        >>> state = OrbitState(
        ...     position_eci_m=np.array([7_000_000.0, 0.0, 0.0]),
        ...     velocity_eci_ms=np.array([0.0, 7546.0, 0.0]),
        ...     acceleration_eci_ms2=np.array([0.0, 0.0, 0.0]),  # Not used
        ... )
        >>> elements = state_to_elements(state, unix_s=0.0)
        >>> print(f"Semi-major axis: {elements.semi_major_axis_m:.0f} m")
    """
    r_vec = np.asarray(state.position_eci_m, dtype=np.float64)
    v_vec = np.asarray(state.velocity_eci_ms, dtype=np.float64)

    # Position and velocity magnitudes
    r_mag = float(np.linalg.norm(r_vec))
    v_mag = float(np.linalg.norm(v_vec))

    # Angular momentum vector: h = r x v
    h_vec = np.cross(r_vec, v_vec)
    h_mag = float(np.linalg.norm(h_vec))

    if h_mag < TOLERANCE:
        msg = "Degenerate orbit: angular momentum is near zero"
        raise ValueError(msg)

    # Node vector: n = k x h (k is unit vector along z-axis)
    k_vec = np.array([0.0, 0.0, 1.0])
    n_vec = np.cross(k_vec, h_vec)
    n_mag = float(np.linalg.norm(n_vec))

    # Eccentricity vector: e_vec = (v x h)/mu - r/|r|
    e_vec = np.cross(v_vec, h_vec) / GM_EARTH - r_vec / r_mag
    e = float(np.linalg.norm(e_vec))

    if e >= 1.0:
        msg = f"Non-elliptical orbit: eccentricity = {e:.6f} >= 1"
        raise ValueError(msg)

    # Specific orbital energy: epsilon = v^2/2 - mu/r
    specific_energy = v_mag * v_mag / 2.0 - GM_EARTH / r_mag

    # Semi-major axis: a = -mu / (2*epsilon)
    a = -GM_EARTH / (2.0 * specific_energy)

    # Inclination: i = acos(h_z / |h|)
    inclination = math.acos(np.clip(h_vec[2] / h_mag, -1.0, 1.0))

    # Right ascension of ascending node (RAAN)
    if n_mag > TOLERANCE:
        # Normal case: orbit is inclined
        raan = math.acos(np.clip(n_vec[0] / n_mag, -1.0, 1.0))
        if n_vec[1] < 0:
            raan = 2.0 * math.pi - raan
    else:
        # Equatorial orbit: RAAN is undefined, set to 0
        raan = 0.0

    # Argument of periapsis
    if e > TOLERANCE:
        if n_mag > TOLERANCE:
            # Normal case
            cos_omega = np.dot(n_vec, e_vec) / (n_mag * e)
            arg_periapsis = math.acos(np.clip(cos_omega, -1.0, 1.0))
            if e_vec[2] < 0:
                arg_periapsis = 2.0 * math.pi - arg_periapsis
        else:
            # Equatorial orbit: measure from x-axis
            arg_periapsis = math.atan2(e_vec[1], e_vec[0])
            if arg_periapsis < 0:
                arg_periapsis += 2.0 * math.pi
    else:
        # Circular orbit: argument of periapsis is undefined, set to 0
        arg_periapsis = 0.0

    # True anomaly
    if e > TOLERANCE:
        cos_nu = np.dot(e_vec, r_vec) / (e * r_mag)
        true_anomaly = math.acos(np.clip(cos_nu, -1.0, 1.0))
        # Check sign using r dot v
        if np.dot(r_vec, v_vec) < 0:
            true_anomaly = 2.0 * math.pi - true_anomaly
    elif n_mag > TOLERANCE:
        # Circular inclined orbit: measure from ascending node
        cos_nu = np.dot(n_vec, r_vec) / (n_mag * r_mag)
        true_anomaly = math.acos(np.clip(cos_nu, -1.0, 1.0))
        if r_vec[2] < 0:
            true_anomaly = 2.0 * math.pi - true_anomaly
    else:
        # Circular equatorial orbit: measure from x-axis
        true_anomaly = math.atan2(r_vec[1], r_vec[0])
        if true_anomaly < 0:
            true_anomaly += 2.0 * math.pi

    # Convert true anomaly to mean anomaly via eccentric anomaly
    # E = 2 * atan(sqrt((1-e)/(1+e)) * tan(nu/2))
    if e > TOLERANCE:
        tan_half_nu = math.tan(true_anomaly / 2.0)
        sqrt_term = math.sqrt((1.0 - e) / (1.0 + e))
        eccentric_anomaly = 2.0 * math.atan(sqrt_term * tan_half_nu)

        # M = E - e*sin(E)
        mean_anomaly = eccentric_anomaly - e * math.sin(eccentric_anomaly)

        # Normalize to [0, 2*pi)
        mean_anomaly = mean_anomaly % (2.0 * math.pi)
    else:
        # Circular orbit: mean anomaly equals true anomaly
        mean_anomaly = true_anomaly

    return OrbitalElements(
        semi_major_axis_m=a,
        eccentricity=e,
        inclination_rad=inclination,
        raan_rad=raan,
        arg_periapsis_rad=arg_periapsis,
        mean_anomaly_rad=mean_anomaly,
        epoch_unix_s=unix_s,
    )


def position_velocity_to_elements(
    position_eci_m: NDArray[np.floating],
    velocity_eci_ms: NDArray[np.floating],
    unix_s: float,
) -> OrbitalElements:
    """Convert position and velocity vectors to orbital elements.

    Convenience function that creates an OrbitState internally and calls
    state_to_elements. The acceleration is set to zero as it's not used.

    Args:
        position_eci_m: Position vector in ECI frame, meters. Shape (3,).
        velocity_eci_ms: Velocity vector in ECI frame, m/s. Shape (3,).
        unix_s: Time of the state as Unix timestamp in seconds.

    Returns:
        OrbitalElements representing the orbit.

    Example:
        >>> import numpy as np
        >>> from gri_utils.orbit import position_velocity_to_elements
        >>> elements = position_velocity_to_elements(
        ...     position_eci_m=np.array([7_000_000.0, 0.0, 0.0]),
        ...     velocity_eci_ms=np.array([0.0, 7546.0, 0.0]),
        ...     unix_s=0.0,
        ... )
    """
    state = OrbitState(
        position_eci_m=np.asarray(position_eci_m),
        velocity_eci_ms=np.asarray(velocity_eci_ms),
        acceleration_eci_ms2=np.zeros(3),
    )
    return state_to_elements(state, unix_s)
