#!/usr/bin/env python3
"""
Precalculate expected values for test_omega_debug.cpp refactoring.
Computes expected orbital elements after a prograde burn at apoapsis.
"""

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


def main():
    earth_mass = 5.972e24
    mu = G * earth_mass

    # Initial orbit: zero inclination, omega = 0, start at apoapsis (nu = pi)
    elements = OrbitalElements(
        a=1.0e7,
        e=0.3,
        nu=math.pi,
        inc=1e-12,
        Omega=0.0,
        omega=0.0,
    )

    pos, vel = orbital_to_cartesian(elements, earth_mass)

    r = vmag(pos)
    v = vmag(vel)

    print("// Initial state")
    print(f"//   pos = ({pos[0]:.15e}, {pos[1]:.15e}, {pos[2]:.15e}) m")
    print(f"//   vel = ({vel[0]:.15e}, {vel[1]:.15e}, {vel[2]:.15e}) m/s")
    print(f"//   r = {r:.15e} m")
    print(f"//   v = {v:.15e} m/s")
    print()

    # Eccentricity vector
    r_dot_v = vdot(pos, vel)
    e_vec = (
        ((v * v - mu / r) * pos[0] - r_dot_v * vel[0]) / mu,
        ((v * v - mu / r) * pos[1] - r_dot_v * vel[1]) / mu,
        ((v * v - mu / r) * pos[2] - r_dot_v * vel[2]) / mu,
    )
    e_mag = vmag(e_vec)

    print(f"//   e_vec_initial = ({e_vec[0]:.15e}, {e_vec[1]:.15e}, {e_vec[2]:.15e})")
    print(f"//   e_initial = {e_mag:.15e}")
    print()

    # Apply prograde burn (1000 m/s)
    burn_dir = get_burn_direction(BurnDirection.PROGRADE, pos, vel)
    dv = 1000.0
    dv_vec = vscale(burn_dir, dv)
    vel_new = vadd(vel, dv_vec)

    v_new = vmag(vel_new)

    print("// After prograde burn")
    print(f"//   burn_dir = ({burn_dir[0]:.15e}, {burn_dir[1]:.15e}, {burn_dir[2]:.15e})")
    print(f"//   vel_new = ({vel_new[0]:.15e}, {vel_new[1]:.15e}, {vel_new[2]:.15e}) m/s")
    print(f"//   v_new = {v_new:.15e} m/s")
    print()

    # Reconstruct orbital elements
    new_elements = cartesian_to_orbital_elements(pos, vel_new, earth_mass)

    print(f"//   new elements:")
    print(f"//     a = {new_elements.a:.15e} m")
    print(f"//     e = {new_elements.e:.15e}")
    print(f"//     nu = {new_elements.nu:.15e} rad ({math.degrees(new_elements.nu):.6f} deg)")
    print(f"//     inc = {new_elements.inc:.15e} rad")
    print(f"//     Omega = {new_elements.Omega:.15e} rad")
    print(f"//     omega = {new_elements.omega:.15e} rad ({math.degrees(new_elements.omega):.6f} deg)")
    print()

    # New eccentricity vector
    r_dot_v_new = vdot(pos, vel_new)
    e_vec_new = (
        ((v_new * v_new - mu / r) * pos[0] - r_dot_v_new * vel_new[0]) / mu,
        ((v_new * v_new - mu / r) * pos[1] - r_dot_v_new * vel_new[1]) / mu,
        ((v_new * v_new - mu / r) * pos[2] - r_dot_v_new * vel_new[2]) / mu,
    )

    print(f"//   e_vec_new = ({e_vec_new[0]:.15e}, {e_vec_new[1]:.15e}, {e_vec_new[2]:.15e})")
    print(f"//   e_new = {vmag(e_vec_new):.15e}")
    print()

    # Verify omega is in [0, 2*pi)
    print("// Omega range check")
    print(f"//   omega = {new_elements.omega:.15e} rad")
    print(f"//   omega in [0, 2*pi)? {0.0 <= new_elements.omega < 2.0 * math.pi}")
    print()

    # After a prograde burn at apoapsis (nu=pi), the eccentricity vector flips
    # direction because the increased velocity raises the opposite side of the orbit.
    # This means omega should change from 0 to approximately pi.
    #
    # Rationale: at apoapsis, position and velocity are perpendicular.
    # A prograde burn adds velocity along the velocity direction, increasing energy.
    # The eccentricity vector formula: e_vec = (v^2 - mu/r)*r/μ - (r·v)*v/μ
    # At apoapsis: r·v = 0, so e_vec = (v^2 - mu/r) * r / mu
    # After prograde burn, v increases, so (v^2 - mu/r) becomes more positive,
    # making e_vec more aligned with r direction.
    # Since at apoapsis, r points opposite to periapsis direction (for ω=0),
    # the eccentricity vector flips, meaning periapsis moves to the opposite side,
    # so ω → π.

    print("// Expected test values")
    print(f"// a_expected = {new_elements.a:.15e}")
    print(f"// e_expected = {new_elements.e:.15e}")
    print(f"// omega_expected = {new_elements.omega:.15e} rad ({math.degrees(new_elements.omega):.6f} deg)")
    print()

    # Also compute expected values using the same check as the original test
    print("// For WithinAbs assertions:")
    print(f"// a := {new_elements.a:.15e}")
    print(f"// e := {new_elements.e:.15e}")
    print(f"// omega := {new_elements.omega:.15e}")
    print(f"// inc := {new_elements.inc:.15e}")
    print(f"// Omega := {new_elements.Omega:.15e}")
    print(f"// nu := {new_elements.nu:.15e}")
    print(f"// r := {r:.15e}")
    print(f"// v_new := {v_new:.15e}")


if __name__ == "__main__":
    main()