#!/usr/bin/env python3
"""
Precalculate expected values for test_analytical_propagation.cpp.
Loads config from tests/test_analytical_propagation.toml, then computes
orbital parameters, propagation results, and error bounds.
"""

import math
import sys
sys.path.insert(0, "scripts")
from sim_engine import Simulator, propagate, orbital_to_cartesian, vmag, G

def main():
    sim = Simulator("tests/test_analytical_propagation.toml", dt=60.0)

    earth = sim.get_body("Earth")
    craft_apsides = sim.get_craft("Apsides_Test_Spacecraft")
    craft_timestep = sim.get_craft("Timestep_Test_Spacecraft")

    earth_mass = earth.mass
    mu = G * earth_mass

    a1 = craft_apsides.orbit.a
    e1 = craft_apsides.orbit.e
    period1 = 2.0 * math.pi * math.sqrt(a1**3 / mu)
    n1 = math.sqrt(mu / a1**3)

    a2 = craft_timestep.orbit.a
    e2 = craft_timestep.orbit.e
    period2 = 2.0 * math.pi * math.sqrt(a2**3 / mu)
    n2 = math.sqrt(mu / a2**3)

    def print_comment_block(title):
        print(f"\n// === {title} ===")

    def print_const(name, value, comment=""):
        c = f"  // {comment}" if comment else ""
        print(f"const double {name} = {value:.15e};{c}")

    # =============================================================================
    # 1. Apsides geometry — both spacecraft
    # =============================================================================

    print_comment_block("Apsides geometry (both spacecraft)")

    # Apsides spacecraft
    r_peri1 = a1 * (1.0 - e1)
    r_apo1 = a1 * (1.0 + e1)
    peri1 = craft_apsides.orbit
    _, vel_peri1 = orbital_to_cartesian(peri1, earth_mass)
    v_peri1 = vmag(vel_peri1)
    apo1_el = type(peri1)(a=a1, e=e1, nu=math.pi, inc=0.0, Omega=0.0, omega=0.0)
    _, vel_apo1 = orbital_to_cartesian(apo1_el, earth_mass)
    v_apo1 = vmag(vel_apo1)
    nu45_1_el = type(peri1)(a=a1, e=e1, nu=math.pi/4.0, inc=0.0, Omega=0.0, omega=0.0)
    _, vel_45_1 = orbital_to_cartesian(nu45_1_el, earth_mass)
    v_45_1 = vmag(vel_45_1)

    print(f"// Apsides spacecraft: a={a1:.0f}, e={e1}, period={period1:.2f}s")
    print(f"// Mean motion: {n1:.15e} rad/s")
    print(f"// r_peri={r_peri1:.3f} m, r_apo={r_apo1:.3f} m")
    print(f"// v_peri={v_peri1:.6f} m/s, v_apo={v_apo1:.6f} m/s")
    print(f"// v_at_pi4={v_45_1:.6f} m/s")

    # Timestep spacecraft
    r_peri2 = a2 * (1.0 - e2)
    r_apo2 = a2 * (1.0 + e2)
    peri2 = craft_timestep.orbit
    _, vel_peri2 = orbital_to_cartesian(peri2, earth_mass)
    v_peri2 = vmag(vel_peri2)
    apo2_el = type(peri2)(a=a2, e=e2, nu=math.pi, inc=0.0, Omega=0.0, omega=0.0)
    _, vel_apo2 = orbital_to_cartesian(apo2_el, earth_mass)
    v_apo2 = vmag(vel_apo2)

    print(f"// Timestep spacecraft: a={a2:.0f}, e={e2}, period={period2:.2f}s")
    print(f"// Mean motion: {n2:.15e} rad/s")
    print(f"// r_peri={r_peri2:.3f} m, r_apo={r_apo2:.3f} m")
    print(f"// v_peri={v_peri2:.6f} m/s, v_apo={v_apo2:.6f} m/s")

    print_const("A1_R_PERI", r_peri1, "m")
    print_const("A1_R_APO", r_apo1, "m")
    print_const("A1_V_PERI", v_peri1, "m/s")
    print_const("A1_V_APO", v_apo1, "m/s")
    print_const("A1_V_AT_PI4", v_45_1, "m/s at nu=pi/4")
    print_const("A1_PERIOD", period1, "seconds")
    print_const("A2_R_PERI", r_peri2, "m")
    print_const("A2_R_APO", r_apo2, "m")
    print_const("A2_V_PERI", v_peri2, "m/s")
    print_const("A2_V_APO", v_apo2, "m/s")
    print_const("A2_PERIOD", period2, "seconds")

    # =============================================================================
    # 2. Vis-viva checks at multiple true anomalies
    # =============================================================================

    print_comment_block("Vis-viva checks at multiple true anomalies")

    true_anomalies = [0.0, math.pi/4.0, math.pi/2.0, 3.0*math.pi/4.0, math.pi]
    for nu in true_anomalies:
        deg = nu * 180.0 / math.pi
        el = type(craft_apsides.orbit)(a=a1, e=e1, nu=nu, inc=0.0, Omega=0.0, omega=0.0)
        pos, vel = orbital_to_cartesian(el, earth_mass)
        r = vmag(pos)
        v = vmag(vel)
        expected_v = math.sqrt(mu * (2.0/r - 1.0/a1))
        v_error = abs(v - expected_v)
        rel_error = v_error / expected_v * 100.0
        print(f"// nu={deg:6.1f}deg: r={r:.3f} m, v={v:.6f} m/s, expected_v={expected_v:.6f} m/s, rel_err={rel_error:.8f}%")

    # =============================================================================
    # 3. Period return — full orbit closure for both spacecraft
    # =============================================================================

    print_comment_block("Period return — full orbit closure")

    # Apsides spacecraft: propagate 1 period from nu=0
    el1 = type(craft_apsides.orbit)(a=a1, e=e1, nu=0.0, inc=0.0, Omega=0.0, omega=0.0)
    _, vel1_init = orbital_to_cartesian(el1, earth_mass)
    prop1 = propagate(el1, period1, earth_mass)
    _, vel1_final = orbital_to_cartesian(prop1, earth_mass)
    vel_change1 = math.sqrt((vel1_final[0]-vel1_init[0])**2 + (vel1_final[1]-vel1_init[1])**2 + (vel1_final[2]-vel1_init[2])**2)
    print(f"// Apsides after 1 period: vel_change={vel_change1:.15e} m/s, final_nu={prop1.nu:.15e} rad")

    # Timestep spacecraft: propagate 1 period from nu=0
    el2 = type(craft_timestep.orbit)(a=a2, e=e2, nu=0.0, inc=0.0, Omega=0.0, omega=0.0)
    _, vel2_init = orbital_to_cartesian(el2, earth_mass)
    prop2 = propagate(el2, period2, earth_mass)
    _, vel2_final = orbital_to_cartesian(prop2, earth_mass)
    vel_change2 = math.sqrt((vel2_final[0]-vel2_init[0])**2 + (vel2_final[1]-vel2_init[1])**2 + (vel2_final[2]-vel2_init[2])**2)
    print(f"// Timestep after 1 period: vel_change={vel_change2:.15e} m/s, final_nu={prop2.nu:.15e} rad")

    # =============================================================================
    # 4. Timestep accuracy
    # =============================================================================

    print_comment_block("Timestep accuracy")

    # Initial state for timestep craft
    init_el = type(craft_timestep.orbit)(a=a2, e=e2, nu=0.0, inc=0.0, Omega=0.0, omega=0.0)
    init_pos, init_vel = orbital_to_cartesian(init_el, earth_mass)
    init_r = vmag(init_pos)
    init_v = vmag(init_vel)

    # Large timestep: 2x period
    large_dt = period2 * 2.0
    prop_large = propagate(init_el, large_dt, earth_mass)
    pos_large, vel_large = orbital_to_cartesian(prop_large, earth_mass)
    r_large = vmag(pos_large)
    v_large = vmag(vel_large)
    r_err_large = abs(r_large - init_r)
    v_err_large = abs(v_large - init_v)
    rel_r_large = r_err_large / init_r * 100.0
    rel_v_large = v_err_large / init_v * 100.0
    print(f"// 2x period: r_err={r_err_large:.6f} m ({rel_r_large:.8f}%), v_err={v_err_large:.6f} m/s ({rel_v_large:.8f}%)")

    # Small timestep: 0.1 s
    small_dt = 0.1
    prop_small = propagate(init_el, small_dt, earth_mass)
    pos_small, vel_small = orbital_to_cartesian(prop_small, earth_mass)
    pos_change = math.sqrt((pos_small[0]-init_pos[0])**2 + (pos_small[1]-init_pos[1])**2 + (pos_small[2]-init_pos[2])**2)
    vel_change = math.sqrt((vel_small[0]-init_vel[0])**2 + (vel_small[1]-init_vel[1])**2 + (vel_small[2]-init_vel[2])**2)
    expected_pos_change = init_v * small_dt
    pos_error_small = abs(pos_change - expected_pos_change)
    print(f"// 0.1s dt: pos_change={pos_change:.6f} m, vel_change={vel_change:.10f} m/s")
    print(f"// expected_pos_change={expected_pos_change:.6f} m, pos_error={pos_error_small:.6f} m")

    # Accuracy at various multiples of period
    dt_ratios = [1.0, 10.0]
    for ratio in dt_ratios:
        dt = period2 * ratio
        prop = propagate(init_el, dt, earth_mass)
        pos_f, vel_f = orbital_to_cartesian(prop, earth_mass)
        pos_err = math.sqrt((pos_f[0]-init_pos[0])**2 + (pos_f[1]-init_pos[1])**2 + (pos_f[2]-init_pos[2])**2)
        vel_err = math.sqrt((vel_f[0]-init_vel[0])**2 + (vel_f[1]-init_vel[1])**2 + (vel_f[2]-init_vel[2])**2)
        print(f"// {ratio:.0f}x period: pos_err={pos_err:.6f} m, vel_err={vel_err:.10f} m/s")

    # =============================================================================
    # 5. Long-term stability (100 periods)
    # =============================================================================

    print_comment_block("Long-term stability (100 periods)")

    prop_100 = propagate(init_el, period2 * 100.0, earth_mass)
    final_nu = prop_100.nu
    expected_delta_nu = n2 * period2 * 100.0
    expected_nu = init_el.nu + expected_delta_nu
    # Normalize both to [0, 2*pi)
    while final_nu < 0:
        final_nu += 2.0 * math.pi
    while final_nu >= 2.0 * math.pi:
        final_nu -= 2.0 * math.pi
    while expected_nu < 0:
        expected_nu += 2.0 * math.pi
    while expected_nu >= 2.0 * math.pi:
        expected_nu -= 2.0 * math.pi

    raw_error = abs(final_nu - expected_nu)
    anomaly_error = min(raw_error, 2.0 * math.pi - raw_error)
    print(f"// Propagation time: {period2*100.0:.2f} s ({100.0} periods)")
    print(f"// final_nu={final_nu:.15e} rad")
    print(f"// expected_nu={expected_nu:.15e} rad")
    print(f"// raw_error={raw_error:.15e} rad")
    print(f"// anomaly_error={anomaly_error:.15e} rad ({anomaly_error*180/math.pi:.10e} degrees)")

    # =============================================================================
    # Output summary
    # =============================================================================

    print("\n// === SUMMARY ===")
    print(f"// Apsides spacecraft: a={a1:.0f}, e={e1}, period={period1:.2f}s")
    print(f"// Timestep spacecraft: a={a2:.0f}, e={e2}, period={period2:.2f}s")
    print(f"// Vis-viva relative errors are all < 0.01%")
    print(f"// Full orbit position/velocity errors are < 0.1%")
    print(f"// Long-term (100 periods) anomaly error: {anomaly_error:.15e} rad")

if __name__ == "__main__":
    main()