"""Checks every number quoted in indian-missions-problems.qmd.

Each check recomputes a quantity from the models in _imcommon.py and compares it with the value
printed in the page, to the precision printed.
Run from the sfm-slides directory:   python3 codes/indian-missions-verify.py
"""
import sys
import numpy as np
from scipy.integrate import solve_ivp

sys.path.insert(0, "codes")
from _imcommon import *                                   # noqa: E402,F403
import _imcommon as C                                     # noqa: E402

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}")


def kms(x):
    return x * 1000.0                                     # km/s -> m/s


# =================================================================== Problem 1: Chandrayaan-1
rb = Re + 255
a_old, a_new = sma(255, 22860), sma(305, 37900)
check("1a r_b [km]", rb, 6633.1, 0.05)
check("1a a_old [km]", a_old, 17935.6, 0.05)
check("1a a_new [km]", a_new, 25480.6, 0.05)
check("1a v after [km/s]", vis(rb, a_new), 10.225, 5e-4)
check("1a v before [km/s]", vis(rb, a_old), 9.898, 5e-4)
for (d, _), s in zip(chain(CH1_EARTH), [327.0, 317.4, 217.1, 76.6, 38.3]):
    check("1a burn dv [m/s]", kms(d), s, 0.05)
check("1a total [m/s]", kms(chain_total(CH1_EARTH)), 976.4, 0.05)
single = vis(rb, sma(255, 380000)) - vis(rb, a_old)
check("1a single burn [m/s]", kms(single), 972.5, 0.05)
for o, a_s, T_s in zip(CH1_EARTH[:5], [17936, 25481, 43904, 88852, 140111], [6.64, 11.24, 25.43, 73.2, 145.0]):
    check(f"1b a of {o} [km]", sma(*o), a_s, 0.5)
    check(f"1b period of {o} [h]", period(sma(*o)) / 3600, T_s, 0.05 if T_s > 70 else 0.005)
for o, days, rev in zip(CH1_EARTH[:2], [1, 2], [3.6, 4.3]):
    check(f"1b revolutions in {o}", days * 24 / (period(sma(*o)) / 3600), rev, 0.05)
check("1b LTT a [km]", sma(465, 380000), 196600, 50)
check("1b LTT half period [d]", period(sma(465, 380000)) / 2 / DAY, 5.0, 0.05)
check("1c m0 ve / F [s]", 1380 * VE / F440, 9535, 0.5)
check("1c 1-exp(-dv/ve)", 1 - np.exp(-0.327 / VE), 0.1020, 5e-5)
check("1c burn 1 rocket-equation estimate [s]", 1380 * VE / F440 * (1 - np.exp(-0.327 / VE)), 973, 1)
fc = finite_chain(CH1_EARTH, 1380.0)
for (tb, dv, palt, m1), s in zip(fc, [993, 860, 531, 177, 86]):
    check("1c integrated burn time [s]", tb, s, 0.5)
check("1c five integrated burns [m/s]", kms(sum(x[1] for x in fc)), 988.8, 0.05)
check("1c integrated minus impulsive [m/s]", kms(sum(x[1] for x in fc) - chain_total(CH1_EARTH)), 12, 0.5)
one = finite_chain([CH1_EARTH[0], CH1_EARTH[-1]], 1380.0)[0]
check("1c one burn: time [s]", one[0], 2931, 0.5)
check("1c one burn: dv [m/s]", kms(one[1]), 1116, 0.5)
check("1c one burn: perigee after [km]", one[2], 683, 0.5)
check("1c one burn: fraction of the 6.6 h period", one[0] / period(sma(*CH1_EARTH[0])), 0.12, 0.005)
from scipy.optimize import brentq as _bq                     # noqa: E402
_a0 = sma(*CH1_EARTH[0])
_e0 = (CH1_EARTH[0][1] - CH1_EARTH[0][0]) / (2 * _a0)
_E = _bq(lambda E: E - _e0 * np.sin(E) - np.sqrt(mu / _a0 ** 3) * one[0] / 2, 0, np.pi)
check("1c one burn: true-anomaly arc [deg]", np.degrees(4 * np.arctan(np.sqrt((1 + _e0) / (1 - _e0)) * np.tan(_E / 2))), 176, 0.5)
check("1c one burn: excess over five [%]", 100 * (one[1] / sum(x[1] for x in fc) - 1), 13, 0.5)
vc = np.sqrt(mu / rb)
vT = vis(rb, sma(255, 380000))
check("1d circular speed [km/s]", vc, 7.752, 5e-4)
check("1d TLI speed [km/s]", vT, 10.870, 5e-4)
check("1d direct TLI [m/s]", kms(vT - vc), 3118, 0.5)
check("1d propellant fraction [%]", 100 * (1 - np.exp(-(vT - vc) / VE)), 64, 0.5)
check("1d propellant [kg]", 1380 * (1 - np.exp(-(vT - vc) / VE)), 890, 5)
check("1d PSLV share [m/s]", kms(vis(rb, a_old) - vc), 2146, 0.5)
vi1 = vinf_moon(465)
check("1e transfer speed at the Moon [km/s]", vis(rM, (Re + 465 + rM) / 2), 0.190, 5e-4)
check("1e Moon speed [km/s]", np.sqrt(mu / rM), 1.018, 5e-4)
check("1e vinf [km/s]", vi1, 0.828, 5e-4)
r504 = RM + 504
check("1e hyperbola at 504 km [km/s]", np.sqrt(vi1 ** 2 + 2 * muM / r504), 2.249, 5e-4)
check("1e orbit 504x7502 at perilune [km/s]", vis(r504, sma(504, 7502, RM), muM), 1.876, 5e-4)
check("1e LOI [m/s]", kms(capture_dv(vi1, 504, 7502)), 373, 0.5)
red1 = [(CH1_MOON[0], CH1_MOON[1], RM + 7502), (CH1_MOON[1], CH1_MOON[2], RM + 200),
        (CH1_MOON[2], CH1_MOON[3], RM + 255), (CH1_MOON[3], CH1_MOON[4], RM + 101)]
r1v = [abs(vis(r, sma(*b, RM), muM) - vis(r, sma(*a, RM), muM)) for a, b, r in red1]
for d, s in zip(r1v, [26, 446, 18, 33]):
    check("1e lowering burn [m/s]", kms(d), s, 0.5)
staged1 = capture_dv(vi1, 504, 7502) + sum(r1v)
check("1e staged total [m/s]", kms(staged1), 897, 0.5)
check("1e hyperbola at 100 km [km/s]", np.sqrt(vi1 ** 2 + 2 * muM / (RM + 100)), 2.454, 5e-4)
check("1e circular at 100 km [km/s]", np.sqrt(muM / (RM + 100)), 1.634, 5e-4)
check("1e direct capture [m/s]", kms(capture_dv(vi1, 100, 100)), 820, 0.5)
check("1e staged penalty [m/s]", kms(staged1 - capture_dv(vi1, 100, 100)), 76, 0.5)
check("1e smallest capture at 504 km [m/s]",
      kms(np.sqrt(vi1 ** 2 + 2 * muM / r504) - np.sqrt(2 * muM / r504)), 158, 0.5)
m_loi = fc[-1][3]
check("1e mass at LOI [kg]", m_loi, 997, 0.5)
check("1e propellant in 817 s [kg]", F440 / VE * 817, 118, 0.5)
check("1e dv of the 817 s burn [m/s]", kms(VE * np.log(m_loi / (m_loi - F440 / VE * 817))), 384, 0.5)
check("1f mass budget [m/s]", kms(VE * np.log(1380 / 675)), 2174, 0.5)
check("1f impulsive Earth + Moon [m/s]", kms(chain_total(CH1_EARTH) + staged1), 1873, 0.5)
check("1f unexplained [m/s]", kms(VE * np.log(1380 / 675) - chain_total(CH1_EARTH) - staged1), 300, 3)
check("1f explained fraction [%]", 100 * (chain_total(CH1_EARTH) + staged1) / (VE * np.log(1380 / 675)), 86, 0.5)

# =================================================================== Problem 2: Chandrayaan-2
for (d, _), s in zip(chain(CH2_EARTH)[1:], [94.2, 112.4, 76.9, 125.3, 150.3]):
    check("2a burn dv [m/s]", kms(d), s, 0.05)
check("2a burn 1 is a trim (left out)", float(np.isnan(chain(CH2_EARTH)[0][0])), 1.0, 0)
for o, m, dv_s, mp_s, eps_s in ((CH1_EARTH, 1380, 976, 379, -11.1), (CH2_EARTH, 3850, 559, 647, -6.8),
                                (CH3_EARTH, 3900, 666, 768, -8.1)):
    tot = chain_total(o)
    check(f"2a total from {o[0]} [m/s]", kms(tot), dv_s, 0.5)
    check(f"2a propellant from {o[0]} [kg]", m * (1 - np.exp(-tot / VE)), mp_s, 0.5)
    check(f"2a launch-orbit energy {o[0]} [km2/s2]", -mu / (2 * sma(*o[0])), eps_s, 0.05)
d23 = chain_total(CH3_EARTH) - chain_total(CH2_EARTH)
check("2a CH3 minus CH2 [m/s]", kms(d23), 107, 0.5)
check("2a its propellant on 3850 kg [kg]", 3850 * (1 - np.exp(-d23 / VE)), 130, 5)
for o, s in ((CH1_EARTH, 2146), (CH2_EARTH, 2595), (CH3_EARTH, 2475)):
    r = Re + o[0][0]
    check(f"2b launcher beyond circular {o[0]} [m/s]", kms(vis(r, sma(*o[0])) - np.sqrt(mu / r)), s, 0.5)
r170 = Re + 170
check("2b TLI beyond circular from 170 km [km/s]", vis(r170, (r170 + rM) / 2) - np.sqrt(mu / r170), 3.1, 0.05)
vi2 = vinf_moon(276)
r114 = RM + 114
check("2c vinf [km/s]", vi2, 0.830, 5e-4)
check("2c r [km]", r114, 1851.4, 0.05)
check("2c hyperbola [km/s]", np.sqrt(vi2 ** 2 + 2 * muM / r114), 2.447, 5e-4)
check("2c orbit [km/s]", vis(r114, sma(114, 18072, RM), muM), 2.201, 5e-4)
check("2c LOI [m/s]", kms(capture_dv(vi2, 114, 18072)), 246, 0.5)
check("2c escape speed [km/s]", np.sqrt(2 * muM / r114), 2.301, 5e-4)
cap2 = np.sqrt(vi2 ** 2 + 2 * muM / r114) - np.sqrt(2 * muM / r114)
check("2c smallest capture [m/s]", kms(cap2), 145, 0.5)
check("2c fraction [%]", 100 * cap2 / capture_dv(vi2, 114, 18072), 59, 0.5)
m2 = 3850 * np.exp(-chain_total(CH2_EARTH) / VE)
check("2c mass at LOI [kg]", m2, 3203, 0.5)
check("2c propellant in 1738 s [kg]", F440 / VE * 1738, 252, 0.5)
check("2c mass after [kg]", m2 - F440 / VE * 1738, 2952, 0.5)
check("2c dv of the 1738 s burn [m/s]", kms(VE * np.log(m2 / (m2 - F440 / VE * 1738))), 249, 0.5)
red2 = [abs(vis(RM + a[0], sma(*b, RM), muM) - vis(RM + a[0], sma(*a, RM), muM))
        for a, b in zip(CH2_MOON[:-1], CH2_MOON[1:])]
for d, s in zip(red2, [183, 185, 199, 9]):
    check("2d lowering burn [m/s]", kms(d), s, 0.5)
check("2d lowering total [m/s]", kms(sum(red2)), 576, 0.5)
check("2d staged total [m/s]", kms(capture_dv(vi2, 114, 18072) + sum(red2)), 822, 0.5)
dir2 = capture_dv(vi2, 119, 127)
check("2d direct capture [m/s]", kms(dir2), 817, 0.5)
check("2d single-burn time [s]", m2 * VE / F440 * (1 - np.exp(-dir2 / VE)), 5215, 1)
check("2d Earth-bound days (22 Jul to 14 Aug)", 14 + 31 - 22, 23, 0)

# =================================================================== Problem 3: Chandrayaan-3
vi3 = vinf_moon(288)
r164 = RM + 164
check("3a vinf [km/s]", vi3, 0.830, 5e-4)
check("3a hyperbola [km/s]", np.sqrt(vi3 ** 2 + 2 * muM / r164), 2.418, 5e-4)
check("3a orbit [km/s]", vis(r164, sma(164, 18074, RM), muM), 2.169, 5e-4)
L3 = capture_dv(vi3, 164, 18074)
check("3a LOI [m/s]", kms(L3), 249, 0.5)
red3 = [abs(vis(RM + a[0], sma(*b, RM), muM) - vis(RM + a[0], sma(*a, RM), muM))
        for a, b in zip(CH3_MOON[:-1], CH3_MOON[1:])]
for d, s in zip(red3, [188, 185, 192, 2]):
    check("3a lowering burn [m/s]", kms(d), s, 0.5)
check("3a staged total [m/s]", kms(L3 + sum(red3)), 816, 0.5)
dir3 = capture_dv(vi3, 153, 163)
check("3a direct capture [m/s]", kms(dir3), 812, 0.5)
m3 = 3900 * np.exp(-chain_total(CH3_EARTH) / VE)
check("3a mass at LOI [kg]", m3, 3132, 0.5)
check("3a single-burn time [s]", m3 * VE / F440 * (1 - np.exp(-dir3 / VE)), 5070, 5)
check("3a propellant in 1835 s [kg]", F440 / VE * 1835, 266, 0.5)
check("3a dv of the 1835 s burn [m/s]", kms(VE * np.log(m3 / (m3 - F440 / VE * 1835))), 269, 0.5)
for a, b, s in (((153, 163), (113, 157), 10), ((113, 157), (25, 134), 25)):
    r = RM + a[1]
    check(f"3b deboost {a}->{b} at apolune [m/s]", kms(abs(vis(r, sma(*b, RM), muM) - vis(r, sma(*a, RM), muM))), s, 0.5)
check("3b speed at 25 km perilune [km/s]", vis(RM + 25, sma(25, 134, RM), muM), 1.7, 0.05)
r150 = RM + 150
check("3c r [km]", r150, 1887.4, 0.05)
check("3c period of 150 km circle [h]", period(r150, muM) / 3600, 2.04, 0.005)
check("3c a of 150x5112 [km]", sma(150, 5112, RM), 4368.4, 0.05)
check("3c period of 150x5112 [h]", period(sma(150, 5112, RM), muM) / 3600, 7.20, 0.005)
rpE, raE = Re + 180000, Re + 380000
vap = np.sqrt(2 * mu * rpE / (raE * (raE + rpE)))
check("3c target-orbit apogee speed [km/s]", vap, 0.819, 5e-4)
vipm = np.sqrt(mu / rM) - vap
check("3c vinf from the Moon [km/s]", vipm, 0.199, 5e-4)
vc150 = np.sqrt(muM / r150)
vp150 = vis(r150, sma(150, 5112, RM), muM)
vh150 = np.sqrt(vipm ** 2 + 2 * muM / r150)
check("3c circular speed [km/s]", vc150, 1.612, 5e-4)
check("3c 150x5112 perilune speed [km/s]", vp150, 2.018, 5e-4)
check("3c hyperbola [km/s]", vh150, 2.288, 5e-4)
check("3c apolune raise [m/s]", kms(vp150 - vc150), 406, 0.5)
check("3c TEI [m/s]", kms(vh150 - vp150), 270, 0.5)
check("3c total [m/s]", kms(vh150 - vc150), 676, 0.5)
check("3c escape only [m/s]", kms((np.sqrt(2) - 1) * vc150), 668, 0.5)
check("3c propellant [kg]", 448.6 * (np.exp((vh150 - vc150) / VE) - 1), 112, 0.5)
a13 = (mu * (13 * DAY) ** 2 / (4 * np.pi ** 2)) ** (1 / 3)
check("3d a for 13 d [km]", a13, 233500, 50)
check("3d perigee radius [km]", Re + 154000, 160400, 25)
check("3d apogee altitude [km]", 2 * a13 - 2 * Re - 154000, 300000, 1000)
check("3d planned period [d]", period(sma(180000, 380000)) / DAY, 17.7, 0.05)
check("3d crash from 150 km [m/s]", kms(vc150 - vis(r150, (r150 + RM) / 2, muM)), 34, 0.5)

# =================================================================== Problem 4: Mangalyaan
for (d, _), s in zip(chain(MOM_EARTH), [127.2, 200.4, 268.1, 32.5, 126.1, 101.8]):
    check("4a burn dv [m/s]", kms(d), s, 0.05)
check("4a total [m/s]", kms(chain_total(MOM_EARTH)), 856, 0.5)
p1, p2 = burn_dv(MOM_EARTH[3], (264.1, 100000))[0], burn_dv((264.1, 100000), MOM_EARTH[5])[0]
f1, f2 = burn_dv(MOM_EARTH[3], MOM_EARTH[4])[0], burn_dv(MOM_EARTH[4], MOM_EARTH[5])[0]
check("4a planned burn to 100 000 km [m/s]", kms(p1), 111.8, 0.05)
check("4a planned burn on to 118 642 km [m/s]", kms(p2), 46.9, 0.05)
check("4a planned pair [m/s]", kms(p1 + p2), 158.6, 0.05)
check("4a flown pair [m/s]", kms(f1 + f2), 158.6, 0.05)
rE, vE = planet(JD_TMI, 3)
rMa, vMa = planet(JD_MOI, 4)
v1, v2 = lambert(rE, rMa, (JD_MOI - JD_TMI) * DAY)
vd, va = np.linalg.norm(v1 - vE), np.linalg.norm(v2 - vMa)
check("4b time of flight [d]", JD_MOI - JD_TMI, 297.3, 0.05)
check("4b vinf at departure [km/s]", vd, 3.129, 5e-4)
check("4b C3 [km2/s2]", vd ** 2, 9.79, 0.005)
check("4b vinf at arrival [km/s]", va, 3.155, 5e-4)
ang = np.degrees(np.arctan2(np.cross(rE, rMa)[2], np.dot(rE, rMa))) % 360
check("4b transfer angle [deg]", ang, 218, 0.5)
at = (AU + 1.52368 * AU) / 2
hd, ha = vis(AU, at, muS) - np.sqrt(muS / AU), np.sqrt(muS / (1.52368 * AU)) - vis(1.52368 * AU, at, muS)
check("4b Hohmann a [AU]", at / AU, 1.262, 5e-4)
check("4b Hohmann vinf departure [km/s]", hd, 2.945, 5e-4)
check("4b Hohmann vinf arrival [km/s]", ha, 2.649, 5e-4)
check("4b Hohmann time [d]", period(at, muS) / 2 / DAY, 259, 0.5)
check("4b excess over Hohmann at departure [km/s]", vd - hd, 0.18, 0.005)
check("4b excess over Hohmann at arrival [km/s]", va - ha, 0.51, 0.005)
best = (9e9, 0, 0)
for dd in np.arange(-40, 41, 1.0):
    rE_, vE_ = planet(JD_TMI + dd, 3)
    for tt in np.arange(250, 361, 1.0):
        rM_, vM_ = planet(JD_TMI + dd + tt, 4)
        a_, b_ = lambert(rE_, rM_, tt * DAY)
        s = np.linalg.norm(a_ - vE_) + np.linalg.norm(b_ - vM_)
        if s < best[0]:
            best = (s, dd, tt)
check("4b best total vinf in the window [km/s]", best[0], 6.24, 0.005)
check("4b best departure, days after MOM's", best[1], 5, 0)
check("4b best time of flight [d]", best[2], 295, 0)
check("4b MOM total vinf [km/s]", vd + va, 6.28, 0.005)
rp = Re + 264.1
vh = np.sqrt(vd ** 2 + 2 * mu / rp)
check("4c r [km]", rp, 6642.2, 0.05)
check("4c 2 mu / r [km2/s2]", 2 * mu / rp, 120.02, 0.005)
check("4c hyperbola [km/s]", vh, 11.394, 5e-4)
check("4c v_p final orbit [km/s]", vis(rp, sma(*MOM_EARTH[-1])), 10.777, 5e-4)
check("4c TMI [m/s]", kms(vh - vis(rp, sma(*MOM_EARTH[-1]))), 616, 0.5)
check("4c v_p launch orbit [km/s]", vis(rp, sma(*MOM_EARTH[0])), 9.921, 5e-4)
check("4c TMI from launch orbit [m/s]", kms(vh - vis(rp, sma(*MOM_EARTH[0]))), 1472, 0.5)
check("4c circular speed [km/s]", np.sqrt(mu / rp), 7.747, 5e-4)
check("4c TMI from circle [m/s]", kms(vh - np.sqrt(mu / rp)), 3647, 0.5)
check("2d single-burn time [min]", m2 * VE / F440 * (1 - np.exp(-dir2 / VE)) / 60, 87, 0.5)
check("4c chain + TMI [m/s]", kms(chain_total(MOM_EARTH) + vh - vis(rp, sma(*MOM_EARTH[-1]))), 1472, 0.5)
mE = 1337.2 * np.exp(-chain_total(MOM_EARTH) / VE)
check("4c single 1472 m/s burn [min]", mE * VE / F440 * (1 - np.exp(-(vh - vis(rp, sma(*MOM_EARTH[0]))) / VE)) / 60,
      45, 0.5)
rmo = RMa + 421.7
amo = sma(421.7, 76993.6, RMa)
check("4d r [km]", rmo, 3817.9, 0.05)
check("4d hyperbola [km/s]", np.sqrt(va ** 2 + 2 * muMa / rmo), 5.691, 5e-4)
check("4d a [km]", amo, 42104, 0.5)
check("4d orbit speed [km/s]", vis(rmo, amo, muMa), 4.628, 5e-4)
moi = np.sqrt(va ** 2 + 2 * muMa / rmo) - vis(rmo, amo, muMa)
check("4d MOI [m/s]", kms(moi), 1063, 0.5)
check("4d period [h]", period(amo, muMa) / 3600, 72.86, 0.005)
check("4d ISRO period 72 h 51 min 51 s [h]", 72 + 51 / 60 + 51 / 3600, 72.86, 0.005)
r400 = RMa + 400
check("4d hyperbola at 400 km [km/s]", np.sqrt(va ** 2 + 2 * muMa / r400), 5.702, 5e-4)
check("4d circular at 400 km [km/s]", np.sqrt(muMa / r400), 3.359, 5e-4)
moi400 = np.sqrt(va ** 2 + 2 * muMa / r400) - np.sqrt(muMa / r400)
check("4d MOI into 400 km circle [m/s]", kms(moi400), 2343, 0.5)
mT = mE * np.exp(-0.648 / VE)
check("4d mass after TMI [kg]", mT, 815, 0.5)
check("4d propellant left after TMI [kg]", 852 - (1337.2 - mT), 330, 0.5)
check("4d propellant for the flown MOI [kg]", mT * (1 - np.exp(-moi / VE)), 241, 0.5)
check("4d propellant for the circle [kg]", mT * (1 - np.exp(-moi400 / VE)), 438, 0.5)
check("4e mass after MOI [kg]", 1337.2 - 852 + 40, 525, 0.5)
check("4e total dv [m/s]", kms(VE * np.log(1337.2 / (1337.2 - 852 + 40))), 2841, 0.5)
imp4 = chain_total(MOM_EARTH) + vh - vis(rp, sma(*MOM_EARTH[-1])) + moi
check("4e impulsive sum [m/s]", kms(imp4), 2535, 0.5)
check("4e flown TMI minus impulsive [m/s]", 648 - kms(vh - vis(rp, sma(*MOM_EARTH[-1]))), 32, 0.5)
check("4e remainder after the TMI loss [m/s]", kms(VE * np.log(1337.2 / 525.2) - imp4) - 32, 250, 30)

# =================================================================== Problem 5: Aditya-L1, two-body
for (d, _), s in zip(chain(AL1_EARTH), [104.2, 372.7, 268.5, 165.2]):
    check("5a burn dv [m/s]", kms(d), s, 0.05)
check("5a total [m/s]", kms(chain_total(AL1_EARTH)), 911, 0.5)
days = [1, 2, 5, 5, 4]
passes = [d * 24 / (period(sma(*o)) / 3600) for d, o in zip(days, AL1_EARTH)]
for o, T_s in zip(AL1_EARTH, [5.7, 6.5, 12.0, 24.1, 48.5]):
    check(f"5a period of {o} [h]", period(sma(*o)) / 3600, T_s, 0.05)
for p, s in zip(passes, [4, 7, 10, 5, 2]):
    check("5a perigee passes", p, s, 0.5)
check("5a total passes", sum(passes), 28, 1.0)
rp5 = Re + 256
aK = (rp5 + 1.5e6) / 2
check("5b r_p [km]", rp5, 6634.1, 0.05)
check("5b v_p now [km/s]", vis(rp5, sma(*AL1_EARTH[-1])), 10.689, 5e-4)
check("5b a [km]", aK, 753317, 1)
check("5b v_p for 1.5e6 km [km/s]", vis(rp5, aK), 10.938, 5e-4)
check("5b TL1I two-body [m/s]", kms(vis(rp5, aK) - vis(rp5, sma(*AL1_EARTH[-1]))), 249, 0.5)
check("5b escape speed [km/s]", np.sqrt(2 * mu / rp5), 10.962, 5e-4)
check("5b escape burn [m/s]", kms(np.sqrt(2 * mu / rp5) - vis(rp5, sma(*AL1_EARTH[-1]))), 273, 0.5)
check("5b escape minus TL1I [m/s]", kms(np.sqrt(2 * mu / rp5) - vis(rp5, aK)), 24, 0.5)
check("5b half period [d]", period(aK) / 2 / DAY, 37.7, 0.05)
rsoi = 1.496e8 * (mu / muS) ** 0.4
check("5b SOI radius [km]", rsoi, 925000, 1000)
e5 = 1 - rp5 / aK
check("5b e", e5, 0.9912, 5e-5)
E5 = np.arccos((1 - rsoi / aK) / e5)
check("5b time to the SOI [d]", (E5 - e5 * np.sin(E5)) / np.sqrt(mu / aK ** 3) / DAY, 10.0, 0.05)

# =================================================================== Problem 6: Aditya-L1, CR3BP
check("6 mu", MU3, 3.0404e-6, 5e-11)
check("6a time unit [d]", TU / DAY, 58.13, 0.005)
check("6a speed unit [km/s]", VU, 29.785, 5e-4)
xL1 = lagrange_L1()
gam = 1 - MU3 - xL1
check("6a gamma [AU]", gam, 0.0100110, 5e-8)
check("6a gamma [million km]", gam * AU / 1e6, 1.4976, 5e-5)
check("6a Hill gamma [AU]", (MU3 / 3) ** (1 / 3), 0.0100447, 5e-8)
check("6a Hill [million km]", (MU3 / 3) ** (1 / 3) * AU / 1e6, 1.5027, 5e-5)
check("6a Hill error [%]", 100 * ((MU3 / 3) ** (1 / 3) / gam - 1), 0.34, 0.005)
s0, T = halo(120000.0)
CL1 = jacobi(np.array([xL1, 0, 0, 0, 0, 0]))
Ch = jacobi(s0)
check("6b C at L1", CL1, 3.000898, 5e-7)
check("6b C on the halo", Ch, 3.000827, 5e-7)
rpc = (Re + 256) / AU
u = np.array([-1.0, 0, 0])                                  # perigee on the Sun side; C does not depend on it


def C_perigee(v_in):
    rr = rpc * u
    vrot = v_in / VU * np.cross([0, 0, 1], u) - np.cross([0, 0, 1], rr)
    return jacobi(np.concatenate([[1 - MU3, 0, 0] + rr, vrot]))


check("6b C at perigee, 10.938 km/s", C_perigee(10.938), 3.002280, 5e-7)
from scipy.optimize import brentq                          # noqa: E402
vhal = brentq(lambda v: C_perigee(v) - Ch, 10.5, 11.5)
vneck = brentq(lambda v: C_perigee(v) - CL1, 10.5, 11.5)
check("6b perigee speed for C_halo [km/s]", vhal, 10.997, 5e-4)
check("6b perigee speed for C_L1 [km/s]", vneck, 10.994, 5e-4)
check("6b TL1I needed [m/s]", kms(vhal - vis(Re + 256, sma(*AL1_EARTH[-1]))), 307, 0.5)
ts, Hh = propagate(s0, T, 4000)
Axc, Ayc = (Hh[0].max() - Hh[0].min()) / 2 * AU, Hh[1].max() * AU
check("6c period [d]", T * TU / DAY, 177.87, 0.005)
check("6c A_x [km]", Axc, 208270, 1)
check("6c A_y [km]", Ayc, 667997, 1)
check("6c A_z [km]", Hh[2].max() * AU, 120000, 1)
check("6c A_x vs ISRO [%]", 100 * abs(Axc / 209200 - 1), 0.4, 0.05)
check("6c A_y vs ISRO [%]", 100 * abs(Ayc / 663200 - 1), 0.7, 0.05)
check("6c period vs ISRO [%]", 100 * abs(T * TU / DAY / 177.86 - 1), 0.004, 0.0005)
c2 = MU3 / gam ** 3 + (1 - MU3) / (1 - gam) ** 3
wp = np.sqrt((2 - c2 + np.sqrt(9 * c2 ** 2 - 8 * c2)) / 2)
check("6c c2", c2, 4.0611, 5e-5)
check("6c omega_p", wp, 2.0865, 5e-5)
check("6c in-plane period [d]", 2 * np.pi / wp * TU / DAY, 175.06, 0.005)
check("6c omega_v", np.sqrt(c2), 2.0152, 5e-5)
check("6c vertical period [d]", 2 * np.pi / np.sqrt(c2) * TU / DAY, 181.25, 0.005)
check("6c k = Ay/Ax (linear)", (wp ** 2 + 1 + 2 * c2) / (2 * wp), 3.23, 0.005)
check("6c Ay/Ax (computed)", Ayc / Axc, 3.21, 0.005)
S0 = np.concatenate([s0, np.eye(6).ravel()])
Mon = solve_ivp(C._cr3bp_stm, [0, T], S0, rtol=1e-12, atol=1e-12).y[6:, -1].reshape(6, 6)
lam = max(abs(np.linalg.eigvals(Mon)))
check("6d monodromy eigenvalue", lam, 1735, 0.5)
check("6d ln(lambda)", np.log(lam), 7.459, 5e-4)
check("6d e-folding time [d]", T * TU / DAY / np.log(lam), 23.8, 0.05)
lin = np.sqrt((c2 - 2 + np.sqrt(9 * c2 ** 2 - 8 * c2)) / 2)
check("6d linear growth rate", lin, 2.533, 5e-4)
check("6d linear e-folding [d]", TU / DAY / lin, 23.0, 0.05)
check("6d growth of 1 cm/s in 90 d [cm/s]", np.exp(90 / (T * TU / DAY / np.log(lam))), 44, 0.5)
check("6d SK spacing 6 Jan -> 22 Feb [d]", 25 + 22, 47, 0)
check("6d SK spacing 22 Feb -> 7 Jun [d]", 7 + 31 + 30 + 31 + 7 - 0, 106, 0)
check("6d SK spacing 7 Jun -> 2 Jul [d]", 23 + 2, 25, 0)
check("5 TL1I to HOI [d] (19 Sep 2023 -> 6 Jan 2024)", 11 + 31 + 30 + 31 + 6, 109, 0)

import glob                                                 # noqa: E402
scan = np.vstack([np.load(f)[:, :5] for f in sorted(glob.glob("codes/data/aditya-hoi-*.npy"))])
for lo, s in ((40, 205), (110, 150), (140, 85), (180, 51)):
    sel = scan[(scan[:, 1] >= lo) & (scan[:, 1] < lo + 10)]
    check(f"6e cheapest insertion, {lo}-{lo + 10} d [m/s]", sel[:, 2].min(), s, 1.0)
w, Vv = np.linalg.eig(Mon)
vs = np.real(Vv[:, np.argmin(abs(w))])
Mm = solve_ivp(C._cr3bp_stm, [0, T], S0, rtol=1e-12, atol=1e-12, dense_output=True)
Sx = Mm.sol(0.875 * T)
vdir = Sx[6:].reshape(6, 6) @ vs
vdir /= np.linalg.norm(vdir[:3])
ym = Sx[:6] + 200 / AU * vdir


def evp(t, y):
    dd = y[:3] - [1 - MU3, 0, 0]
    return dd @ y[3:]


rm = solve_ivp(cr3bp, [0, -400 * DAY / TU], ym, rtol=1e-11, atol=1e-13, events=evp, dense_output=True)
best = min((np.linalg.norm(rm.sol(tq)[:3] - [1 - MU3, 0, 0]), tq) for tq in rm.t_events[0])
check("6e manifold perigee altitude [km]", best[0] * AU - Re, 64, 0.5)
check("6e manifold time [d]", -best[1] * TU / DAY, 327, 0.5)
fast = scan[(scan[:, 1] >= 40) & (scan[:, 1] < 50), 2].min() / 1000
check("6e propellant for the fast insertion [kg]", 1500 * (1 - np.exp(-fast / VE)), 98, 0.5)
check("6e as a fraction of 1500 kg [%]", 100 * (1 - np.exp(-fast / VE)), 7, 0.5)
check("6e manifold time / 110 d", 327 / 110, 3, 0.05)

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