"""
Tests for orbital mechanics, SOI detection, and frame transitions.

Run with: python3 test_orbital.py
or:       python3 -m pytest test_orbital.py -v
"""

import math
import sys
from orbital_elements import OrbitalElements, orbital_elements_from_cartesian, TAU
from soi_calculator import find_soi_crossings
from orbit_drawer import generate_orbit_points
from frame_transition import (
    sun_state_to_planet_orbit,
    planet_orbit_to_sun_state,
    planet_orbit_to_sun_orbit,
)

_assertions = 0
_failures = 0


def check(condition: bool, msg: str = "") -> None:
    global _assertions, _failures
    _assertions += 1
    if not condition:
        _failures += 1
        print(f"  FAIL: {msg}", file=sys.stderr)
    else:
        print(f"  ok: {msg}")


def approx(a: float, b: float, tol: float = 1e-9) -> bool:
    return abs(a - b) < tol


# ═══════════════════════════════════════════════════════════════
#  Elliptical orbit tests (unchanged)
# ═══════════════════════════════════════════════════════════════

def test_circular_orbit_position() -> None:
    print("test_circular_orbit_position")
    orbit = OrbitalElements(a=100.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=1_000_000.0)
    for t in [0.0, 10.0, 50.0, orbit.period / 4, orbit.period / 2]:
        x, y = orbit.position_at(t)
        check(approx(math.hypot(x, y), 100.0), f"t={t:.1f} dist={math.hypot(x,y):.6f}")


def test_elliptical_periapsis_apoapsis() -> None:
    print("test_elliptical_periapsis_apoapsis")
    a, e = 200.0, 0.3
    mu = 500_000.0
    orbit = OrbitalElements(a=a, e=e, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    x, y = orbit.position_at(0.0)
    check(approx(x, a*(1-e)), f"periapsis x={x} == {a*(1-e)}")
    check(approx(y, 0.0), f"periapsis y={y}")
    half_T = orbit.period / 2
    x2, y2 = orbit.position_at(half_T)
    check(approx(x2, -a*(1+e)), f"apoapsis x={x2} == {-a*(1+e)}")
    check(approx(y2, 0.0, 1e-8), f"apoapsis y={y2}")


def test_omega_rotation() -> None:
    print("test_omega_rotation")
    orbit = OrbitalElements(a=100.0, e=0.2, omega=math.pi/2, M0=0.0, epoch=0.0, mu=1_000_000.0)
    x, y = orbit.position_at(0.0)
    check(approx(x, 0.0, 1e-8), f"omega=pi/2: x={x} ≈ 0")
    check(approx(y, 80.0), f"omega=pi/2: y={y} == 80")


def test_period() -> None:
    print("test_period")
    a, mu = 200.0, 1_000_000.0
    orbit = OrbitalElements(a=a, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    expected = 2*math.pi*math.sqrt(a**3/mu)
    check(approx(orbit.period, expected), f"period {orbit.period} == {expected}")
    x0, y0 = orbit.position_at(0.0)
    x1, y1 = orbit.position_at(orbit.period)
    check(approx(x0, x1) and approx(y0, y1), "closed after 1 period")


def test_cartesian_roundtrip() -> None:
    print("test_cartesian_roundtrip")
    orbit1 = OrbitalElements(a=300.0, e=0.25, omega=1.2, M0=2.5, epoch=1000.0, mu=1_500_000.0)
    x, y, vx, vy = orbit1.compute_state_at(2000.0)
    orbit2 = orbital_elements_from_cartesian(x, y, vx, vy, orbit1.mu, epoch=2000.0)
    check(approx(orbit2.a, orbit1.a, 1e-6), f"a: {orbit2.a} == {orbit1.a}")
    check(approx(orbit2.e, orbit1.e, 1e-6), f"e: {orbit2.e} == {orbit1.e}")
    omega_diff = abs(orbit2.omega - orbit1.omega) % TAU
    check(omega_diff < 1e-6 or abs(omega_diff-TAU) < 1e-6, f"ω: {orbit2.omega} ≈ {orbit1.omega}")
    x2, y2, vx2, vy2 = orbit2.compute_state_at(2000.0)
    check(approx(x, x2, 1e-6) and approx(y, y2, 1e-6), "position at epoch")
    check(approx(vx, vx2, 1e-6) and approx(vy, vy2, 1e-6), "velocity at epoch")


def test_orbit_points() -> None:
    print("test_orbit_points")
    orbit = OrbitalElements(a=100.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=100_000.0)
    points = generate_orbit_points(orbit, 128)
    check(len(points) == 128, f"{len(points)} == 128")
    for x, y in points:
        if not approx(math.hypot(x, y), 100.0, 1e-8):
            check(False, f"point at {math.hypot(x,y):.9f} ≠ 100.0")
            break
    else:
        check(True, "all 128 points at correct distance")


def test_radial_overlap_quick_reject() -> None:
    print("test_radial_overlap_quick_reject")
    mu = 1_000_000.0
    inner = OrbitalElements(a=100.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    outer = OrbitalElements(a=500.0, e=0.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    crossings = find_soi_crossings(inner, outer, 30.0, 0.0, 200.0)
    check(len(crossings) == 0, f"non-overlapping: {len(crossings)} crossings")


# ═══════════════════════════════════════════════════════════════
#  SOI crossing tests
# ═══════════════════════════════════════════════════════════════

def test_soi_crossing_detection() -> None:
    print("test_soi_crossing_detection")
    mu = 1_500_000.0
    planet = OrbitalElements(a=350.0, e=0.02, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    body = OrbitalElements(a=320.0, e=0.25, omega=math.pi*0.7, M0=1.0, epoch=0.0, mu=mu)
    soi_r = 50.0
    window = 3*max(body.period, planet.period)
    crossings = find_soi_crossings(body, planet, soi_r, 0.0, window)
    check(len(crossings) > 0, f"found {len(crossings)} crossing(s)")
    for i, (enter_t, exit_t) in enumerate(crossings):
        dur = exit_t - enter_t
        check(enter_t <= exit_t, f"crossing {i}: {enter_t:.3f} ≤ {exit_t:.3f}")
        for label, t in [("enter", enter_t), ("exit", exit_t)]:
            bx, by = body.position_at(t)
            px, py = planet.position_at(t)
            d = math.hypot(bx-px, by-py)
            check(approx(d, soi_r, 0.5), f"{label} d={d:.3f} ≈ {soi_r}")
        if dur > 0.2:
            mid_t = (enter_t+exit_t)/2
            bx, by = body.position_at(mid_t)
            px, py = planet.position_at(mid_t)
            d_mid = math.hypot(bx-px, by-py)
            check(d_mid < soi_r-1.0, f"deep crossing midpoint d={d_mid:.1f}")
        else:
            print(f"  (crossing {i} is a graze, dur={dur:.3f}s)")


def test_deep_soi_crossing() -> None:
    print("test_deep_soi_crossing")
    mu = 1_500_000.0
    planet = OrbitalElements(a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=mu)
    body = OrbitalElements(a=345.0, e=0.02, omega=0.5, M0=-0.35, epoch=0.0, mu=mu)
    soi = 80.0
    window = 5*max(body.period, planet.period)
    crossings = find_soi_crossings(body, planet, soi, 0.0, window, n_steps=800)
    check(len(crossings) >= 1, f"found {len(crossings)} crossing(s)")
    for i, (enter_t, exit_t) in enumerate(crossings):
        dur = exit_t-enter_t
        check(enter_t < exit_t, f"crossing {i}: {enter_t:.2f} < {exit_t:.2f} dur={dur:.2f}s")
        for label, t in [("enter", enter_t), ("exit", exit_t)]:
            bx, by = body.position_at(t)
            px, py = planet.position_at(t)
            d = math.hypot(bx-px, by-py)
            check(approx(d, soi, 1.0), f"{label} d={d:.2f} ≈ {soi}")
        if dur > 0.5:
            mid_t = (enter_t+exit_t)/2
            bx, by = body.position_at(mid_t)
            px, py = planet.position_at(mid_t)
            d_mid = math.hypot(bx-px, by-py)
            check(d_mid < soi-2.0, f"midpoint d={d_mid:.1f} < {soi}")


# ═══════════════════════════════════════════════════════════════
#  Hyperbolic orbit tests
# ═══════════════════════════════════════════════════════════════

def test_hyperbolic_orbit_position() -> None:
    """Body at periapsis at t=0 on a hyperbolic trajectory."""
    print("test_hyperbolic_orbit_position")
    mu = 50_000.0
    # a < 0 for hyperbolas, e > 1
    orbit = OrbitalElements(a=-40.0, e=2.5, omega=0.0, M0=0.0, epoch=0.0, mu=mu)

    # At t=0 (H=0), should be at periapsis
    x, y = orbit.position_at(0.0)
    r_peri = abs(orbit.a)*(orbit.e - 1.0)  # = 40*1.5 = 60
    check(approx(x, r_peri), f"hyperbolic periapsis x={x:.3f} == {r_peri}")
    check(approx(y, 0.0), f"y={y:.3f} ≈ 0")


def test_hyperbolic_velocity() -> None:
    """Velocity at periapsis should match the vis-viva equation."""
    print("test_hyperbolic_velocity")
    mu = 50_000.0
    orbit = OrbitalElements(a=-40.0, e=2.5, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    _, _, vx, vy = orbit.compute_state_at(0.0)

    r_p = abs(orbit.a)*(orbit.e - 1.0)  # 60
    # v_p = sqrt(mu * (2/r_p + 1/|a|))
    v_expected = math.sqrt(mu * (2.0/r_p + 1.0/abs(orbit.a)))
    v_actual = math.hypot(vx, vy)
    check(approx(v_actual, v_expected, 1e-6),
          f"v_peri: {v_actual:.3f} == {v_expected:.3f}")


def test_hyperbolic_cartesian_roundtrip() -> None:
    """Hyperbolic orbit → cartesian → same hyperbolic orbit."""
    print("test_hyperbolic_cartesian_roundtrip")
    mu = 50_000.0
    orbit1 = OrbitalElements(a=-30.0, e=3.0, omega=1.0, M0=0.5, epoch=0.0, mu=mu)
    x, y, vx, vy = orbit1.compute_state_at(10.0)
    orbit2 = orbital_elements_from_cartesian(x, y, vx, vy, mu, epoch=10.0)
    check(approx(orbit2.a, orbit1.a, 1e-6), f"a: {orbit2.a} == {orbit1.a}")
    check(approx(orbit2.e, orbit1.e, 1e-6), f"e: {orbit2.e} == {orbit1.e}")
    omega_diff = abs(orbit2.omega - orbit1.omega) % TAU
    check(omega_diff < 1e-6 or abs(omega_diff-TAU) < 1e-6, f"ω: {orbit2.omega} ≈ {orbit1.omega}")
    x2, y2 = orbit2.position_at(10.0)
    check(approx(x, x2, 1e-6) and approx(y, y2, 1e-6), "position at epoch")


def test_hyperbolic_orbit_points() -> None:
    """generate_orbit_points should truncate at max_radius."""
    print("test_hyperbolic_orbit_points")
    mu = 50_000.0
    orbit = OrbitalElements(a=-20.0, e=2.0, omega=0.0, M0=0.0, epoch=0.0, mu=mu)
    points = generate_orbit_points(orbit, num_points=100, max_radius=80.0)
    check(len(points) > 0, f"generated {len(points)} hyperbolic points")
    for x, y in points:
        d = math.hypot(x, y)
        check(d <= 80.0 + 1e-6, f"point at r={d:.3f} ≤ 80.0")


# ═══════════════════════════════════════════════════════════════
#  Frame transition tests
# ═══════════════════════════════════════════════════════════════

def test_frame_transition_roundtrip() -> None:
    """
    Body in sun frame → planet frame → back → should match original.
    """
    print("test_frame_transition_roundtrip")
    mu_sun = 1_500_000.0
    mu_planet = 30_000.0

    planet = OrbitalElements(a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=mu_sun)
    body_sun = OrbitalElements(a=380.0, e=0.25, omega=2.0, M0=1.0, epoch=0.0, mu=mu_sun)

    t_transition = 50.0

    # Original body position in sun frame
    bx, by, bvx, bvy = body_sun.compute_state_at(t_transition)
    px, py, pvx, pvy = planet.compute_state_at(t_transition)

    # Transition to planet frame
    body_planet = sun_state_to_planet_orbit(
        bx, by, bvx, bvy, px, py, pvx, pvy, mu_planet, t_transition,
    )

    # Body's planet-relative position should be the vector difference
    bx_rel, by_rel = body_planet.position_at(t_transition)
    check(approx(bx_rel, bx-px, 1e-6), f"rel x: {bx_rel:.3f} == {bx-px:.3f}")
    check(approx(by_rel, by-py, 1e-6), f"rel y: {by_rel:.3f} == {by-py:.3f}")

    # Transition back to sun frame
    x_back, y_back, vx_back, vy_back = planet_orbit_to_sun_state(
        body_planet, px, py, pvx, pvy, t_transition,
    )
    check(approx(x_back, bx, 1e-6) and approx(y_back, by, 1e-6), "position roundtrip")
    check(approx(vx_back, bvx, 1e-6) and approx(vy_back, bvy, 1e-6), "velocity roundtrip")


def test_frame_transition_hyperbolic_flyby() -> None:
    """
    Body on planet-crossing orbit → planet frame at SOI entry.
    Verify the planet-relative orbit is hyperbolic (flyby, not capture).
    """
    print("test_frame_transition_hyperbolic_flyby")
    mu_sun = 1_500_000.0
    mu_planet = 10_000.0   # v_rel ≈ 20.4, critical ≈ 16.9 → hyperbolic ✓

    # Planet on near-circular orbit
    planet = OrbitalElements(a=350.0, e=0.02, omega=0.5, M0=0.0, epoch=0.0, mu=mu_sun)

    # Asteroid: crossing orbit (verified to intersect planet's SOI)
    asteroid = OrbitalElements(a=380.0, e=0.35, omega=2.5, M0=0.8, epoch=0.0, mu=mu_sun)

    # Find first SOI crossing
    soi = 70.0
    window = 4 * max(asteroid.period, planet.period)
    crossings = find_soi_crossings(asteroid, planet, soi, 0.0, window, n_steps=1000)

    check(len(crossings) > 0, f"found {len(crossings)} SOI crossing(s)")

    if crossings:
        enter_t = crossings[0][0]

        # Get states at SOI entry
        ax, ay, avx, avy = asteroid.compute_state_at(enter_t)
        px, py, pvx, pvy = planet.compute_state_at(enter_t)

        # Transition to planet frame
        body_planet = sun_state_to_planet_orbit(
            ax, ay, avx, avy, px, py, pvx, pvy, mu_planet, enter_t,
        )

        # This SHOULD be hyperbolic (flyby, energy > 0)
        check(body_planet.is_hyperbolic,
              f"planet-relative orbit is hyperbolic (e={body_planet.e:.3f})")

        # Verify position matches relative to planet
        bx_rel, by_rel = body_planet.position_at(enter_t)
        d_rel = math.hypot(bx_rel, by_rel)
        check(approx(d_rel, soi, 2.0),
              f"rel distance at SOI entry: {d_rel:.1f} ≈ {soi}")

        # Verify the orbit goes out past the SOI (will exit)
        check(body_planet.apoapsis > soi,
              f"hyperbolic: path extends past SOI (max r={body_planet.apoapsis})")


# ═══════════════════════════════════════════════════════════════
#  Runner
# ═══════════════════════════════════════════════════════════════

if __name__ == "__main__":
    print("=== Orbital Mechanics Tests ===\n")

    test_circular_orbit_position()
    test_elliptical_periapsis_apoapsis()
    test_omega_rotation()
    test_period()
    test_cartesian_roundtrip()
    test_orbit_points()
    test_radial_overlap_quick_reject()
    test_soi_crossing_detection()
    test_deep_soi_crossing()
    test_hyperbolic_orbit_position()
    test_hyperbolic_velocity()
    test_hyperbolic_cartesian_roundtrip()
    test_hyperbolic_orbit_points()
    test_frame_transition_roundtrip()
    test_frame_transition_hyperbolic_flyby()

    print(f"\n---\n{_assertions} assertions, {_failures} failures")
    sys.exit(0 if _failures == 0 else 1)