#!/usr/bin/env python3
"""
Precalculate expected values for test_periapsis_burn.cpp refactoring.
Uses sim_engine.py for physics propagation with maneuver trigger support.
"""

import math
import sys
sys.path.insert(0, "/home/agent/dev/claudes_game")
from scripts.sim_engine import *


def main():
    dt = 60.0
    earth = None
    for b in sim.bodies:
        if b.name == "Earth":
            earth = b
            break

    # =========================================================================
    # Scenario 1: TestSatellite - starting at periapsis, two sequential burns
    # =========================================================================
    sim1 = Simulator("tests/test_periapsis_burn.toml", dt=dt)
    craft1 = sim1.spacecraft[0]  # TestSatellite

    # Initial orbit state
    r0 = vmag(craft1.local_pos)
    v0 = vmag(craft1.local_vel)
    a0 = craft1.orbit.a
    e0 = craft1.orbit.e
    periapsis0 = a0 * (1.0 - e0)
    apoapsis0 = a0 * (1.0 + e0)
    period0 = 2.0 * math.pi * math.sqrt(a0**3 / (G * earth.mass))

    print("// === Scenario 1: TestSatellite - Two sequential periapsis burns ===")
    print(f"// Initial orbit:")
    print(f"//   a = {a0:.4f} m")
    print(f"//   e = {e0:.10f}")
    print(f"//   periapsis = {periapsis0:.4f} m")
    print(f"//   apoapsis = {apoapsis0:.4f} m")
    print(f"//   period = {period0:.4f} s ({period0/3600:.4f} hours)")
    print(f"//   r0 = {r0:.4f} m (should equal periapsis)")
    print(f"//   v0 = {v0:.4f} m/s")
    print(f"//   nu0 = {math.degrees(craft1.orbit.nu):.4f} deg")
    print()

    # First burn fires immediately (nu=0, trigger=0)
    # After burn, orbit changes - compute new elements
    craft1_before = Spacecraft(
        name=craft1.name, mass=craft1.mass, parent_index=craft1.parent_index,
        orbit=OrbitalElements(a=craft1.orbit.a, e=craft1.orbit.e, nu=craft1.orbit.nu,
                              inc=craft1.orbit.inc, Omega=craft1.orbit.Omega, omega=craft1.orbit.omega),
        local_pos=craft1.local_pos, local_vel=craft1.local_vel,
        global_pos=craft1.global_pos, global_vel=craft1.global_vel,
    )

    # Simulate first orbit: first burn fires immediately, then propagate full orbit
    # Need to run enough steps to capture both burns
    # First burn: immediate (step 0)
    # Second burn: after ~1 full orbit from first burn
    total_steps = int(2.5 * period0 / dt)  # ~2.5 orbits

    burn1_time = -1.0
    burn1_radius = -1.0
    burn1_a = -1.0
    burn1_e = -1.0
    burn1_v = -1.0

    burn2_time = -1.0
    burn2_radius = -1.0
    burn2_a = -1.0
    burn2_e = -1.0
    burn2_v = -1.0

    for step in range(total_steps):
        sim1._step()

        # Check if first burn executed
        if sim1.maneuvers[0].executed and burn1_time < 0:
            burn1_time = sim1.time
            burn1_radius = vmag(craft1.local_pos)
            burn1_a = craft1.orbit.a
            burn1_e = craft1.orbit.e
            burn1_v = vmag(craft1.local_vel)
            burn1_pos = craft1.local_pos
            burn1_vel = craft1.local_vel
            b1x, b1y, b1z = burn1_pos
            b1vx, b1vy, b1vz = burn1_vel
            print(f"// First burn at step {step}, t={burn1_time:.1f}s")
            print(f"//   radius = {burn1_radius:.4f} m")
            print(f"//   velocity = {burn1_v:.4f} m/s")
            print(f"//   new a = {burn1_a:.4f} m")
            print(f"//   new e = {burn1_e:.10f}")
            print(f"//   pos = ({b1x:.4f}, {b1y:.4f}, {b1z:.4f}) m")
            print(f"//   vel = ({b1vx:.4f}, {b1vy:.4f}, {b1vz:.4f}) m/s")

        # Check if second burn executed
        if sim1.maneuvers[1].executed and burn2_time < 0:
            burn2_time = sim1.time
            burn2_radius = vmag(craft1.local_pos)
            burn2_a = craft1.orbit.a
            burn2_e = craft1.orbit.e
            burn2_v = vmag(craft1.local_vel)
            burn2_pos = craft1.local_pos
            burn2_vel = craft1.local_vel
            b2x, b2y, b2z = burn2_pos
            b2vx, b2vy, b2vz = burn2_vel
            print(f"// Second burn at step {step}, t={burn2_time:.1f}s")
            print(f"//   radius = {burn2_radius:.4f} m")
            print(f"//   velocity = {burn2_v:.4f} m/s")
            print(f"//   new a = {burn2_a:.4f} m")
            print(f"//   new e = {burn2_e:.10f}")
            print(f"//   pos = ({b2x:.4f}, {b2y:.4f}, {b2z:.4f}) m")
            print(f"//   vel = ({b2vx:.4f}, {b2vy:.4f}, {b2vz:.4f}) m/s")

    print()

    # =========================================================================
    # Scenario 2: TestSatelliteCrossing - starts at nu=pi/2, one burn
    # =========================================================================
    sim2 = Simulator("tests/test_periapsis_burn.toml", dt=dt)
    craft2 = sim2.spacecraft[1]  # TestSatelliteCrossing
    r0_cross = vmag(craft2.local_pos)
    v0_cross = vmag(craft2.local_vel)
    a0_cross = craft2.orbit.a
    e0_cross = craft2.orbit.e
    periapsis_cross = a0_cross * (1.0 - e0_cross)
    period_cross = 2.0 * math.pi * math.sqrt(a0_cross**3 / (G * earth.mass))

    print("// === Scenario 2: TestSatelliteCrossing - Burn crossing from nu=pi/2 ===")
    print(f"// Initial orbit:")
    print(f"//   a = {a0_cross:.4f} m")
    print(f"//   e = {e0_cross:.10f}")
    print(f"//   periapsis = {periapsis_cross:.4f} m")
    print(f"//   period = {period_cross:.4f} s")
    print(f"//   r0 = {r0_cross:.4f} m")
    print(f"//   v0 = {v0_cross:.4f} m/s")
    print(f"//   nu0 = {math.degrees(craft2.orbit.nu):.4f} deg")
    print()

    burn_cross_time = -1.0
    burn_cross_radius = -1.0
    burn_cross_a = -1.0
    burn_cross_e = -1.0
    burn_cross_v = -1.0

    max_steps = int(2.0 * period_cross / dt)
    for step in range(max_steps):
        sim2._step()

        if sim2.maneuvers[2].executed and burn_cross_time < 0:
            burn_cross_time = sim2.time
            burn_cross_radius = vmag(craft2.local_pos)
            burn_cross_a = craft2.orbit.a
            burn_cross_e = craft2.orbit.e
            burn_cross_v = vmag(craft2.local_vel)
            burn_cross_pos = craft2.local_pos
            burn_cross_vel = craft2.local_vel
            bcx, bcy, bcz = burn_cross_pos
            bcvx, bcvy, bcvz = burn_cross_vel
            print(f"// Burn at step {step}, t={burn_cross_time:.1f}s")
            print(f"//   radius = {burn_cross_radius:.4f} m")
            print(f"//   velocity = {burn_cross_v:.4f} m/s")
            print(f"//   new a = {burn_cross_a:.4f} m")
            print(f"//   new e = {burn_cross_e:.10f}")
            print(f"//   pos = ({bcx:.4f}, {bcy:.4f}, {bcz:.4f}) m")
            print(f"//   vel = ({bcvx:.4f}, {bcvy:.4f}, {bcvz:.4f}) m/s")

    print()

    # =========================================================================
    # Summary: Expected values for C++ test embedding
    # =========================================================================
    print("// === SUMMARY: Values for C++ test embedding ===")
    print()
    print("// --- TestSatellite initial orbit ---")
    print(f"// initial_periapsis = {periapsis0:.4f}")
    print(f"// initial_apoapsis = {apoapsis0:.4f}")
    print(f"// initial_radius = {r0:.4f}")
    print(f"// initial_velocity = {v0:.4f}")
    print(f"// initial_period = {period0:.4f}")
    print()

    if burn1_time >= 0:
        print("// --- First burn (TestSatellite) ---")
        print(f"// burn1_time = {burn1_time:.4f}")
        print(f"// burn1_radius = {burn1_radius:.4f}")
        print(f"// burn1_velocity = {burn1_v:.4f}")
        print(f"// burn1_a = {burn1_a:.4f}")
        print(f"// burn1_e = {burn1_e:.10f}")
        print(f"// burn1_pos = ({b1x:.4f}, {b1y:.4f}, {b1z:.4f}) m")
        print(f"// burn1_vel = ({b1vx:.4f}, {b1vy:.4f}, {b1vz:.4f}) m/s")
        print()

    if burn2_time >= 0:
        print("// --- Second burn (TestSatellite) ---")
        print(f"// burn2_time = {burn2_time:.4f}")
        print(f"// burn2_radius = {burn2_radius:.4f}")
        print(f"// burn2_velocity = {burn2_v:.4f}")
        print(f"// burn2_a = {burn2_a:.4f}")
        print(f"// burn2_e = {burn2_e:.10f}")
        print(f"// burn2_pos = ({b2x:.4f}, {b2y:.4f}, {b2z:.4f}) m")
        print(f"// burn2_vel = ({b2vx:.4f}, {b2vy:.4f}, {b2vz:.4f}) m/s")
        if burn1_time >= 0:
            time_between = burn2_time - burn1_time
            print(f"// time_between_burns = {time_between:.4f}")
        print()

    if burn_cross_time >= 0:
        print("// --- Cross burn (TestSatelliteCrossing) ---")
        print(f"// burn_cross_time = {burn_cross_time:.4f}")
        print(f"// burn_cross_radius = {burn_cross_radius:.4f}")
        print(f"// burn_cross_velocity = {burn_cross_v:.4f}")
        print(f"// burn_cross_a = {burn_cross_a:.4f}")
        print(f"// burn_cross_e = {burn_cross_e:.10f}")
        print(f"// burn_cross_pos = ({bcx:.4f}, {bcy:.4f}, {bcz:.4f}) m")
        print(f"// burn_cross_vel = ({bcvx:.4f}, {bcvy:.4f}, {bcvz:.4f}) m/s")
        print()

    # Key assertions
    print("// === Key assertions for test ===")
    print(f"// Periapsis preserved: initial_periapsis ~= final_periapsis (within 1.0)")
    print(f"// Initial radius ~= periapsis: {r0:.4f} ~= {periapsis0:.4f}")
    print(f"// Burn radius ~= periapsis: burn radii should be close to {periapsis0:.4f}")
    print(f"// Two burns at same location: burn1_radius ~= burn2_radius")
    print(f"// Time between burns ~= orbital period")
    print()
    # State vector separation errors (compared to C++ test output)
    def state_vec_dist(p1, v1, p2, v2):
        dr = math.sqrt(sum((a-b)**2 for a,b in zip(p1,p2)))
        dv = math.sqrt(sum((a-b)**2 for a,b in zip(v1,v2)))
        return dr, dv
    burn1_r = (b1x, b1y, b1z)
    burn1_v = (b1vx, b1vy, b1vz)
    burn2_r = (b2x, b2y, b2z)
    burn2_v = (b2vx, b2vy, b2vz)
    ddr1, ddv1 = state_vec_dist(burn1_r, burn1_v, burn1_r, burn1_v)
    print(f"// State vector self-check (burn1 vs burn1): dr={ddr1:.2e} m, dv={ddv1:.2e} m/s")


if __name__ == "__main__":
    # Quick sanity: need to create a dummy sim first to test config loading
    sim = Simulator("tests/test_periapsis_burn.toml", dt=60.0)
    main()