"""
2D Keplerian orbital elements — elliptical and hyperbolic.

Provides OrbitalElements dataclass with cached derived quantities for
efficient position/velocity computation, plus cartesian ↔ orbital
conversion.

Supports both elliptical (e < 1, a > 0) and hyperbolic (e > 1, a < 0)
orbits using the standard Keplerian formulation.
"""

import math
from dataclasses import dataclass
from typing import Tuple

TAU = 2.0 * math.pi
KEPLER_TOL = 1e-12
KEPLER_MAX_ITER = 50


@dataclass
class OrbitalElements:
    """
    2D Keplerian orbital elements.

    a       — semi-major axis.  a > 0 for ellipses,  a < 0 for hyperbolas.
    e       — eccentricity.  e ≥ 0.  e < 1 = ellipse,  e > 1 = hyperbola.
    omega   — argument of periapsis (rad), +x → periapsis.
    M0      — mean anomaly at epoch (rad).  Wrapped to [0, 2π) for ellipses;
              unbounded (can be any real) for hyperbolas.
    epoch   — reference time.
    mu      — gravitational parameter of central body.
    """
    a: float
    e: float
    omega: float
    M0: float
    epoch: float
    mu: float

    def __post_init__(self) -> None:
        if self.e < 0.0:
            raise ValueError("Eccentricity must be ≥ 0")

        self._hyperbolic: bool = self.e > 1.0

        if self._hyperbolic:
            if self.a > 0:
                raise ValueError("Hyperbolic orbits require a < 0")
            self._abs_a = abs(self.a)
            self._p = self._abs_a * (self.e * self.e - 1.0)  # semi-latus rectum
            self._sqrt_mu_over_p = math.sqrt(self.mu / self._p)
            self._n = math.sqrt(self.mu / self._abs_a ** 3)  # mean motion
            self._sqrt_em1 = math.sqrt(self.e * self.e - 1.0)
        else:
            if self.a <= 0:
                raise ValueError("Elliptical orbits require a > 0")
            self._p = self.a * (1.0 - self.e * self.e)  # semi-latus rectum
            self._sqrt_mu_over_p = math.sqrt(self.mu / self._p)
            self._n = math.sqrt(self.mu / self.a ** 3)   # mean motion
            self._sqrt_1me2 = math.sqrt(1.0 - self.e * self.e)
            self._sqrt_mu_a = math.sqrt(self.mu * self.a)

    # ── Convenience properties ──────────────────────────────────

    @property
    def period(self) -> float:
        """Orbital period (s).  +inf for parabolas, meaningless for hyperbolas."""
        if self._hyperbolic:
            return float("inf")
        return TAU / self._n

    @property
    def periapsis(self) -> float:
        """Distance at closest approach."""
        if self._hyperbolic:
            return self._abs_a * (self.e - 1.0)
        return self.a * (1.0 - self.e)

    @property
    def apoapsis(self) -> float:
        """Distance at farthest point (ellipses only)."""
        if self._hyperbolic:
            return float("inf")
        return self.a * (1.0 + self.e)

    @property
    def is_hyperbolic(self) -> bool:
        return self._hyperbolic

    def radius_range(self) -> Tuple[float, float]:
        """(minimum, maximum) distance.  Max is inf for hyperbolas."""
        return self.periapsis, self.apoapsis

    # ── Kepler solvers ──────────────────────────────────────────

    def solve_kepler(self, M: float) -> float:
        """Solve Kepler's equation.  Returns E (ellipse) or H (hyperbola)."""
        if self._hyperbolic:
            return self._solve_kepler_hyperbolic(M)
        return self._solve_kepler_elliptic(M)

    def _solve_kepler_elliptic(self, M: float) -> float:
        """M = E - e·sin(E)  →  E."""
        e = self.e
        if e < 1e-14:
            return M
        if e < 0.8:
            E = M
        else:
            E = M + 0.85 * e * math.sin(M)
        for _ in range(KEPLER_MAX_ITER):
            dE = (E - e * math.sin(E) - M) / (1.0 - e * math.cos(E))
            E -= dE
            if abs(dE) < KEPLER_TOL:
                break
        return E

    def _solve_kepler_hyperbolic(self, M: float) -> float:
        """M = e·sinh(H) − H  →  H."""
        e = self.e
        if abs(M) < 1e-14:
            return 0.0
        # Robust initial guess
        H = math.asinh(M / e)
        for _ in range(KEPLER_MAX_ITER):
            dH = (e * math.sinh(H) - H - M) / (e * math.cosh(H) - 1.0)
            H -= dH
            if abs(dH) < KEPLER_TOL:
                break
        return H

    # ── Mean anomaly at time ────────────────────────────────────

    def mean_anomaly_at(self, t: float) -> float:
        """Mean anomaly M at time t."""
        M = self.M0 + self._n * (t - self.epoch)
        if self._hyperbolic:
            return M           # unbounded — can be any real
        return M % TAU          # wrap to [0, 2π)

    # ── Position / state computation ────────────────────────────

    def compute_state_at(self, t: float) -> Tuple[float, float, float, float]:
        """World-space (x, y, vx, vy) at time t."""
        M = self.mean_anomaly_at(t)
        anomaly = self.solve_kepler(M)  # E for ellipse, H for hyperbola
        if self._hyperbolic:
            return self._state_from_H(anomaly)
        return self._state_from_E(anomaly)

    def position_at(self, t: float) -> Tuple[float, float]:
        """World-space (x, y) at time t."""
        x, y, _, _ = self.compute_state_at(t)
        return x, y

    # ── Elliptical: state from eccentric anomaly E ─────────────

    def _state_from_E(self, E: float) -> Tuple[float, float, float, float]:
        cosE = math.cos(E)
        sinE = math.sin(E)

        # Perifocal position
        x_pf = self.a * (cosE - self.e)
        y_pf = self.a * self._sqrt_1me2 * sinE

        # Perifocal velocity
        inv_denom = 1.0 / (1.0 - self.e * cosE)
        sqrt_mu_over_a = math.sqrt(self.mu / self.a)
        vx_pf = -sqrt_mu_over_a * sinE * inv_denom
        vy_pf = sqrt_mu_over_a * self._sqrt_1me2 * cosE * inv_denom

        return self._rotate_to_world(x_pf, y_pf, vx_pf, vy_pf)

    def position_from_E(self, E: float) -> Tuple[float, float]:
        cosE = math.cos(E)
        sinE = math.sin(E)
        x_pf = self.a * (cosE - self.e)
        y_pf = self.a * self._sqrt_1me2 * sinE
        cos_om = math.cos(self.omega)
        sin_om = math.sin(self.omega)
        return (
            x_pf * cos_om - y_pf * sin_om,
            x_pf * sin_om + y_pf * cos_om,
        )

    # ── Hyperbolic: state from hyperbolic anomaly H ────────────

    def _state_from_H(self, H: float) -> Tuple[float, float, float, float]:
        coshH = math.cosh(H)
        sinhH = math.sinh(H)

        # Perifocal position  x = |a|(e − coshH),  y = |a|√(e²−1)·sinhH
        x_pf = self._abs_a * (self.e - coshH)
        y_pf = self._abs_a * self._sqrt_em1 * sinhH

        # Distance  r = |a|(e·coshH − 1)
        r = self._abs_a * (self.e * coshH - 1.0)

        # True anomaly trig
        cos_nu = x_pf / r
        sin_nu = y_pf / r

        # Perifocal velocity (unified ν-formula, works for both orbit types)
        vx_pf = -self._sqrt_mu_over_p * sin_nu
        vy_pf = self._sqrt_mu_over_p * (cos_nu + self.e)

        return self._rotate_to_world(x_pf, y_pf, vx_pf, vy_pf)

    def position_from_H(self, H: float) -> Tuple[float, float]:
        coshH = math.cosh(H)
        sinhH = math.sinh(H)
        x_pf = self._abs_a * (self.e - coshH)
        y_pf = self._abs_a * self._sqrt_em1 * sinhH
        cos_om = math.cos(self.omega)
        sin_om = math.sin(self.omega)
        return (
            x_pf * cos_om - y_pf * sin_om,
            x_pf * sin_om + y_pf * cos_om,
        )

    # ── Rotate perifocal → world frame ─────────────────────────

    def _rotate_to_world(
        self, x_pf: float, y_pf: float, vx_pf: float, vy_pf: float,
    ) -> Tuple[float, float, float, float]:
        cos_om = math.cos(self.omega)
        sin_om = math.sin(self.omega)
        return (
            x_pf * cos_om - y_pf * sin_om,
            x_pf * sin_om + y_pf * cos_om,
            vx_pf * cos_om - vy_pf * sin_om,
            vx_pf * sin_om + vy_pf * cos_om,
        )


# ── Cartesian → Orbital Elements ─────────────────────────────────


def orbital_elements_from_cartesian(
    x: float,
    y: float,
    vx: float,
    vy: float,
    mu: float,
    epoch: float = 0.0,
) -> OrbitalElements:
    """
    Convert (x, y, vx, vy) to Keplerian orbital elements.

    Handles both elliptical (energy < 0) and hyperbolic (energy > 0)
    trajectories.
    """
    r = math.hypot(x, y)
    v_sq = vx * vx + vy * vy
    energy = 0.5 * v_sq - mu / r
    r_dot_v = x * vx + y * vy

    # Semi-major axis (negative for hyperbolas)
    if abs(energy) < 1e-15:
        raise ValueError("Parabolic trajectory not supported (energy ≈ 0)")
    a = -mu / (2.0 * energy)   # negative for hyperbolic, positive for elliptical

    # Eccentricity vector
    factor = v_sq / mu - 1.0 / r
    ex = factor * x - r_dot_v * vx / mu
    ey = factor * y - r_dot_v * vy / mu
    e = math.hypot(ex, ey)

    # Argument of periapsis
    omega = math.atan2(ey, ex) % TAU

    # True anomaly
    if e > 1e-14:
        cos_nu = (ex * x + ey * y) / (e * r)
        cos_nu = max(-1.0, min(1.0, cos_nu))
        nu = math.acos(cos_nu)
        if r_dot_v < 0:
            nu = -nu
    else:
        nu = 0.0

    # Mean anomaly
    if e < 1.0:  # elliptical
        cos_E = (e + math.cos(nu)) / (1.0 + e * math.cos(nu))
        cos_E = max(-1.0, min(1.0, cos_E))
        sin_E = math.sqrt(max(0.0, 1.0 - cos_E * cos_E))
        if nu < 0:
            sin_E = -sin_E
        E = math.atan2(sin_E, cos_E)
        M0 = (E - e * math.sin(E)) % TAU
    else:  # hyperbolic
        # cosh(H) = (e + cos(ν)) / (1 + e·cos(ν))
        cosh_H = (e + math.cos(nu)) / (1.0 + e * math.cos(nu))
        cosh_H = max(1.0, cosh_H)
        H = math.acosh(cosh_H)
        if nu < 0:
            H = -H
        M0 = e * math.sinh(H) - H  # unbounded, no wrap

    return OrbitalElements(
        a=a, e=e, omega=omega, M0=M0, epoch=epoch, mu=mu,
    )