"""Checks every number quoted in patched-conics.qmd, including the two numerical integrations of
the section "The method against the truth".
Run from the sfm-slides directory:   python3 codes/patched-conics-verify.py
"""
import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq

RESULTS = []


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


muS, AU = 1.32712440018e11, 1.495978707e8
mu, Re = 398600.4418, 6378.137
muM, RM, D = 4902.800066, 1737.4, 384400.0
muMa, RMa, aMa, TMa = 42828.37, 3396.2, 1.52368 * AU, 686.98
muJ, RJ, aJ = 1.26686534e8, 71492.0, 5.2026 * AU
muV, RV, aV = 324858.6, 6051.8, 0.72333 * AU
DAY = 86400.0
EPS = np.radians(23.4392911)


def vis(r, a, m):
    return np.sqrt(m * (2 / r - 1 / a))


# ------------------------------------------------------------------ section 3: accelerations
check("3 Sun's pull at 1 AU [mm/s2]", muS / AU ** 2 * 1e6, 5.93, 0.005)
check("3 Sun's pull at 1.26 AU [mm/s2]", muS / (1.26 * AU) ** 2 * 1e6, 3.7, 0.05)
check("3 Earth's pull at 0.6 AU [um/s2]", mu / (0.6 * AU) ** 2 * 1e9, 0.05, 0.005)
check("3 their ratio (order 1e5)", np.log10((muS / (1.26 * AU) ** 2) / (mu / (0.6 * AU) ** 2)), 4.9, 0.1)

# ------------------------------------------------------------------ section 4: the sphere of influence
mE = mu / muS
check("4 m/M for the Earth", mE, 3.0035e-6, 5e-10)
check("4 equal-pull radius [km]", AU * np.sqrt(mE) / (1 + np.sqrt(mE)), 259000, 500)
check("4 Sun/Earth pull on the Moon", (muS / AU ** 2) / (mu / D ** 2), 2.2, 0.05)
check("4 angle factor 4^(-1/10)", 4 ** -0.1, 0.871, 5e-4)
check("4 variation of the boundary [%]", 100 * (1 - 4 ** -0.1), 13, 0.5)
for name, m, R, Mc, Rb, s_soi, s_rad, s_frac in (
        ("Venus", muV, aV, muS, RV, 616000, 102, 0.0057), ("Earth", mu, AU, muS, Re, 925000, 145, 0.0062),
        ("Moon", muM, D, mu, RM, 66200, 38, 0.17), ("Mars", muMa, aMa, muS, RMa, 577000, 170, 0.0025),
        ("Jupiter", muJ, aJ, muS, RJ, 48.2e6, 674, 0.062)):
    rs = R * (m / Mc) ** 0.4
    check(f"4 SOI {name} [km]", rs, s_soi, 0.0012 * s_soi)
    check(f"4 SOI {name} in radii", rs / Rb, s_rad, 0.5)
    check(f"4 SOI {name} / R", rs / R, s_frac, 0.5 * 10 ** (np.floor(np.log10(s_frac)) - 1))
check("4 Hill radius of the Earth [million km]", AU * (mE / 3) ** (1 / 3) / 1e6, 1.50, 0.005)
check("4 ratio at the boundary (m/M)^(1/5)", mE ** 0.2, 0.079, 5e-4)
check("4 ratio on the Sun-Earth line 2(m/M)^(1/5) [%]", 100 * 2 * mE ** 0.2, 16, 0.5)
check("4 ratio on the exact boundary, along the line", (mE * 4) ** 0.2, 0.104, 5e-4)
for th in (0.0, 0.7, np.pi / 2):                  # the two ratios agree on the exact boundary
    f4 = 1 + 3 * np.cos(th) ** 2
    rL = AU * mE ** 0.4 * f4 ** -0.1
    check(f"4 Q_planet = Q_Sun on the exact boundary, theta={th:.2f}",
          (rL / AU) ** 3 / mE * np.sqrt(f4) - mE * (AU / rL) ** 2, 0, 1e-12)
check("4 sphere beyond the exact boundary on the line [%]", 100 * (4 ** 0.1 - 1), 15, 0.5)
for name, m, Mc, s_q, tol in (("Venus", muV, muS, 0.075, 5e-4), ("Earth", mu, muS, 0.079, 5e-4),
                              ("Moon", muM, mu, 0.41, 5e-3), ("Mars", muMa, muS, 0.050, 5e-4),
                              ("Jupiter", muJ, muS, 0.25, 5e-3)):
    check(f"4 Q at the boundary, {name}", (m / Mc) ** 0.2, s_q, tol)
check("4 Jupiter/Earth mass", muJ / mu, 318, 0.5)
check("4 318^(1/5)", (muJ / mu) ** 0.2, 3.2, 0.05)
check("4 Earth/Moon mass", mu / muM, 81, 0.5)
check("4 Moon SOI [km]", D * (muM / mu) ** 0.4, 66183, 0.5)

# ------------------------------------------------------------------ Example 1
r1, r2 = AU, aMa
at = (r1 + r2) / 2
VE, Vt1, VMa, Vt2 = np.sqrt(muS / r1), vis(r1, at, muS), np.sqrt(muS / r2), vis(r2, at, muS)
check("E1 a_t [AU]", at / AU, 1.26184, 5e-6)
check("E1 V_Earth [km/s]", VE, 29.785, 5e-4)
check("E1 V_t1 [km/s]", Vt1, 32.729, 5e-4)
check("E1 V_Mars [km/s]", VMa, 24.129, 5e-4)
check("E1 V_t2 [km/s]", Vt2, 21.480, 5e-4)
check("E1 flight time [d]", np.pi * np.sqrt(at ** 3 / muS) / DAY, 258.9, 0.05)
vi1, vi2 = Vt1 - VE, VMa - Vt2
check("E1 vinf departure [km/s]", vi1, 2.945, 5e-4)
check("E1 vinf arrival [km/s]", vi2, 2.649, 5e-4)
check("E1 V_t2 = V_M sqrt(2r1/(r1+r2))", VMa * np.sqrt(2 * AU / (AU + aMa)) - Vt2, 0, 1e-9)
rp = Re + 300
vc = np.sqrt(mu / rp)
vp1 = np.sqrt(vi1 ** 2 + 2 * mu / rp)
check("E1 r_p [km]", rp, 6678.1, 0.05)
check("E1 v_c [km/s]", vc, 7.726, 5e-4)
check("E1 vinf^2 [km2/s2]", vi1 ** 2, 8.67, 0.005)
check("E1 2mu/r_p [km2/s2]", 2 * mu / rp, 119.37, 0.005)
check("E1 v_p1 [km/s]", vp1, 11.316, 5e-4)
check("E1 departure burn [km/s]", vp1 - vc, 3.590, 5e-4)
e1 = 1 + rp * vi1 ** 2 / mu
check("E1 e_1", e1, 1.1453, 5e-5)
check("E1 f_inf [deg]", np.degrees(np.arccos(-1 / e1)), 150.8, 0.05)
aH = mu / vi1 ** 2
rsoiE = AU * mE ** 0.4
F = np.arccosh((1 + rsoiE / aH) / e1)
check("E1 |a| [km]", aH, 45960, 10)
check("E1 F at the SOI", F, 3.61, 0.005)
tS = np.sqrt(aH ** 3 / mu) * (e1 * np.sinh(F) - F) / DAY
check("E1 time to the SOI [d]", tS, 3.2, 0.05)
check("E1 fraction of the trip [%]", 100 * tS / 258.9, 1.2, 0.05)
vS = np.sqrt(vi1 ** 2 + 2 * mu / rsoiE)
check("E1 speed at the SOI [km/s]", vS, 3.088, 5e-4)
check("E1 excess over vinf [%]", 100 * (vS / vi1 - 1), 4.9, 0.05)
rpM = RMa + 300
vp2 = np.sqrt(vi2 ** 2 + 2 * muMa / rpM)
check("E1 Mars r_p [km]", rpM, 3696.2, 0.05)
check("E1 v_p2 [km/s]", vp2, 5.495, 5e-4)
check("E1 Mars circular [km/s]", np.sqrt(muMa / rpM), 3.404, 5e-4)
check("E1 capture into circle [km/s]", vp2 - np.sqrt(muMa / rpM), 2.091, 5e-4)
aE = (rpM + RMa + 33000) / 2
check("E1 ellipse periapsis speed [km/s]", vis(rpM, aE, muMa), 4.587, 5e-4)
check("E1 capture into ellipse [km/s]", vp2 - vis(rpM, aE, muMa), 0.908, 5e-4)
check("E1 ellipse period [h]", 2 * np.pi * np.sqrt(aE ** 3 / muMa) / 3600, 23.9, 0.05)
check("E1 aim radius [km]", rpM * np.sqrt(1 + 2 * muMa / (rpM * vi2 ** 2)), 7667, 0.5)
nMa = 360 / TMa
check("E1 Mars moves [deg]", nMa * 258.866, 135.7, 0.05)
check("E1 phase angle [deg]", 180 - nMa * 258.866, 44.3, 0.05)
check("E1 synodic period [d]", 1 / (1 / 365.256 - 1 / TMa), 780, 0.5)

# ------------------------------------------------------------------ Example 2
Rs = D * (muM / mu) ** 0.4
vM = np.sqrt((mu + muM) / D)
nM = vM / D
r0, v0, lam = Re + 300, 10.84, np.radians(60.0)
check("E2 v_Moon [km/s]", vM, 1.0245, 5e-5)
eps0 = v0 ** 2 / 2 - mu / r0
h0 = r0 * v0
p0 = h0 ** 2 / mu
e0 = p0 / r0 - 1
a0 = -mu / (2 * eps0)
check("E2 eps0 [km2/s2]", eps0, -0.9346, 5e-5)
check("E2 a0 [km]", a0, 213253, 0.5)
check("E2 h0 [km2/s]", h0, 72391, 0.5)
check("E2 p0 [km]", p0, 13147, 0.5)
check("E2 e0", e0, 0.9687, 5e-5)
check("E2 apogee [km]", a0 * (1 + e0), 419828, 0.5)


def lunar(lam):
    r1 = np.sqrt(D ** 2 + Rs ** 2 - 2 * D * Rs * np.cos(lam))
    v1 = np.sqrt(2 * (eps0 + mu / r1))
    phi1 = np.arccos(h0 / (r1 * v1))
    gam = np.arcsin(Rs / r1 * np.sin(lam))
    f1 = np.arccos((p0 / r1 - 1) / e0)
    E1 = 2 * np.arctan(np.sqrt((1 - e0) / (1 + e0)) * np.tan(f1 / 2))
    t1 = (E1 - e0 * np.sin(E1)) * np.sqrt(a0 ** 3 / mu)
    thM = f1 - gam
    P1 = r1 * np.array([np.cos(f1), np.sin(f1)])
    M = D * np.array([np.cos(thM), np.sin(thM)])
    V1 = mu / h0 * np.array([-np.sin(f1), e0 + np.cos(f1)])
    VMv = vM * np.array([-np.sin(thM), np.cos(thM)])
    rho, v2 = P1 - M, V1 - VMv
    h2 = rho[0] * v2[1] - rho[1] * v2[0]
    eps2 = v2 @ v2 / 2 - muM / np.linalg.norm(rho)
    e2 = np.sqrt(1 + 2 * eps2 * h2 ** 2 / muM ** 2)
    rpl = h2 ** 2 / muM / (1 + e2)
    return dict(r1=r1, v1=v1, phi1=phi1, gam=gam, f1=f1, E1=E1, t1=t1, thM=thM, v2=np.linalg.norm(v2),
                ang=np.arccos(V1 @ VMv / np.linalg.norm(V1) / vM), h2=h2, eps2=eps2, e2=e2, rp=rpl, rho=np.linalg.norm(rho))


L = lunar(lam)
check("E2 r1 [km]", L["r1"], 355953, 0.5)
check("E2 gamma1 [deg]", np.degrees(L["gam"]), 9.27, 0.005)
check("E2 v1 [km/s]", L["v1"], 0.6087, 5e-5)
check("E2 phi1 [deg]", np.degrees(L["phi1"]), 70.48, 0.005)
check("E2 f1 [deg]", np.degrees(L["f1"]), 173.83, 0.005)
check("E2 E1 [rad]", L["E1"], 2.3334, 5e-5)
check("E2 t1 [h]", L["t1"] / 3600, 70.75, 0.005)
check("E2 angle between velocities [deg]", np.degrees(L["ang"]), 61.21, 0.005)
check("E2 phi1 - gamma1 [deg]", np.degrees(L["phi1"] - L["gam"]), 61.21, 0.005)
check("E2 v2 [km/s]", L["v2"], 0.9053, 5e-5)
check("E2 |rho| = SOI [km]", L["rho"], Rs, 1e-6)
check("E2 eps2 [km2/s2]", L["eps2"], 0.3357, 5e-5)
check("E2 vinf [km/s]", np.sqrt(2 * L["eps2"]), 0.819, 5e-4)
check("E2 |h2| [km2/s]", abs(L["h2"]), 6371, 0.5)
check("E2 h2 is retrograde (negative)", float(L["h2"] < 0), 1, 0)
check("E2 e2", L["e2"], 1.4607, 5e-5)
check("E2 perilune radius [km]", L["rp"], 3364, 0.5)
check("E2 perilune altitude [km]", L["rp"] - RM, 1627, 0.5)
a2 = muM / (2 * L["eps2"])
F2 = np.arccosh((1 + Rs / a2) / L["e2"])
t2 = np.sqrt(a2 ** 3 / muM) * (L["e2"] * np.sinh(F2) - F2) / 3600
check("E2 |a2| [km]", a2, 7302, 0.5)
check("E2 SOI to perilune [h]", t2, 18.2, 0.05)
check("E2 total [h]", L["t1"] / 3600 + t2, 88.9, 0.05)
check("E2 Moon moves [deg]", np.degrees(nM * L["t1"]), 38.9, 0.05)
check("E2 Moon angle at the patch [deg]", np.degrees(L["thM"]), 164.56, 0.005)
M0 = L["thM"] - nM * L["t1"]
check("E2 Moon lead at injection [deg]", np.degrees(M0), 125.66, 0.005)
slope = lunar(np.radians(60.5))["rp"] - lunar(np.radians(59.5))["rp"]
check("E2 perilune per degree of lambda [km]", slope, 600, 10)
lam200 = brentq(lambda x: lunar(np.radians(x))["rp"] - RM - 200, 55, 60)
check("E2 lambda for a 200 km perilune [deg]", lam200, 57.49, 0.005)


def rhs_moon(tt, y):
    rr = y[:2]
    th = M0 + nM * tt
    rm = D * np.array([np.cos(th), np.sin(th)])
    acc = -mu * rr / np.linalg.norm(rr) ** 3 - muM * ((rr - rm) / np.linalg.norm(rr - rm) ** 3 + rm / D ** 3)
    return [y[2], y[3], acc[0], acc[1]]


def ev_moon(tt, y):
    th = M0 + nM * tt
    dd = y[:2] - D * np.array([np.cos(th), np.sin(th)])
    return dd @ (y[2:] - vM * np.array([-np.sin(th), np.cos(th)]))


num = solve_ivp(rhs_moon, [0, 6 * DAY], [r0, 0, 0, v0], rtol=1e-12, atol=1e-9, events=ev_moon,
                dense_output=True, max_step=600)
best = min((np.linalg.norm(num.sol(tq)[:2] - D * np.array([np.cos(M0 + nM * tq), np.sin(M0 + nM * tq)])), tq)
           for tq in num.t_events[0])
check("E2 numerical perilune altitude [km]", best[0] - RM, 1030, 10)
check("E2 numerical perilune time [h]", best[1] / 3600, 87.0, 0.05)
check("E2 patched minus numerical [km]", L["rp"] - best[0], 600, 15)

# ------------------------------------------------------------------ Example 3
atJ = (AU + aJ) / 2
VJ, Vm = np.sqrt(muS / aJ), vis(aJ, atJ, muS)
vi = VJ - Vm
check("E3 V- [km/s]", Vm, 7.415, 5e-4)
check("E3 V_J [km/s]", VJ, 13.058, 5e-4)
check("E3 vinf [km/s]", vi, 5.643, 5e-4)
check("E3 trip [yr]", np.pi * np.sqrt(atJ ** 3 / muS) / DAY / 365.25, 2.73, 0.005)
check("E3 5 R_J [km]", 5 * RJ, 357460, 0)
check("E3 vinf^2", vi ** 2, 31.85, 0.005)
eJ = 1 + 5 * RJ * vi ** 2 / muJ
dl = 2 * np.arcsin(1 / eJ)
check("E3 e", eJ, 1.0899, 5e-5)
check("E3 delta [deg]", np.degrees(dl), 133.1, 0.05)
vout = vi * np.array([np.sin(dl), -np.cos(dl)])
Vp = vout + np.array([0, VJ])
check("E3 vinf+ x [km/s]", vout[0], 4.118, 5e-4)
check("E3 vinf+ y [km/s]", vout[1], 3.859, 5e-4)
check("E3 V+ y [km/s]", Vp[1], 16.917, 5e-4)
check("E3 |V+| [km/s]", np.linalg.norm(Vp), 17.41, 0.005)
check("E3 energy gain [km2/s2]", VJ * (vout[1] + vi), 124.1, 0.05)
check("E3 energy gain = change of V^2/2", (np.linalg.norm(Vp) ** 2 - Vm ** 2) / 2, VJ * (vout[1] + vi), 1e-9)
En = np.linalg.norm(Vp) ** 2 / 2 - muS / aJ
check("E3 new energy [km2/s2]", En, -18.95, 0.005)
an = -muS / (2 * En)
hn = aJ * Vp[1]
en = np.sqrt(1 + 2 * En * hn ** 2 / muS ** 2)
check("E3 new a [AU]", an / AU, 23.4, 0.05)
check("E3 new e", en, 0.79, 0.005)
check("E3 new perihelion [AU]", an * (1 - en) / AU, 4.87, 0.005)
check("E3 new aphelion [AU]", an * (1 + en) / AU, 42, 0.5)
vesc = np.sqrt(2 * muS / aJ)
check("E3 solar escape at Jupiter [km/s]", vesc, 18.47, 0.005)
cd = (VJ ** 2 + vi ** 2 - vesc ** 2) / (2 * VJ * vi)
check("E3 cos delta for escape", cd, -0.941, 5e-4)
de = np.arccos(cd)
check("E3 delta for escape [deg]", np.degrees(de), 160.2, 0.05)
ee = 1 / np.sin(de / 2)
check("E3 e for escape", ee, 1.0151, 5e-5)
check("E3 r_p for escape [km]", (ee - 1) * muJ / vi ** 2, 60100, 50)
eg = 1 + RJ * vi ** 2 / muJ
dg = 2 * np.arcsin(1 / eg)
check("E3 grazing delta [deg]", np.degrees(dg), 158.4, 0.05)
check("E3 grazing |V+| [km/s]", np.sqrt(VJ ** 2 + vi ** 2 - 2 * VJ * vi * np.cos(dg)), 18.42, 0.005)

# ------------------------------------------------------------------ Example 4
dla = lambda ls: np.degrees(np.arcsin(-np.sin(EPS) * np.cos(np.radians(ls))))
check("E4 sin(eps)", np.sin(EPS), 0.3978, 5e-5)
check("E4 sin(delta) on 30 Nov", -np.sin(EPS) * np.cos(np.radians(248)), 0.149, 5e-4)
check("E4 DLA on 30 Nov [deg]", dla(248), 8.57, 0.005)
check("E4 DLA on 20 Mar [deg]", dla(0), -23.44, 0.005)
sb = np.cos(EPS) / np.cos(np.radians(13.72))
check("E4 sin(azimuth)", sb, 0.9444, 5e-5)
check("E4 azimuth [deg]", np.degrees(np.arcsin(sb)), 70.8, 0.05)
check("E4 other azimuth [deg]", 180 - np.degrees(np.arcsin(sb)), 109.2, 0.05)
al = np.radians(23.44 - 13.72)
dvt = np.sqrt(vc ** 2 + vp1 ** 2 - 2 * vc * vp1 * np.cos(al))
check("E4 tilted burn [km/s]", dvt, 3.924, 5e-4)
check("E4 extra [m/s]", 1000 * (dvt - (vp1 - vc)), 334, 0.5)
check("E4 extra [%]", 100 * (dvt / (vp1 - vc) - 1), 9.3, 0.05)

# ------------------------------------------------------------------ Example 5
beta = np.radians(1.85)
for th, i_s, v_s, d_s in ((150, 3.70, 3.57, 3.77), (160, 5.42, 4.17, 3.97), (170, 10.71, 6.53, 5.00), (175, 21.74, 12.14, 8.61)):
    it = np.arcsin(np.sin(beta) / np.sin(np.radians(th)))
    vinf = np.sqrt(Vt1 ** 2 + VE ** 2 - 2 * Vt1 * VE * np.cos(it))
    check(f"E5 i_t at {th} deg [deg]", np.degrees(it), i_s, 0.005)
    check(f"E5 vinf at {th} deg [km/s]", vinf, v_s, 0.005)
    check(f"E5 LEO dv at {th} deg [km/s]", np.sqrt(vinf ** 2 + 2 * mu / rp) - vc, d_s, 0.005)
eT = (r2 - r1) / (r2 + r1)
pT = at * (1 - eT ** 2)
hT = np.sqrt(muS * pT)
check("E5 break-point radius p_t [AU]", pT / AU, 1.2075, 5e-5)
check("E5 transverse speed [km/s]", hT / pT, 27.10, 0.005)
dvb = 2 * hT / pT * np.sin(beta / 2)
check("E5 broken-plane burn [km/s]", dvb, 0.875, 5e-4)
check("E5 broken-plane total [km/s]", vp1 - vc + dvb, 4.47, 0.005)


def tilt_dv(th):
    it = np.arcsin(np.sin(beta) / np.sin(np.radians(th)))
    vinf = np.sqrt(Vt1 ** 2 + VE ** 2 - 2 * Vt1 * VE * np.cos(it))
    return np.sqrt(vinf ** 2 + 2 * mu / rp) - vc


check("E5 crossover transfer angle [deg]", brentq(lambda x: tilt_dv(x) - (vp1 - vc + dvb), 150, 176), 167, 0.5)
check("E5 extra at 170 deg [km/s]", tilt_dv(170) - (vp1 - vc), 1.4, 0.05)

# ------------------------------------------------------------------ Example 6 (MOM arrival, from the Lambert arc)
import sys                                                    # noqa: E402
sys.path.insert(0, "codes")
from _imcommon import planet, lambert, JD_TMI, JD_MOI         # noqa: E402

rEa, vEa = planet(JD_TMI, 3)
rMa_, vMa_ = planet(JD_MOI, 4)
w1, w2 = lambert(rEa, rMa_, (JD_MOI - JD_TMI) * DAY)
via = w2 - vMa_
q = np.array([via[0], via[1] * np.cos(EPS) - via[2] * np.sin(EPS), via[1] * np.sin(EPS) + via[2] * np.cos(EPS)])
a0p, d0p = np.radians(317.68143), np.radians(52.88650)          # IAU pole of Mars (J2000)
pole = np.array([np.cos(d0p) * np.cos(a0p), np.cos(d0p) * np.sin(a0p), np.sin(d0p)])
dinf = np.degrees(np.arcsin(q @ pole / np.linalg.norm(q)))
check("E6 vinf [km/s]", np.linalg.norm(via), 3.155, 5e-4)
check("E6 declination to Mars's equator [deg]", dinf, -30.25, 0.005)
check("E6 cos(delta)", np.cos(np.radians(dinf)), 0.864, 5e-4)
check("E6 largest inclination [deg]", 180 - abs(dinf), 149.75, 0.005)
check("E6 MOM inside the range by [deg]", abs(150.0 - (180 - abs(dinf))), 0.25, 0.005)
rpm, ram = RMa + 421.7, RMa + 76993.6
am = (rpm + ram) / 2
vhm = np.sqrt(np.linalg.norm(via) ** 2 + 2 * muMa / rpm)
vpm = vis(rpm, am, muMa)
vam = vis(ram, am, muMa)
check("E6 hyperbola at periapsis [km/s]", vhm, 5.691, 5e-4)
check("E6 orbit at periapsis [km/s]", vpm, 4.628, 5e-4)
check("E6 coplanar MOI [km/s]", vhm - vpm, 1.063, 5e-4)
check("E6 MOI with a 30.25 deg turn [km/s]",
      np.sqrt(vhm ** 2 + vpm ** 2 - 2 * vhm * vpm * np.cos(np.radians(30.25))), 2.881, 5e-4)
check("E6 extra [km/s]", np.sqrt(vhm ** 2 + vpm ** 2 - 2 * vhm * vpm * np.cos(np.radians(30.25))) - (vhm - vpm), 1.8, 0.05)
check("E6 apoapsis speed [km/s]", vam, 0.220, 5e-4)
check("E6 half-angle [deg]", 30.25 / 2, 15.125, 1e-9)
check("E6 plane change at apoapsis [km/s]", 2 * vam * np.sin(np.radians(30.25 / 2)), 0.115, 5e-4)

# ------------------------------------------------------------------ section 8: Earth-Mars integration
nE = VE / AU


def rhs_em(tt, y):
    r = y[:2]
    th = nE * tt
    d = r - AU * np.array([np.cos(th), np.sin(th)])
    a = -muS * r / np.linalg.norm(r) ** 3 - mu * d / np.linalg.norm(d) ** 3
    return [y[2], y[3], a[0], a[1]]


ts = np.linspace(0, 320 * DAY, 64001)


def aphelion(vinf):
    e = 1 + rp * vinf ** 2 / mu
    w = np.pi / 2 - np.arccos(-1 / e)
    vpk = np.sqrt(vinf ** 2 + 2 * mu / rp)
    rr = rp * np.array([np.cos(w), np.sin(w)])
    vv = vpk * np.array([-np.sin(w), np.cos(w)])
    s = solve_ivp(rhs_em, [0, 320 * DAY], [AU + rr[0], rr[1], vv[0], VE + vv[1]], rtol=1e-12, atol=1e-6,
                  dense_output=True, max_step=6 * 3600)
    return np.linalg.norm(s.sol(ts)[:2], axis=0).max() / AU


Q = aphelion(vi1)
check("8 aphelion with the patched-conic burn [AU]", Q, 1.5271, 5e-5)
check("8 overshoot [km]", (Q - 1.52368) * AU, 510000, 10000)
vneed = brentq(lambda v: aphelion(v) - 1.52368, 2.8, 3.1, xtol=1e-7)
dvneed = np.sqrt(vneed ** 2 + 2 * mu / rp) - vc
check("8 burn that reaches Mars [km/s]", dvneed, 3.586, 5e-4)
check("8 patched-conic excess [m/s]", 1000 * ((vp1 - vc) - dvneed), 3.8, 0.05)
check("8 excess [%]", 100 * ((vp1 - vc) / dvneed - 1), 0.1, 0.05)

print()
nf = sum(not ok for _, ok in RESULTS)
print(f"{len(RESULTS) - nf} PASS, {nf} FAIL")
