"""Checks every identity and number of complications-to-impulsive-transfers-chobotov.qmd.

Part 1: the algebra (sympy, and explicit vectors): the velocity triangle with a plane change, the composition of two plane changes, the angle between
        two orbit planes, Edelbaum's pseudo-velocity argument, the gravity-loss equation.
Part 2: fixed-impulse transfers: the solid-motor example, the non-Hohmann family, the region of parking orbits, the launch window, problems 6.1 to 6.3.
Part 3: very low thrust (Edelbaum), the LEO-to-GEO numbers.
Part 4: finite burns by numerical integration: the gravity-loss integral, the effect of thrust level and of splitting the burn.

Run from the sfm-slides directory:   python3 codes/complications-verify.py
"""
import sys
import json
import numpy as np
import sympy as sp
from scipy.integrate import solve_ivp
from scipy.optimize import minimize_scalar
sys.path.insert(0, "codes")
from _complic import *

RESULTS = []


def check(name, value, stated, tol):
    ok = abs(value - stated) <= tol
    RESULTS.append(ok)
    print(f"{'PASS' if ok else 'FAIL'}  {name:<80} {value:16.8g}   note: {stated}")


def truth(name, cond):
    RESULTS.append(bool(cond))
    print(f"{'PASS' if cond else 'FAIL'}  {name}")


# ------------------------------------------------------------------ Part 1
print("Part 1: the algebra")
Vc, Vt, al, ga = sp.symbols("V_c V_t alpha gamma", positive=True)
vc = sp.Matrix([Vc, 0, 0])
vt = sp.Matrix([Vt * sp.cos(ga) * sp.cos(al), Vt * sp.cos(ga) * sp.sin(al), Vt * sp.sin(ga)])
dv2 = sp.simplify((vt - vc).dot(vt - vc))
truth("|v_t - v_c|^2 = Vc^2 + Vt^2 - 2 Vc Vt cos(gamma) cos(alpha)", sp.simplify(dv2 - (Vc ** 2 + Vt ** 2 - 2 * Vc * Vt * sp.cos(ga) * sp.cos(al))) == 0)

rng = np.random.default_rng(11)


def rot(v, axis, ang):
    axis = axis / np.linalg.norm(axis)
    return v * np.cos(ang) + np.cross(axis, v) * np.sin(ang) + axis * (axis @ v) * (1 - np.cos(ang))


worst = 0.0
for _ in range(300):
    nt = rng.normal(size=3); nt /= np.linalg.norm(nt)
    e1 = np.cross(nt, rng.normal(size=3)); e1 /= np.linalg.norm(e1)          # radius vector at the first burn (in the transfer plane)
    f2 = rng.uniform(0.2, np.pi)
    e2 = rot(e1, nt, f2)                                                        # radius vector at the second burn
    A1, A2 = rng.uniform(0.0, 0.5), rng.uniform(0.0, 0.5)
    for s in (1.0, -1.0):
        n2 = rot(nt, e2, s * A2)           # final-orbit normal: the transfer normal turned about the second burn's radius vector
        n1 = rot(nt, e1, A1)               # parking-orbit normal: turned about the first burn's radius vector
        th = np.arccos(np.clip(n1 @ n2, -1, 1))
        # the two senses: s = +1 here corresponds to the two turns made in the sense that makes cos f2 enter with + or -; test both signs
        c_pm = np.cos(A1) * np.cos(A2) + np.sin(A1) * np.sin(A2) * np.cos(f2)
        c_mm = np.cos(A1) * np.cos(A2) - np.sin(A1) * np.sin(A2) * np.cos(f2)
        worst = max(worst, min(abs(np.cos(th) - c_pm), abs(np.cos(th) - c_mm)))
check("cos(theta) = cos a1 cos a2 +- sin a1 sin a2 cos f2: explicit vectors, 600 cases, worst error", worst, 0.0, 1e-12)
truth("Hohmann (f2 = 180 deg): the formula gives theta = a1 + a2 and |a1 - a2|",
      abs(total_plane_angle(0.05, 0.4, np.pi, True) - 0.45) < 1e-12 and abs(total_plane_angle(0.05, 0.4, np.pi, False) - 0.35) < 1e-12)
# plane tilt equals the rotation of the velocity direction about the radius vector
r_hat = np.array([0.0, 0.0, 1.0]); v_dir = np.array([1.0, 0.0, 0.0]); a_ = 0.3
v_new = rot(v_dir, r_hat, a_)
n_old = np.cross(r_hat, v_dir); n_new = np.cross(r_hat, v_new)
check("rotating the velocity about r by alpha tilts the orbit plane by alpha", np.arccos(n_old @ n_new), a_, 1e-12)
# orbit-plane normals and cos(theta) = cos i1 cos i2 + sin i1 sin i2 cos dOmega
i1_, i2_, O1, O2 = sp.symbols("i1 i2 O1 O2")
n1v = sp.Matrix([sp.sin(i1_) * sp.sin(O1), -sp.sin(i1_) * sp.cos(O1), sp.cos(i1_)])
n2v = sp.Matrix([sp.sin(i2_) * sp.sin(O2), -sp.sin(i2_) * sp.cos(O2), sp.cos(i2_)])
truth("n1 . n2 = cos i1 cos i2 + sin i1 sin i2 cos(O1 - O2)",
      sp.simplify(n1v.dot(n2v) - (sp.cos(i1_) * sp.cos(i2_) + sp.sin(i1_) * sp.sin(i2_) * sp.cos(O1 - O2))) == 0)
# Edelbaum: the pseudo-velocity u = V (cos phi, sin phi) with phi = (pi/2) di moves at the constant rate a
t, a_T = sp.symbols("t a_T", positive=True)
Vf, phif, bf = sp.Function("V")(t), sp.Function("phi")(t), sp.Function("beta")(t)
uvec = sp.Matrix([Vf * sp.cos(phif), Vf * sp.sin(phif)])
duvec = uvec.diff(t).subs({Vf.diff(t): -a_T * sp.cos(bf), phif.diff(t): a_T * sp.sin(bf) / Vf})
truth("Edelbaum: |d u/dt|^2 = a^2 for any steering angle beta", sp.simplify(duvec.dot(duvec) - a_T ** 2) == 0)
truth("Edelbaum: d u/dt = -a (cos(beta - phi), -sin(beta - phi))",
      sp.simplify(duvec - (-a_T) * sp.Matrix([sp.cos(bf - phif), -sp.sin(bf - phif)])) == sp.zeros(2, 1))
# integrate the averaged equations with the steering beta = phi + beta_c and check the end conditions
V0, V1, di = 7.7385, 3.8736, 28.5 * D2R
phi1 = 0.5 * np.pi * di
aT = 1e-6
bc = np.arctan2(V1 * np.sin(phi1), V0 - V1 * np.cos(phi1))
tf = edelbaum_dv(V0, V1, di) / aT
f = lambda tt, y: [-aT * np.cos(y[1] + bc), aT * np.sin(y[1] + bc) / y[0]]
sol = solve_ivp(f, [0, tf], [V0, 0.0], rtol=1e-12, atol=1e-14)
check("Edelbaum: final speed after the optimal-steering integration [km/s]", sol.y[0, -1], V1, 1e-8)
check("Edelbaum: final (pi/2) di after the integration [rad]", sol.y[1, -1], phi1, 1e-9)

# gravity loss in a vertical ascent
u_ = 300 * 9.80665
for tw, Rm in ((1.5, 3.0), (1.2, 2.0), (3.0, 2.0)):
    g = 9.80665; tb = (1 - 1 / Rm) * u_ / (tw * g)
    s = solve_ivp(lambda tt, y: [tw * g / y[1] - g, -tw * g / u_], [0, tb], [0.0, 1.0], rtol=1e-12, atol=1e-12)
    check(f"vertical ascent T/W0 = {tw}, m0/mf = {Rm}: loss u ln R - v equals g t_b [m/s]", u_ * np.log(Rm) - s.y[0, -1], g * tb, 1e-3)
check("  loss for T/W0 = 1.5, R = 3 [m/s]", g * (1 - 1 / 3.0) * u_ / (1.5 * g), 1307.55, 0.005)
check("  loss for T/W0 = 3.0, R = 2 [m/s]", (1 - 1 / 2.0) * u_ / 3.0, 490.33, 0.005)

# ------------------------------------------------------------------ Part 2
print("\nPart 2: fixed-impulse transfers")
dv1 = rocket_dv(907 + 91 + 898 + 181 + 2174, 907 + 91 + 898 + 181, 300)
dv2 = rocket_dv(907 + 91 + 898, 907 + 91, 300)
check("rocket equation, second burn dV2 [m/s]", dv2 * 1000, 1888.0, 0.05); check("  first burn dV1 [m/s]", dv1 * 1000, 2107.1, 0.05)
check("  mass ratio of the second burn Wi/Wf", (907 + 91 + 898) / (907 + 91), 1.8998, 5e-5)
check("  first burn: Wi/Wf", (907 + 91 + 898 + 181 + 2174) / (907 + 91 + 898 + 181), 4251 / 2077, 1e-12)
r1, r2 = RE + 278, RE + 20187
h = hohmann(r1, r2)
check("r1 [km]", r1, 6656.137, 5e-4); check("r2 [km]", r2, 26565.137, 5e-4); check("Vc1 [km/s]", h["vc1"], 7.7385, 5e-5); check("Vc2 [km/s]", h["vc2"], 3.8736, 5e-5)
check("a_tr [km]", h["a"], 16610.6, 0.05); check("V_PTR [km/s]", h["vp"], 9.7863, 5e-5); check("V_ATR [km/s]", h["va"], 2.4521, 5e-5)
check("pure Hohmann: dV1 [m/s]", h["dv1"] * 1000, 2047.8, 0.05); check("  dV2 [m/s]", h["dv2"] * 1000, 1421.5, 0.05); check("  total [m/s]", (h["dv1"] + h["dv2"]) * 1000, 3469.4, 0.05)
check("  the excess the motors carry [m/s]", (dv1 + dv2 - h["dv1"] - h["dv2"]) * 1000, 525.8, 0.05)
a1 = plane_angle(h["vc1"], h["vp"], dv1); a2 = plane_angle(h["vc2"], h["va"], dv2)
check("alpha1 [deg]", a1 / D2R, 3.269, 5e-4); check("alpha2 [deg]", a2 / D2R, 23.259, 5e-4)
check("theta_a = alpha1 + alpha2 [deg]", total_plane_angle(a1, a2, np.pi, True) / D2R, 26.528, 5e-4)
check("theta_b = |alpha1 - alpha2| [deg]", total_plane_angle(a1, a2, np.pi, False) / D2R, 19.990, 5e-4)
check("check: dV1 from alpha1", dv_from_angle(h["vc1"], h["vp"], a1), dv1, 1e-12); check("check: dV2 from alpha2", dv_from_angle(h["vc2"], h["va"], a2), dv2, 1e-12)
i2 = 55 * D2R
for lab, th, want in (("i1 = i2 - theta_a", i2 - 26.528141 * D2R, 28.47), ("i1 = i2 + theta_a", i2 + 26.528141 * D2R, 81.53),
                      ("i1 = i2 - theta_b", i2 - 19.990303 * D2R, 35.01), ("i1 = i2 + theta_b", i2 + 19.990303 * D2R, 74.99)):
    check(f"  {lab} [deg]", th / D2R, want, 0.005)
f_split = lambda x: dv_from_angle(h["vc1"], h["vp"], x) + dv_from_angle(h["vc2"], h["va"], (a1 + a2) - x)
rs = minimize_scalar(f_split, bounds=(0, a1 + a2), method="bounded")
check("least total dV for the plane change theta_a: alpha1 [deg]", rs.x / D2R, 3.170, 0.005); check("  that least total [m/s]", rs.fun * 1000, 3995.11, 0.005)
check("  the motors' total [m/s]", (dv1 + dv2) * 1000, 3995.16, 0.005)
allapo = dv_from_angle(h["vc1"], h["vp"], 0.0) + dv_from_angle(h["vc2"], h["va"], a1 + a2)
check("  the same plane change done all at the apogee burn [m/s]", allapo * 1000, 4053.0, 0.05)

# non-Hohmann
vp_max = h["vc1"] + dv1
a_max = 1 / (2 / r1 - vp_max ** 2 / MU); ra_max = 2 * a_max - r1
check("largest apogee radius (alpha1 -> 0) [km]", ra_max, 28259.4, 0.05); check("  above the Hohmann apogee [km]", ra_max - r2, 1694.3, 0.05)
ras = np.linspace(r2 + 1e-3, ra_max - 1e-6, 6000); tha, thb = [], []
for ra in ras:
    vp, v2, g2, f2, a = transfer_state(r1, r2, ra)
    A1 = plane_angle(h["vc1"], vp, dv1); A2 = plane_angle(h["vc2"], v2, dv2, g2)
    tha.append(total_plane_angle(A1, A2, f2, True) / D2R); thb.append(total_plane_angle(A1, A2, f2, False) / D2R)
tha, thb = np.array(tha), np.array(thb)
check("family: theta_a falls from [deg]", tha[0], 26.528, 5e-3); check("  to [deg]", tha[-1], 13.063, 5e-3)
check("family: theta_b smallest value [deg]", thb.min(), 12.885, 5e-3); check("  reached at apogee radius above the Hohmann one [km]", ras[np.argmin(thb)] - r2, 1671, 3.0)
vp, v2, g2, f2, a = transfer_state(r1, r2, ra_max - 1e-6)
check("at the limit: flight-path angle change gamma2 [deg]", g2 / D2R, 22.95, 0.005); check("  true anomaly of the second burn [deg]", f2 / D2R, 163.88, 0.005)
check("  alpha2 [deg]", plane_angle(h["vc2"], v2, dv2, g2) / D2R, 13.063, 5e-3)
vp, v2, g2, f2, a = transfer_state(r1, r2, r2 + 100.0)
check("a 100 km higher apogee: alpha1 [deg]", plane_angle(h["vc1"], vp, dv1) / D2R, 3.167, 5e-3); check("  alpha2 [deg]", plane_angle(h["vc2"], v2, dv2, g2) / D2R, 22.754, 5e-3)
check("  gamma2 [deg]", g2 / D2R, 6.046, 5e-3)
th_out, th_in = tha[0], thb.min()
print("Launch window for the 55 deg final orbit (outer theta %.3f, inner theta %.3f)" % (th_out, th_in))
for i1d, key in ((55.0, "55"), (35.0, "35"), (28.5, "28.5")):
    o = np.degrees(node_change(i1d * D2R, i2, th_out * D2R)); n = node_change(i1d * D2R, i2, th_in * D2R)
    print("   i1 =", i1d, "outer dOmega", o, "inner", np.degrees(n) if not np.isnan(n) else None)
o55 = np.degrees(node_change(55 * D2R, i2, th_out * D2R)); n55 = np.degrees(node_change(55 * D2R, i2, th_in * D2R))
check("i1 = 55: outer boundary dOmega [deg]", o55, 32.53, 0.005); check("  inner boundary [deg]", n55, 15.75, 0.005)
check("  times: outer [h]", o55 / OMEGA_E, 2.163, 0.0005); check("  inner [h]", n55 / OMEGA_E, 1.047, 0.0005)
o35 = np.degrees(node_change(35 * D2R, i2, th_out * D2R)); check("i1 = 35: outer dOmega [deg]", o35, 25.27, 0.005); check("  time [h]", o35 / OMEGA_E, 1.680, 0.0005)
check("  plane angle at dOmega = 0 [deg] (inside the region: no gap)", abs(55 - 35), 20.0, 1e-9)
o285 = np.degrees(node_change(28.5 * D2R, i2, th_out * D2R)); check("i1 = 28.5: outer dOmega [deg]", o285, 1.92, 0.005); check("  time [min]", o285 / OMEGA_E * 60, 7.66, 0.01)
check("the Earth turns 15.041 deg per hour: 15 deg per hour times 23.9345 h", 15.0410686 * 23.9344696, 360.0, 1e-3)

print("Problem 6.1 (payload 990 kg)")
d1b = rocket_dv(990 + 91 + 898 + 181 + 2174, 990 + 91 + 898 + 181, 300); d2b = rocket_dv(990 + 91 + 898, 990 + 91, 300)
check("  dV1 [m/s]", d1b * 1000, 2048.8, 0.05); check("  dV2 [m/s]", d2b * 1000, 1779.0, 0.05); check("  tangential Hohmann dV1 [m/s]", h["dv1"] * 1000, 2047.8, 0.05)
b1 = plane_angle(h["vc1"], h["vp"], d1b); b2 = plane_angle(h["vc2"], h["va"], d2b)
check("  alpha1 [deg]", b1 / D2R, 0.405, 5e-4); check("  alpha2 [deg]", b2 / D2R, 19.988, 5e-4)
th1, th2 = (b1 + b2) / D2R, abs(b1 - b2) / D2R
check("  theta_a [deg]", th1, 20.393, 5e-3); check("  theta_b [deg]", th2, 19.583, 5e-3)
vpm = h["vc1"] + d1b; am = 1 / (2 / r1 - vpm ** 2 / MU)
check("  the apogee may rise by at most [km]", 2 * am - r1 - r2, 25.05, 0.05)
c1 = np.degrees(node_change(55 * D2R, i2, th1 * D2R)); c2 = np.degrees(node_change(55 * D2R, i2, th2 * D2R))
check("  i1 = 55: dOmega for theta_b [deg]", c2, 23.96, 0.005); check("  dOmega for theta_a [deg]", c1, 24.96, 0.005)
check("  the two instants are +-[h] from the conodal time", c2 / OMEGA_E, 1.593, 5e-4); check("  and +-[h]", c1 / OMEGA_E, 1.660, 5e-4)

print("Problem 6.2 (ion engine)")
Vc1_, Vc2_ = h["vc1"], h["vc2"]
dV62 = edelbaum_dv(Vc1_, Vc2_, 0.0)
check("  dV = V0 - V1 [m/s]", dV62 * 1000, 3864.9, 0.05); check("  the Hohmann total [m/s]", (h["dv1"] + h["dv2"]) * 1000, 3469.4, 0.05)
check("  loss [m/s]", (dV62 - h["dv1"] - h["dv2"]) * 1000, 395.6, 0.05)
W0 = 1361 * np.exp(dV62 / (G0 * 3000)); check("  initial weight [kg]", W0, 1552.1, 0.05); check("  propellant [kg]", W0 - 1361, 191.1, 0.05)
mdot = 4.45 / (9.80665 * 3000)
check("  mass flow [kg/s]", mdot * 1e4, 1.5126, 5e-5); check("  transfer time [days]", (W0 - 1361) / mdot / 86400, 14.62, 0.005)

print("Problem 6.3 (two motors, 296 km to 741 km)")
rr1, rr2 = 6378.0 + 296, 6378.0 + 741
hc = hohmann(rr1, rr2)
d2c = rocket_dv(1452 + 136 + 1044, 1452 + 136, 285); d1c = rocket_dv(1452 + 136 + 1044 + 136 + 1044, 1452 + 136 + 1044 + 136, 285)
check("  dV1 [m/s]", d1c * 1000, 894.4, 0.05); check("  dV2 [m/s]", d2c * 1000, 1412.2, 0.05)
c1_ = plane_angle(hc["vc1"], hc["vp"], d1c); c2_ = plane_angle(hc["vc2"], hc["va"], d2c)
check("  alpha1 [deg]", c1_ / D2R, 6.519, 5e-4); check("  alpha2 [deg]", c2_ / D2R, 10.878, 5e-4)
check("  largest destination inclination [deg]", 57 + (c1_ + c2_) / D2R, 74.40, 0.005); check("  transfer-orbit inclination [deg]", 57 + c1_ / D2R, 63.52, 0.005)
check("  V_PTR [km/s] (the book gives 7851 m/s)", hc["vp"], 7.8518, 5e-4); check("  V_ATR [km/s] (7361 m/s)", hc["va"], 7.3610, 5e-4)

# ------------------------------------------------------------------ Part 3
print("\nPart 3: low thrust")
rl, rg = RE + 278, RE + 35786
hg = hohmann(rl, rg)
check("coplanar LEO to GEO, Hohmann [m/s]", (hg["dv1"] + hg["dv2"]) * 1000, 3900.0, 1.5)
print("      coplanar Hohmann LEO-GEO:", (hg["dv1"] + hg["dv2"]) * 1000)
th28 = 28.5 * D2R
check("LEO to GEO with 28.5 deg: all the plane change at the apogee burn [m/s]", (hg["dv1"] + dv_from_angle(hg["vc2"], hg["va"], th28)) * 1000, 4263.6, 0.05)
fsp = lambda x: dv_from_angle(hg["vc1"], hg["vp"], x) + dv_from_angle(hg["vc2"], hg["va"], th28 - x)
rsp = minimize_scalar(fsp, bounds=(0, th28), method="bounded")
check("  with the plane change split best: alpha1 [deg]", rsp.x / D2R, 2.193, 5e-4); check("  total [m/s]  (Chobotov: 4237)", rsp.fun * 1000, 4239.1, 0.05)
check("Edelbaum LEO to GEO, 28.5 deg [m/s]", edelbaum_dv(hg["vc1"], hg["vc2"], th28) * 1000, 5962.7, 0.05)
check("  coplanar spiral V0 - V1 [m/s]", (hg["vc1"] - hg["vc2"]) * 1000, 4663.9, 0.05)
check("  Edelbaum minus the best impulsive [m/s]", (edelbaum_dv(hg["vc1"], hg["vc2"], th28) - rsp.fun) * 1000, 1723.6, 0.1)
check("Edelbaum for equal circular speeds is 2 V sin(pi di/4) [m/s] (di = 10 deg, V = 7.7385)", edelbaum_dv(7.7385, 7.7385, 10 * D2R) * 1000, 2 * 7.7385 * np.sin(np.pi * 10 * D2R / 4) * 1000, 1e-6)
check("  against the impulsive plane change 2 V sin(di/2) [m/s]", 2 * 7.7385 * np.sin(10 * D2R / 2) * 1000, 1349.0, 0.5)
check("  low thrust for a pure plane change costs pi/2 times more, small angles", edelbaum_dv(7.7, 7.7, 0.01) / (2 * 7.7 * np.sin(0.005)), np.pi / 2, 1e-4)

# ------------------------------------------------------------------ Part 4
print("\nPart 4: finite burns")
u = G0 * 300.0
m0 = 1000.0; T = 0.1 * m0 * G0
st = np.array([rl, 0.0, 0.0, hg["vc1"]])
vt_ = vis_viva(rl, 0.5 * (rl + rg)); v1_ = hg["vc1"] + (vt_ - hg["vc1"]) / 2
ra1 = 2 / (2 / rl - v1_ ** 2 / MU) - rl
s1 = burn(st, m0, T, 300.0, ra1); y1 = s1.y_events[0][0]
s0 = np.array(y1[:4]); s0 = coast(s0, time_to_perigee(s0))
tb = burn(s0, y1[4], T, 300.0, rg).t_events[0][0]
back = solve_ivp(lambda tt, y: [y[2], y[3], -MU * y[0] / np.hypot(y[0], y[1]) ** 3, -MU * y[1] / np.hypot(y[0], y[1]) ** 3], [0, -tb / 2], list(s0), rtol=1e-11, atol=1e-9, max_step=30).y[:, -1]
s2 = burn(back, y1[4], T, 300.0, rg, with_integrals=True); yy = s2.y_events[0][0]
dur = s2.t_events[0][0]
ideal = u * np.log(y1[4] / yy[4]); vstart = np.hypot(*back[2:]); vend = np.hypot(yy[2], yy[3])
check("second perigee burn (T/W0 = 0.1, centred): duration [min]", dur / 60, 11.20, 0.005)
check("  ideal dV = u ln(Wi/Wf) [m/s]", ideal * 1000, 1231.9, 0.05); check("  gravity integral [m/s]", yy[5] * 1000, 152.6, 0.05)
check("  speed gain [m/s]", (vend - vstart) * 1000, 1079.3, 0.05)
check("  equation 6.5: speed gain minus (ideal - gravity integral) [m/s]", ((vend - vstart) - (ideal - yy[5])) * 1000, 0.0, 1e-6)
r_s = np.hypot(*back[:2]); r_e = np.hypot(*yy[:2])
gam_s = np.degrees(np.arcsin((back[0] * back[2] + back[1] * back[3]) / (r_s * vstart)))
gam_e = np.degrees(np.arcsin((yy[0] * yy[2] + yy[1] * yy[3]) / (r_e * vend)))
check("  flight-path angle at the start of the burn [deg]", gam_s, -6.17, 0.005); check("  at the end [deg]", gam_e, 11.13, 0.005)
eps_gain = (vend ** 2 / 2 - MU / r_e) - (vstart ** 2 / 2 - MU / r_s)
check("  energy gain [km^2/s^2]", eps_gain, 11.54, 0.005); check("  = integral of v dV", yy[6], eps_gain, 1e-8)
vp_before = np.hypot(*s0[2:])
check("  energy gain of an impulsive kick of the same ideal dV at the perigee speed [km^2/s^2]", vp_before * ideal + ideal ** 2 / 2, 11.66, 0.005)

cache = json.load(open("codes/_cache/ci-burns.json"))
rows = {}
for tw in (0.1, 0.03, 0.01):
    for N in (1, 2, 3, 4, 6, 8):
        en, ref, rp, ra, tt = perigee_burns(rl, rg, N, tw, centred=True)
        rows[(tw, N)] = (en, ref, rp, ra, tt)
for (tw, N), want_loss, want_t in (((0.1, 1), 107.0, 0.49), ((0.1, 2), 30.9, 3.06), ((0.1, 4), 7.9, 9.51), ((0.1, 8), 2.0, 23.51),
                                   ((0.03, 1), 499.8, 1.89), ((0.03, 2), 278.9, 4.05), ((0.03, 4), 84.2, 10.00), ((0.03, 8), 21.8, 23.76),
                                   ((0.01, 1), 640.3, 5.97), ((0.01, 2), 547.9, 7.91), ((0.01, 4), 387.6, 13.10), ((0.01, 8), 200.2, 25.66)):
    en, ref, rp, ra, tt = rows[(tw, N)]
    check(f"T/W0 = {tw}, {N} burns: loss against the best impulsive transfer to the same orbit [m/s]", (en - ref) * 1000, want_loss, 0.06)
    check(f"   elapsed time [h]", tt / 3600, want_t, 0.006)
check("single burn T/W0 = 0.1: final perigee radius [km] (start 6656)", rows[(0.1, 1)][2], 7226, 1.0)
check("single burn T/W0 = 0.01: final perigee radius [km]", rows[(0.01, 1)][2], 15252, 1.0)
json.dump({f"{k[0]}_{k[1]}": list(v) for k, v in rows.items()}, open("codes/_cache/ci-burns.json", "w"))
for tw in (1.0, 0.3, 0.003):
    en, ref, rp, ra, tt = perigee_burns(rl, rg, 1, tw, centred=False)
    print(f"      single burn T/W0 = {tw}: loss {(en-ref)*1000:.1f} m/s, final perigee {rp:.0f}, time {tt/3600:.2f} h")

# ---- extra checks quoted in the note
print("\nExtra checks")
from scipy.optimize import brentq


def hoh_ratio(R):
    return (np.sqrt(2 * R / (1 + R)) - 1) + (1 / np.sqrt(R)) * (1 - np.sqrt(2 / (1 + R)))


def bielliptic_ratio(R, rb):
    a1, a2 = (1 + rb) / 2, (R + rb) / 2
    v = lambda r, a: np.sqrt(2 / r - 1 / a)
    return (v(1, a1) - 1) + (v(rb, a2) - v(rb, a1)) + abs(1 / np.sqrt(R) - v(R, a2))


check("Hohmann dV/Vc1 at R = 10", hoh_ratio(10), 0.5298, 5e-5); check("  bi-elliptic, r_b -> infinity", bielliptic_ratio(10, 1e9), 0.5452, 5e-5)
check("R = 15.58: Hohmann", hoh_ratio(15.58), 0.5363, 5e-5); check("  bi-elliptic", bielliptic_ratio(15.58, 1e6 * 15.58), 0.5192, 5e-5)
check("R = 20: Hohmann", hoh_ratio(20), 0.5347, 5e-5); check("  bi-elliptic", bielliptic_ratio(20, 1e9), 0.5068, 5e-5)
check("R = 100: Hohmann", hoh_ratio(100), 0.4931, 5e-5); check("  bi-elliptic", bielliptic_ratio(100, 1e9), 0.4556, 5e-5)
check("  relative saving at R = 100", 1 - bielliptic_ratio(100, 1e9) / hoh_ratio(100), 0.076, 5e-4)
check("radius ratio where the r_b -> infinity bi-elliptic ties Hohmann", brentq(lambda R: bielliptic_ratio(R, 1e9 * R) - hoh_ratio(R), 10, 13), 11.939, 5e-4)
slope = lambda R: (bielliptic_ratio(R, R * (1 + 1e-6)) - bielliptic_ratio(R, R)) / (R * 1e-6)
check("radius ratio where the bi-elliptic starts to beat Hohmann for every r_b", brentq(slope, 14, 16), 15.582, 5e-4)
truth("  R = 14: a bi-elliptic with r_b = 1.001 R is worse, with r_b = 1e6 R is better", bielliptic_ratio(14, 14.014) > hoh_ratio(14) > bielliptic_ratio(14, 1.4e7))
truth("  R = 16: a bi-elliptic with r_b = 1.001 R is already better", bielliptic_ratio(16, 16.016) < hoh_ratio(16))
check("bi-elliptic with r_b = R is the Hohmann transfer (R = 7)", bielliptic_ratio(7, 7.0), hoh_ratio(7), 1e-12)
pay = 1200
e1 = rocket_dv(pay + 91 + 898 + 181 + 2174, pay + 91 + 898 + 181, 300) * 1000; e2 = rocket_dv(pay + 91 + 898, pay + 91, 300) * 1000
check("Ex 1: dV1 with a 1200 kg payload [m/s]", e1, 1915.0, 0.05); check("  dV2 [m/s]", e2, 1553.5, 0.05)
check("  the sum minus the Hohmann total [m/s]", e1 + e2 - (h["dv1"] + h["dv2"]) * 1000, -0.9, 0.05)
dv2_all = dv_from_angle(h["vc2"], h["va"], a1 + a2) * 1000
check("Ex 3: dV2 with the whole 26.53 deg at the apogee burn [m/s]", dv2_all, 2005.2, 0.05); check("  total [m/s]", h["dv1"] * 1000 + dv2_all, 4053.0, 0.05)
check("  above the stage's total [m/s]", h["dv1"] * 1000 + dv2_all - (dv1 + dv2) * 1000, 57.9, 0.05)
check("Ex 4: burn time [s]", (1 - 1 / 3.0) * 300 * 9.80665 / (1.5 * 9.80665), 133.3, 0.05)
check("  speed gained [m/s]", 300 * 9.80665 * np.log(3.0) - (1 - 1 / 3.0) * 300 * 9.80665 / 1.5, 1924.6, 0.05)
check("vertical ascent: ideal dV for m0/mf = 3 [m/s]", 300 * 9.80665 * np.log(3.0), 3232.1, 0.05)
check("  loss at T/W0 = 3 [m/s]", (1 - 1 / 3.0) * 300 * 9.80665 / 3.0, 653.8, 0.05)
check("burn 2: radius at the start [km]", r_s, 6934.8, 0.05); check("  radius at the end [km]", r_e, 7130.2, 0.05)
check("  energy shortfall over the perigee speed [m/s]", (vp_before * ideal + ideal ** 2 / 2 - eps_gain) / vp_before * 1000, 13.9, 0.05)
check("  gravity integral over ideal dV", yy[5] / ideal, 0.124, 5e-4)
check("  energy shortfall in per cent", (1 - eps_gain / (vp_before * ideal + ideal ** 2 / 2)) * 100, 1.05, 0.005)
check("Problem 6.1: window width of the pair of instants [min]", (c1 - c2) / OMEGA_E * 60, 4.0, 0.05)
for dra, a1w, a2w, g2w, f2w, tw_a, tw_b in ((500.0, 2.726, 20.655, 13.228, 171.013, 23.351, 17.967), (1000.0, 2.065, 17.794, 18.232, 167.431, 19.814, 15.785), (100.0, 3.167, 22.754, 6.046, 175.944, 25.914, 19.597)):
    vp, v2, g2, f2, a = transfer_state(r1, r2, r2 + dra)
    xa1 = plane_angle(h["vc1"], vp, dv1); xa2 = plane_angle(h["vc2"], v2, dv2, g2)
    check(f"table row {dra:g} km: alpha1", xa1 / D2R, a1w, 5e-4); check("   alpha2", xa2 / D2R, a2w, 5e-4); check("   gamma2", g2 / D2R, g2w, 5e-4)
    check("   f2", f2 / D2R, f2w, 5e-4); check("   theta_a", total_plane_angle(xa1, xa2, f2, True) / D2R, tw_a, 5e-4); check("   theta_b", total_plane_angle(xa1, xa2, f2, False) / D2R, tw_b, 5e-4)
check("limit transfer perigee speed Vc1 + dV1 [m/s]", (h["vc1"] + dv1) * 1000, 9845.7, 0.05)
check("chemical LEO-GEO mass ratio exp(4239/2942)", np.exp(4.2391 / (9.80665e-3 * 300)), 4.22, 0.005)
check("  fraction delivered", 1 / np.exp(4.2391 / (9.80665e-3 * 300)), 0.237, 5e-4)
check("ion LEO-GEO mass ratio exp(5963/29420)", np.exp(5.9627 / (9.80665e-3 * 3000)), 1.2247, 5e-4)
check("  fraction delivered", 1 / np.exp(5.9627 / (9.80665e-3 * 3000)), 0.8165, 5e-4)
print("\n%d checks, %d failed" % (len(RESULTS), RESULTS.count(False)))