#!/usr/bin/env python3
"""
Full analytical propagation simulation of the Hohmann rendezvous scenario.
Replicates the exact physics from src/orbital_mechanics.cpp and src/maneuver.cpp.

Step-by-step trace to find where the 11,578 km separation comes from.

Usage: python3 tests/simulate_rendezvous.py
"""

import math
import sys

G = 6.67430e-11
MU = G * 5.972e24  # Earth

# ---- Vector operations ----
def vadd(a, b): return (a[0]+b[0], a[1]+b[1], a[2]+b[2])
def vsub(a, b): return (a[0]-b[0], a[1]-b[1], a[2]-b[2])
def vscale(v, s): return (v[0]*s, v[1]*s, v[2]*s)
def vmag(v): return math.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
def vdot(a, b): return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]
def vcross(a, b): return (
    a[1]*b[2] - a[2]*b[1],
    a[2]*b[0] - a[0]*b[2],
    a[0]*b[1] - a[1]*b[0]
)
def vnorm(v):
    m = vmag(v)
    if m < 1e-15: return (0, 0, 0)
    return (v[0]/m, v[1]/m, v[2]/m)

def normalize_angle(angle):
    while angle < 0.0: angle += 2*math.pi
    while angle >= 2*math.pi: angle -= 2*math.pi
    return angle

def normalize_angle_2pi(angle):
    while angle < 0.0: angle += 2*math.pi
    while angle >= 2*math.pi: angle -= 2*math.pi
    return angle

def normalize_angle_pi(angle):
    angle = normalize_angle_2pi(angle)
    while angle > math.pi: angle -= 2*math.pi
    while angle < -math.pi: angle += 2*math.pi
    return angle

# ---- Kepler equation solvers (exact C++ logic) ----
def get_initial_trial_value(mean_anomaly, eccentricity):
    return (mean_anomaly + eccentricity * math.sin(mean_anomaly)
            + ((eccentricity**2 / 2.0) * math.sin(2.0 * mean_anomaly)))

def solve_kepler_elliptical(mean_anomaly, eccentricity):
    E = get_initial_trial_value(mean_anomaly, eccentricity)
    E_prev = E + 2.0e-10
    for _ in range(50):
        if abs(E - E_prev) < 1e-10:
            break
        E_prev = E
        sin_E = math.sin(E)
        E = E - (E - eccentricity * sin_E - mean_anomaly) / (1.0 - eccentricity * math.cos(E))
    return E

def eccentric_to_true_anomaly(eccentric_anomaly, eccentricity):
    if abs(1.0 - eccentricity) < 0.01:
        E = eccentric_anomaly
        e = eccentricity
        cos_E = math.cos(E)
        sin_E = math.sin(E)
        denom = 1.0 - e * cos_E
        cos_nu = max(-1.0, min(1.0, (cos_E - e) / denom))
        sin_nu = max(-1.0, min(1.0, sin_E * math.sqrt(1.0 - e*e) / denom))
        return math.atan2(sin_nu, cos_nu)
    tan_half_E = math.tan(eccentric_anomaly / 2.0)
    tan_half_nu = math.sqrt((1.0 + eccentricity) / (1.0 - eccentricity)) * tan_half_E
    return 2.0 * math.atan(tan_half_nu)

# ---- Propagation (exact C++ propagate_orbital_elements) ----
def propagate(elements, dt, parent_mass):
    a = elements['a']
    e = elements['e']
    nu = elements['nu']
    mu = MU  # fixed for this sim

    if e < 1.0:
        n = math.sqrt(mu / a**3)
        E = 2.0 * math.atan(math.sqrt((1.0 - e) / (1.0 + e)) * math.tan(nu / 2.0))
        M = E - e * math.sin(E)
        M = M + n * dt
        E_new = get_initial_trial_value(M, e)
        E_prev = E_new + 2.0e-10
        for _ in range(50):
            if abs(E_new - E_prev) < 1e-10:
                break
            E_prev = E_new
            sin_E = math.sin(E_new)
            E_new = E_new - (E_new - e * sin_E - M) / (1.0 - e * math.cos(E_new))
        nu_new = 2.0 * math.atan(math.sqrt((1.0 + e) / (1.0 - e)) * math.tan(E_new / 2.0))
        result = dict(elements)
        result['nu'] = nu_new
        return result
    else:
        # Hyperbolic (not needed for this test)
        raise NotImplementedError("hyperbolic propagation not needed")

# ---- Cartesian from orbital elements ----
def orbital_to_cartesian(elements, parent_mass):
    a = elements['a']
    e = elements['e']
    nu = elements['nu']
    inc = elements['inc']
    Omega = elements['Omega']
    omega = elements['omega']
    mu = MU

    p = a * (1.0 - e*e)
    r = p / (1.0 + e * math.cos(nu))

    # Orbital plane position/velocity
    x_orb = r * math.cos(nu)
    y_orb = r * math.sin(nu)

    vx_orb = -math.sqrt(mu / p) * math.sin(nu)
    vy_orb = math.sqrt(mu / p) * (e + math.cos(nu))

    # z-x-z rotation: Rz(Omega) * Rx(inc) * Rz(omega)
    # Apply Rz(omega) first
    cos_w = math.cos(omega)
    sin_w = math.sin(omega)
    x1 = x_orb * cos_w - y_orb * sin_w
    y1 = x_orb * sin_w + y_orb * cos_w

    # Then Rx(inc)
    cos_i = math.cos(inc)
    sin_i = math.sin(inc)
    x2 = x1
    y2 = y1 * cos_i
    z2 = y1 * sin_i

    # Then Rz(Omega)
    cos_O = math.cos(Omega)
    sin_O = math.sin(Omega)
    pos = (x2 * cos_O - y2 * sin_O,
           x2 * sin_O + y2 * cos_O,
           z2)

    # Same rotation for velocity
    vx1 = vx_orb * cos_w - vy_orb * sin_w
    vy1 = vx_orb * sin_w + vy_orb * cos_w
    vx2 = vx1
    vy2 = vy1 * cos_i
    vz2 = vy1 * sin_i
    vel = (vx2 * cos_O - vy2 * sin_O,
           vx2 * sin_O + vy2 * cos_O,
           vz2)

    return pos, vel

# ---- Cartesian to orbital elements ----
def cartesian_to_elements(pos, vel, parent_mass):
    mu = MU
    r = vmag(pos)
    v = vmag(vel)

    # Specific orbital energy
    specific_energy = -mu / r + v**2 / 2.0

    # Semi-major axis
    if abs(specific_energy) < 1e-10:
        a = 1e10
    else:
        a = -mu / (2.0 * specific_energy)

    # Angular momentum
    h_vec = vcross(pos, vel)
    h = vmag(h_vec)

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

    # True anomaly
    if e < 1e-10:
        nu = 0.0
    else:
        cos_nu = vdot(pos, e_vec) / (r * e)
        cos_nu = max(-1.0, min(1.0, cos_nu))
        if abs(cos_nu) > 1.0 - 1e-10:
            h_cross_e = vcross(h_vec, e_vec)
            denom = r * e * h
            sin_nu = vdot(pos, h_cross_e) / denom if denom > 1e-10 else 0.0
        else:
            r_cross_h = vcross(pos, h_vec)
            denom = r * e * h
            sin_nu = vdot(r_cross_h, e_vec) / denom if denom > 1e-10 else 0.0
        nu = math.atan2(sin_nu, cos_nu)
        if nu == -math.pi:
            nu = math.pi
        nu = normalize_angle(nu)

    # Inclination
    if h > 1e-10:
        i = math.acos(h_vec[2] / h)
    else:
        i = 0.0

    # RAAN
    n_vec = (0, 0, 1)
    n = vcross(n_vec, h_vec)
    n_mag = vmag(n)
    if n_mag > 1e-10:
        Omega = math.acos(n[0] / n_mag)
        if n[1] < 0.0:
            Omega = 2*math.pi - Omega
    else:
        Omega = 0.0

    # Argument of periapsis
    if e > 1e-10 and n_mag > 1e-10 and i > 0.01:
        cos_omega = vdot(e_vec, n) / (e * n_mag)
        n_cross_e = vcross(n, e_vec)
        sin_omega = vdot(n_cross_e, h_vec) / (e * n_mag * h)
        omega = math.atan2(sin_omega, cos_omega)
        if omega < 0: omega += 2*math.pi
    elif e > 1e-10:
        omega = math.atan2(e_vec[1], e_vec[0])
        if omega < 0: omega += 2*math.pi
    else:
        omega = 0.0

    return {'a': a, 'e': e, 'nu': nu, 'inc': i, 'Omega': Omega, 'omega': omega}

# ---- Hohmann transfer calculations ----
def hohmann_transfer_time(r1, r2):
    a_t = (r1 + r2) / 2.0
    T = 2*math.pi * math.sqrt(a_t**3 / MU)
    return T / 2.0

def required_separation(r1, r2):
    tt = hohmann_transfer_time(r1, r2)
    n2 = math.sqrt(MU / r2**3)
    target_angle = n2 * tt
    return target_angle - math.pi

def calc_mean_motion(radius):
    return math.sqrt(MU / radius**3)

def calculate_wait_time_for_hohmann(r1, r2, angular_separation):
    required_sep = required_separation(r1, r2)
    n1 = calc_mean_motion(r1)
    n2 = calc_mean_motion(r2)
    rel_angular_vel = n1 - n2

    current_sep = normalize_angle_pi(angular_separation)
    required_sep = normalize_angle_pi(required_sep)

    angle_to_close = required_sep - current_sep
    return angle_to_close / rel_angular_vel

def relative_orbit_period(r1, r2):
    n1 = calc_mean_motion(r1)
    n2 = calc_mean_motion(r2)
    return 2*math.pi / abs(n1 - n2)

def calculate_next_hohmann_wait_time(r1, r2, angular_sep, dt):
    wait_time = calculate_wait_time_for_hohmann(r1, r2, angular_sep)
    rel_period = relative_orbit_period(r1, r2)
    while wait_time < dt:
        wait_time += rel_period
    return wait_time

# ---- Burn application ----
def apply_burn(pos, vel, direction, delta_v, parent_mass):
    """Apply impulsive burn in local orbital frame."""
    # direction: 'prograde', 'retrograde', 'normal'
    if direction == 'prograde':
        d = vnorm(vel)
    elif direction == 'retrograde':
        d = vscale(vnorm(vel), -1)
    elif direction == 'normal':
        h = vcross(pos, vel)
        d = vnorm(h)
    else:
        raise ValueError(f"Unknown direction: {direction}")

    new_vel = vadd(vel, vscale(d, delta_v))
    return pos, new_vel

# ---- Dump state helper ----
def dump_state(label, chaser, target, chaser_pos, chaser_vel, target_pos, target_vel, sim_time):
    """Print state at key simulation milestones, matching test_rendezvous.cpp dump_state."""
    c_r = vmag(chaser_pos)
    t_r = vmag(target_pos)
    c_sep = vmag(vsub(chaser_pos, target_pos))
    print(f"\n*** {label} (t={sim_time:.1f}s) ***")
    print(f"  Chaser: r={c_r:.0f} m, nu={chaser['nu']:.6f} rad ({math.degrees(chaser['nu']):.1f}°)  "
          f"a={chaser['a']:.0f} e={chaser['e']:.6f}")
    print(f"    pos={chaser_pos}, vel={chaser_vel}")
    print(f"  Target: r={t_r:.0f} m, nu={target['nu']:.6f} rad ({math.degrees(target['nu']):.1f}°)  "
          f"a={target['a']:.0f} e={target['e']:.6f}")
    print(f"    pos={target_pos}, vel={target_vel}")
    print(f"  Separation: {c_sep:.0f} m")


# ---- Full rendezvous scenario ----
def main():
    # Initial conditions from test_rendezvous.toml
    TARGET_R = 6.771e6
    TARGET_NU = 0.0
    CHASER_R = 6.671e6
    CHASER_NU = 4.71238898038469  # 270 degrees

    print("=== INITIAL STATE ===")
    print(f"Chaser: r={CHASER_R:.1f} m, nu={math.degrees(CHASER_NU):.1f} deg")
    print(f"Target: r={TARGET_R:.1f} m, nu={math.degrees(TARGET_NU):.1f} deg")

    # Create orbital elements (coplanar, circular)
    chaser = {'a': CHASER_R, 'e': 0.0, 'nu': CHASER_NU,
              'inc': 0.0, 'Omega': 0.0, 'omega': 0.0}
    target = {'a': TARGET_R, 'e': 0.0, 'nu': TARGET_NU,
              'inc': 0.0, 'Omega': 0.0, 'omega': 0.0}

    chaser_pos, chaser_vel = orbital_to_cartesian(chaser, 5.972e24)
    target_pos, target_vel = orbital_to_cartesian(target, 5.972e24)

    print(f"Chaser pos: {chaser_pos}, vel: {vmag(chaser_vel):.1f} m/s")
    print(f"Target pos: {target_pos}, vel: {vmag(target_vel):.1f} m/s")

    # Angular separation
    angular_sep = chaser['nu'] - target['nu']
    angular_sep = normalize_angle_pi(angular_sep)
    print(f"\nAngular separation (chaser - target): {math.degrees(angular_sep):.1f} deg")

    # Hohmann parameters
    hohmann_tt = hohmann_transfer_time(CHASER_R, TARGET_R)
    dv1 = math.sqrt(MU * (2/CHASER_R - 2/(CHASER_R + TARGET_R))) - math.sqrt(MU/CHASER_R)
    dv2 = math.sqrt(MU/TARGET_R) - math.sqrt(MU * (2/TARGET_R - 2/(CHASER_R + TARGET_R)))
    print(f"\nHohmann transfer: tt={hohmann_tt:.1f} s, dv1={dv1:.2f} m/s, dv2={dv2:.2f} m/s")

    # Phasing
    dt = 0.1
    wait_time = calculate_next_hohmann_wait_time(CHASER_R, TARGET_R, angular_sep, dt)
    arrival_time = wait_time + hohmann_tt
    print(f"Wait time: {wait_time:.2f} s")
    print(f"Arrival time: {arrival_time:.2f} s")
    print(f"Steps: {int(arrival_time/dt)}")

    # ---- Run simulation ----
    print(f"\n=== SIMULATION (dt={dt}) ===")
    sim_time = 0.0
    steps = 0
    chaser_executed = False
    arrival_executed = False

    while steps < int(arrival_time / dt) + 1000:
        chaser, target, chaser_pos, chaser_vel, target_pos, target_vel = \
            update_simulation(chaser, target, sim_time, dt, dv1, dv2, wait_time, arrival_time,
                              chaser_pos, chaser_vel, target_pos, target_vel,
                              chaser_executed, arrival_executed)
        sim_time += dt
        steps += 1

        if steps == 1:
            dump_state("T=0 (initial)", chaser, target, chaser_pos, chaser_vel, target_pos, target_vel, sim_time)
        if steps == int(wait_time / dt):
            dump_state("JUST BEFORE DEPARTURE", chaser, target, chaser_pos, chaser_vel, target_pos, target_vel, sim_time)
        if steps == int(wait_time / dt) + 1:
            dump_state("AFTER DEPARTURE BURN", chaser, target, chaser_pos, chaser_vel, target_pos, target_vel, sim_time)
        if steps == int(arrival_time / dt) - 1:
            dump_state("JUST BEFORE ARRIVAL BURN", chaser, target, chaser_pos, chaser_vel, target_pos, target_vel, sim_time)

        if not chaser_executed and sim_time >= wait_time:
            # Execute departure burn
            print(f"\n  *** DEPARTURE BURN at t={sim_time:.1f}s ***")
            print(f"    Before: pos={chaser_pos}, vel={chaser_vel}")
            chaser_pos, chaser_vel = apply_burn(chaser_pos, chaser_vel, 'prograde', dv1, 5.972e24)
            print(f"    After:  pos={chaser_pos}, vel={chaser_vel}")
            chaser = cartesian_to_elements(chaser_pos, chaser_vel, 5.972e24)
            chaser_executed = True
            print(f"    Chaser: r={vmag(chaser_pos):.0f} nu={math.degrees(chaser['nu']):.1f}°  "
                  f"a={chaser['a']:.0f} e={chaser['e']:.6f}")

        if not arrival_executed and sim_time >= arrival_time:
            # Execute arrival burn
            chaser_pos, chaser_vel = apply_burn(chaser_pos, chaser_vel, 'prograde', dv2, 5.972e24)
            chaser = cartesian_to_elements(chaser_pos, chaser_vel, 5.972e24)
            arrival_executed = True
            print(f"\n  *** ARRIVAL BURN at t={sim_time:.1f}s ***")
            print(f"    Chaser: r={vmag(chaser_pos):.0f} nu={math.degrees(chaser['nu']):.1f}°  "
                  f"a={chaser['a']:.0f} e={chaser['e']:.6f}")

            dump_state("AFTER ARRIVAL BURN", chaser, target, chaser_pos, chaser_vel, target_pos, target_vel, sim_time)

            # Final comparison
            c_sep = vmag(vsub(chaser_pos, target_pos))
            c_r = vmag(chaser_pos)
            t_r = vmag(target_pos)
            c_vel = vmag(chaser_vel)
            t_vel = vmag(target_vel)
            print(f"\n=== FINAL STATE ===")
            print(f"Chaser: r={c_r:.0f} m, nu={chaser['nu']:.6f} rad ({math.degrees(chaser['nu']):.1f}°)")
            print(f"  pos={chaser_pos}, vel={chaser_vel}")
            print(f"Target: r={t_r:.0f} m, nu={target['nu']:.6f} rad ({math.degrees(target['nu']):.1f}°)")
            print(f"  pos={target_pos}, vel={target_vel}")
            print(f"Separation: {c_sep:.0f} m")
            print(f"Speed: chaser={c_vel:.2f} target={t_vel:.2f} m/s")
            print(f"Radius error: {abs(c_r - t_r):.6f} m")
            print(f"Chaser eccentricity: {chaser['e']:.15f}")
            print(f"Target eccentricity: {target['e']:.15f}")
            break


def update_simulation(chaser, target, sim_time, dt, dv1, dv2, wait_time, arrival_time,
                      chaser_pos, chaser_vel, target_pos, target_vel,
                      chaser_executed, arrival_executed):
    """Propagate one timestep for both spacecraft."""
    chaser = propagate(chaser, dt, 5.972e24)
    target = propagate(target, dt, 5.972e24)

    chaser_pos, chaser_vel = orbital_to_cartesian(chaser, 5.972e24)
    target_pos, target_vel = orbital_to_cartesian(target, 5.972e24)

    return chaser, target, chaser_pos, chaser_vel, target_pos, target_vel


if __name__ == '__main__':
    main()