#!/usr/bin/env python3
"""
Precalculate expected values for test_orbital_period.cpp.

Measures:
1. Earth orbital period (seconds, days) — track global angle for circular orbit
2. Mars orbital period (seconds, days)
3. Direction test: prograde check over 1 day
"""

import sys
import math

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

MAX_STEPS = 1_100_000  # safety limit (687 days × 1440 steps/day)
DT = 60.0


def measure_period(sim, body_name, parent_mass, analytical_days):
    """
    Measure period by tracking global angle for one full revolution.
    For circular orbits, nu stays at 0 so we track atan2(y, x) instead.
    """
    body = sim.get_body(body_name)
    parent = sim.get_body(body.parent_index) if body.parent_index >= 0 else None

    # Track global angle
    if parent:
        angle_start = math.atan2(
            body.global_pos[1] - parent.global_pos[1],
            body.global_pos[0] - parent.global_pos[0]
        )
    else:
        angle_start = math.atan2(body.global_pos[1], body.global_pos[0])

    total_angle = 0.0
    prev_angle = angle_start

    for step in range(1, MAX_STEPS + 1):
        sim._step()

        if parent:
            angle = math.atan2(
                body.global_pos[1] - parent.global_pos[1],
                body.global_pos[0] - parent.global_pos[0]
            )
        else:
            angle = math.atan2(body.global_pos[1], body.global_pos[0])

        # Accumulate angle (handle wrap)
        delta = angle - prev_angle
        if delta > math.pi:
            delta -= 2 * math.pi
        elif delta < -math.pi:
            delta += 2 * math.pi
        total_angle += delta
        prev_angle = angle

        if total_angle >= 2 * math.pi:
            break

    if step >= MAX_STEPS:
        print(f"  TIMEOUT after {MAX_STEPS} steps ({sim.time/86400:.1f} days)")
        return None

    period_s = sim.time
    period_days = period_s / 86400.0
    print(f"  Measured: {period_s:.1f}s = {period_days:.4f} days")
    print(f"  Analytical: {analytical_days:.4f} days")
    print(f"  Error: {abs(period_days - analytical_days):.4f} days ({abs(period_days - analytical_days)/analytical_days*100:.4f}%)")
    print(f"  e after: {body.orbit.e:.15f}")

    return period_days


def main():
    print("=== Earth Period ===")
    sim = Simulator("tests/test_orbital_period.toml", dt=DT)
    earth_a = 1.496e11
    earth_mu = G * 1.989e30  # Sun mass
    earth_analytical = 2.0 * math.pi * math.sqrt(earth_a**3 / earth_mu) / 86400.0
    measure_period(sim, "Earth", 1.989e30, earth_analytical)

    print("\n=== Mars Period ===")
    sim = Simulator("tests/test_orbital_period.toml", dt=DT)
    mars_a = 2.244e11
    mars_mu = G * 1.989e30  # Sun mass
    mars_analytical = 2.0 * math.pi * math.sqrt(mars_a**3 / mars_mu) / 86400.0
    measure_period(sim, "Mars", 1.989e30, mars_analytical)

    print("\n=== Direction Test (1 day) ===")
    sim = Simulator("tests/test_orbital_period.toml", dt=DT)
    earth = sim.get_body("Earth")
    sun = sim.get_body("Sun")
    theta_start = math.atan2(earth.global_pos[1] - sun.global_pos[1],
                             earth.global_pos[0] - sun.global_pos[0])

    sim.run(steps=1440)  # 1 day = 86400s / 60s

    theta_end = math.atan2(earth.global_pos[1] - sun.global_pos[1],
                           earth.global_pos[0] - sun.global_pos[0])
    delta = theta_end - theta_start
    print(f"  theta_start: {theta_start:.10f} rad")
    print(f"  theta_end:   {theta_end:.10f} rad")
    print(f"  delta:       {delta:.10f} rad")
    print(f"  prograde:    {delta > 0}")

    # Expected delta for 1 day of Earth orbit
    expected_delta = math.sqrt(earth_mu / earth_a**3) * 86400.0
    print(f"  expected:    {expected_delta:.10f} rad")
    print(f"  error:       {abs(delta - expected_delta):.10f} rad")


if __name__ == "__main__":
    main()