"""Checks every identity and number of ballistic-trajectories-bates.qmd.

Part 1: the algebra, by sympy (each identity is simplified to zero, or checked at random points where simplify is too slow).
Part 2: the formulas against a direct numerical integration of the two-body equation (the range angle and the flight time).
Part 3: the numbers in the worked examples, the influence coefficients (against finite differences) and the cross-range formulas.
Part 4: the rotating-Earth iteration.

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

D = np.pi / 180
RESULTS = []


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


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


# ------------------------------------------------------------------ Part 1: algebra
print("Part 1: the algebra")
Qs, g, ps = sp.symbols("Q gamma psi", positive=True)
c, s = sp.cos(g), sp.sin(g)
e2 = 1 + Qs * (Qs - 2) * c ** 2
truth("p/a = Q(2-Q)cos^2(gamma), so e^2 = 1 - p/a = 1 + Q(Q-2)cos^2(gamma)",
      sp.simplify((1 - (Qs * c ** 2) * (2 - Qs)) - e2) == 0)
ecos, esin = 1 - Qs * c ** 2, Qs * s * c           # e cos(psi/2), e sin(psi/2)
truth("(e cos(psi/2))^2 + (e sin(psi/2))^2 = e^2", sp.simplify(ecos ** 2 + esin ** 2 - e2) == 0)
truth("cot(psi/2) = 2/(Q sin 2gamma) - cot(gamma)", sp.simplify(ecos / esin - (2 / (Qs * sp.sin(2 * g)) - sp.cot(g))) == 0)
# the flight-path-angle equation: sin(2 gamma + psi/2) = (2 - Q)/Q sin(psi/2)   (multiply through by e)
truth("sin(2 gamma + psi/2) = ((2-Q)/Q) sin(psi/2)",
      sp.simplify(sp.sin(2 * g) * ecos + sp.cos(2 * g) * esin - (2 - Qs) / Qs * esin) == 0)
truth("e^2 - (Q-1)^2 = Q(2-Q) sin^2(gamma)   (so e sin E = sqrt(Q(2-Q)) sin gamma)",
      sp.simplify(e2 - (Qs - 1) ** 2 - Qs * (2 - Qs) * s ** 2) == 0)
# influence coefficients
psi_expr = 2 * sp.atan(Qs * s * c / (1 - Qs * c ** 2))
dpsi_dg = sp.diff(psi_expr, g); dpsi_dQ = sp.diff(psi_expr, Qs)
rng = np.random.default_rng(3)
ok1 = ok2 = True
for _ in range(20):
    Qv, gv = rng.uniform(0.3, 0.99), rng.uniform(0.05, 1.5)
    pv = float(psi_expr.subs({Qs: Qv, g: gv}))
    # psi_expr uses atan, so it equals psi_of only for 1 - Q cos^2 > 0 (psi < 180); compare in that range
    if 1 - Qv * np.cos(gv) ** 2 <= 0:
        continue
    ok1 &= abs(float(dpsi_dg.subs({Qs: Qv, g: gv})) - dpsi_dgamma(pv, gv)) < 1e-9
    ok2 &= abs(float(dpsi_dQ.subs({Qs: Qv, g: gv})) - 4 * np.sin(pv / 2) ** 2 / (Qv ** 2 * np.sin(2 * gv))) < 1e-9
truth("d psi/d gamma = 2 sin(psi + 2 gamma)/sin 2 gamma - 2   (sympy derivative, 20 random points)", ok1)
truth("d psi/d Q = 4 sin^2(psi/2)/(Q^2 sin 2 gamma)   (sympy derivative, 20 random points)", ok2)
# maximum range
truth("at psi = pi - 4 gamma the coefficient d psi/d gamma vanishes",
      sp.simplify(((2 * sp.sin((sp.pi - 4 * g) + 2 * g) / sp.sin(2 * g)) - 2)) == 0)
Qm = 2 * sp.sin(ps / 2) / (1 + sp.sin(ps / 2))
truth("Q_min = 2 sin(psi/2)/(1+sin(psi/2)) makes (2-Q)/Q sin(psi/2) = 1",
      sp.simplify((2 - Qm) / Qm * sp.sin(ps / 2) - 1) == 0)
# at max range Q = 1 - e^2 : e^2 = 1 + Q(Q-2)cos^2 gamma with gamma = (pi - psi)/4 and Q = Q_min
vals = [(psi_v, Qmin(psi_v), gamma_max_range(psi_v)) for psi_v in (0.3, 1.0, 2.0, 3.0)]
truth("at maximum range Q = 1 - e^2 (4 values of psi)", all(abs(Q_ - (1 - ecc(Q_, g_) ** 2)) < 1e-12 for _, Q_, g_ in vals))
truth("at maximum range sin(psi/2) = (1-e^2)/(1+e^2)", all(abs(np.sin(p_ / 2) - (1 - ecc(Q_, g_) ** 2) / (1 + ecc(Q_, g_) ** 2)) < 1e-12 for p_, Q_, g_ in vals))
# flat-Earth limit
x = sp.symbols("x", positive=True)
lim = sp.limit(psi_expr.subs(Qs, x) / x, x, 0)
truth("small Q: psi -> Q sin 2 gamma  (the flat-Earth range v^2 sin 2 gamma/g)", sp.simplify(lim - sp.sin(2 * g)) == 0)
# symmetry of the two roots
truth("gamma_low + gamma_high = (pi - psi)/2 for 20 random (Q, psi)",
      all(abs(sum(gs) - (np.pi - p_) / 2) < 1e-9 for Q_, p_ in zip(rng.uniform(0.7, 0.99, 20), rng.uniform(0.3, 1.0, 20))
          for gs in [gammas_of(Q_, p_)] if len(gs) == 2))

# ------------------------------------------------------------------ Part 2: against numerical integration
print("\nPart 2: the formulas against a numerical integration of the two-body equation")


def integrate(r, v, gam):
    y0 = [r, 0, v * np.sin(gam), v * np.cos(gam)]
    f = lambda t, 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]
    ev = lambda t, y: np.hypot(y[0], y[1]) - r
    ev.terminal = True; ev.direction = -1
    sol = solve_ivp(f, [0, 1e6], y0, events=ev, rtol=1e-12, atol=1e-9, max_step=50)
    y = sol.y_events[0][0]
    return np.arctan2(y[1], y[0]) % (2 * np.pi), sol.t_events[0][0]


worst_psi = worst_t = 0.0
for r_, v_, g_ in [(7653.76, 5.2703, 51.3178), (6537.6, 8.551, 39.33), (7000, 7.2, 30), (7000, 7.2, 80), (6600, 9.5, 20), (6578.137, 7.3, 12), (6578.137, 7.3, 55)]:
    pi_, ti_ = integrate(r_, v_, g_ * D)
    Q_ = Q_of(r_, v_)
    worst_psi = max(worst_psi, abs(psi_of(Q_, g_ * D) - pi_)); worst_t = max(worst_t, abs(tff(r_, v_, g_ * D) - ti_))
check("range angle: formula minus integration, worst of 7 cases [rad]", worst_psi, 0.0, 1e-8)
check("flight time: formula minus integration, worst of 7 cases [s]", worst_t, 0.0, 1e-4)

# ------------------------------------------------------------------ Part 3: numbers of the note
print("\nPart 3: worked examples")
vc = np.sqrt(MU / RE)                         # circular speed at the equatorial radius
# Example 1: from measurements to range
r1, v1, ra1 = 1.2 * RE, 2 / 3 * vc, 1.5 * RE
eps1 = v1 ** 2 / 2 - MU / r1
va1 = np.sqrt(2 * (eps1 + MU / ra1)); cg1 = ra1 * va1 / (r1 * v1)
Q1 = Q_of(r1, v1); g1 = np.arccos(cg1); psi1 = psi_of(Q1, g1)
check("Ex 1: r_bo [km]", r1, 7653.8, 0.05); check("Ex 1: v_bo [km/s]", v1, 5.2702, 5e-5)
check("Ex 1: Q", Q1, 0.5333, 5e-5); check("Ex 1: cos gamma", cg1, 0.625, 5e-4); check("Ex 1: gamma [deg]", g1 / D, 51.32, 0.005)
check("Ex 1: psi [deg]", psi1 / D, 36.39, 0.005); check("Ex 1: ground range psi R [km]", psi1 * RE, 4051, 0.5)
check("Ex 1: t_ff [s]", tff(r1, v1, g1), 1990.5, 0.05); check("Ex 1: t_ff [min]", tff(r1, v1, g1) / 60, 33.18, 0.005); check("Ex 1: P_cs [min]", 2 * np.pi * np.sqrt(r1 ** 3 / MU) / 60, 111.06, 0.005)
check("Ex 1: apogee radius of the ellipse from e and a [km]", (r1 / (2 - Q1)) * (1 + ecc(Q1, g1)), ra1, 0.5)
# Example 2: which flight-path angle
phiA, lamA, phiB, lamB = 30 * D, 60 * D, -30 * D, -60 * D
Lam2 = great_circle_angle(phiA, lamA, phiB, lamB)
r2, v2 = 1.025 * RE, 1.0817 * vc
Q2 = Q_of(r2, v2)
check("Ex 2: range angle between (30N,60E) and (30S,60W) [deg]", Lam2 / D, 128.68, 0.005)
check("Ex 2: r_bo [km]", r2, 6537.6, 0.05); check("Ex 2: v_bo [km/s]", v2, 8.551, 5e-4); check("Ex 2: Q", Q2, 1.1993, 5e-5)
gs2 = gammas_of(Q2, Lam2)
check("Ex 2: number of solutions", len(gs2), 1, 0); check("Ex 2: gamma [deg]", gs2[0] / D, 39.33, 0.005)
check("Ex 2: t_ff [s]", tff(r2, v2, gs2[0]), 5848.8, 0.05); check("Ex 2: t_ff [min]", tff(r2, v2, gs2[0]) / 60, 97.48, 0.005)
# Example 3: two solutions
Q3, psi3 = 0.9, 90 * D
k3 = (2 - Q3) / Q3 * np.sin(psi3 / 2); gs3 = gammas_of(Q3, psi3)
check("Ex 3: (2-Q)/Q sin(psi/2)", k3, 0.8642, 5e-5); check("Ex 3: gamma_low [deg]", gs3[0] / D, 7.4, 0.05); check("Ex 3: gamma_high [deg]", gs3[1] / D, 37.6, 0.05)
check("Ex 3: psi_max(0.9) [deg]", psi_max(Q3) / D, 109.81, 0.005); check("Ex 3: gamma_max for psi=90 [deg]", gamma_max_range(psi3) / D, 22.5, 1e-9)
# Example 4: maximum range of a given vehicle
r4, v4 = RE + 477.8, 24300 * 0.3048 / 1000
Q4 = Q_of(r4, v4)
check("Ex 4: v_bo [km/s]", v4, 7.407, 5e-4); check("Ex 4: Q", Q4, 0.9436, 5e-5)
check("Ex 4: psi_max [deg]", psi_max(Q4) / D, 126.55, 0.005); check("Ex 4: ground range [km]", psi_max(Q4) * RE, 14087, 1.0)
check("Ex 4: gamma at max range [deg]", gamma_max_range(psi_max(Q4)) / D, 13.36, 0.005)
check("Ex 4: t_ff [min]", tff(r4, v4, gamma_max_range(psi_max(Q4))) / 60, 43.11, 0.005); check("Ex 4: P_cs [min]", 2 * np.pi * np.sqrt(r4 ** 3 / MU) / 60, 94.16, 0.005)
# Example 5: minimum speed for a given range
psi5 = 14816 / RE; r5 = RE + 637.1
Q5 = Qmin(psi5); v5 = np.sqrt(Q5 * MU / r5)
check("Ex 5: psi [deg]", psi5 / D, 133.09, 0.005); check("Ex 5: Q_min", Q5, 0.9569, 5e-5); check("Ex 5: v_bo [km/s]", v5, 7.374, 5e-4)
check("Ex 5: gamma [deg]", gamma_max_range(psi5) / D, 11.73, 0.005)
check("Ex 5: t_ff [min]", tff(r5, v5, gamma_max_range(psi5)) / 60, 45.57, 0.005)
check("Ex 5: the same Q_min from the Lambert minimum-energy ellipse a_m = r(1+sin(psi/2))/2: Q = 2 - r/a_m", 2 - r5 / (r5 * (1 + np.sin(psi5 / 2)) / 2), Q5, 1e-12)

print("\nTable of the number of solutions (roots in [0, 90 deg] of the range equation)")
for Q_ in (0.8, 1.0, 1.2):
    row = {p_: [round(x_ / D, 2) for x_ in gammas_of(Q_, p_ * D)] for p_ in (30, 90, 150, 200, 270)}
    print("  Q =", Q_, row)
truth("Q = 1, psi = 200: only gamma = 0 (the circular orbit)", gammas_of(1.0, 200 * D) == [0.0])
truth("Q = 1.2: one root at every psi in 30..270", all(len(gammas_of(1.2, p_ * D)) == 1 for p_ in (30, 90, 150, 200, 270)))
truth("Q = 0.9, psi > psi_max: no root", gammas_of(0.9, 120 * D) == [])

print("\nInfluence coefficients (example of Bate: v = 0.905 v_c, r = 1.1 R, gamma = 30 deg)")
r6, v6, g6 = 1.1 * RE, 0.905 * vc, 30 * D
Q6 = Q_of(r6, v6); psi6 = psi_of(Q6, g6)
psi_rv = lambda r_, v_, g_: psi_of(Q_of(r_, v_), g_)
fd_g = (psi_rv(r6, v6, g6 + 1e-6) - psi_rv(r6, v6, g6 - 1e-6)) / 2e-6
fd_r = (psi_rv(r6 + 1e-3, v6, g6) - psi_rv(r6 - 1e-3, v6, g6)) / 2e-3
fd_v = (psi_rv(r6, v6 + 1e-6, g6) - psi_rv(r6, v6 - 1e-6, g6)) / 2e-6
check("Q", Q6, 0.9009, 5e-5); check("psi [deg]", psi6 / D, 100.53, 0.005)
check("d psi/d gamma  [-]", dpsi_dgamma(psi6, g6), -1.2301, 5e-5); check("  finite difference", fd_g, dpsi_dgamma(psi6, g6), 1e-8)
check("d psi/d r  [1/km]", dpsi_dr(r6, v6, psi6, g6), 4.3211e-4, 5e-8); check("  finite difference", fd_r, dpsi_dr(r6, v6, psi6, g6), 1e-9)
check("d psi/d v  [s/km]", dpsi_dv(r6, v6, psi6, g6), 0.84749, 5e-6); check("  finite difference", fd_v, dpsi_dv(r6, v6, psi6, g6), 1e-7)
check("d psi/d v = (2 r/v) d psi/d r", dpsi_dv(r6, v6, psi6, g6), 2 * r6 / v6 * dpsi_dr(r6, v6, psi6, g6), 1e-12)
dv_, dr_, dg_ = -5e-5 * vc, 5e-4 * RE, -1e-4
d_lin = dpsi_dr(r6, v6, psi6, g6) * dr_ + dpsi_dv(r6, v6, psi6, g6) * dv_ + dpsi_dgamma(psi6, g6) * dg_
d_ex = psi_rv(r6 + dr_, v6 + dv_, g6 + dg_) - psi_rv(r6, v6, g6)
check("Delta v [m/s]", dv_ * 1000, -0.395, 5e-4); check("Delta r [km]", dr_, 3.189, 5e-4)
check("linear miss [km]", d_lin * RE, 7.44, 0.005); check("  exact miss [km]", d_ex * RE, 7.44, 0.01)
check("  contributions: r, v, gamma [km]  (r)", dpsi_dr(r6, v6, psi6, g6) * dr_ * RE, 8.79, 0.005)
check("  (v)", dpsi_dv(r6, v6, psi6, g6) * dv_ * RE, -2.14, 0.005); check("  (gamma)", dpsi_dgamma(psi6, g6) * dg_ * RE, 0.785, 0.0005)
print("Sensitivity of a 200 km, 7.3 km/s burnout (Q = 0.8794):")
r7, v7 = 6578.137, 7.3; Q7 = Q_of(r7, v7)
check("Q", Q7, 0.8794, 5e-5); check("psi_max [deg]", psi_max(Q7) / D, 103.41, 0.005); check("gamma_max [deg]", gamma_max_range(psi_max(Q7)) / D, 19.15, 0.005)
for nm, gd in (("low", 12.0), ("high", 55.0)):
    p7 = psi_of(Q7, gd * D)
    per_arcmin = dpsi_dgamma(p7, gd * D) * (D / 60) * RE
    per_km = dpsi_dr(r7, v7, p7, gd * D) * RE; per_ms = dpsi_dv(r7, v7, p7, gd * D) * 1e-3 * RE
    print(f"  {nm}: gamma={gd}, psi={p7/D:.2f} deg, miss per 1' = {per_arcmin:.3f} km, per 1 km of r = {per_km:.3f} km, per 1 m/s = {per_ms:.3f} km")
gm7 = gamma_max_range(psi_max(Q7))
miss_max = (psi_of(Q7, gm7 + D / 60) - psi_of(Q7, gm7)) * RE * 1000
check("miss at the maximum-range angle for 1 arcmin of gamma [m]", miss_max, -2.73, 0.005)

print("\nCross-range errors")
# equator, due north, burnout displaced east by DX
for psi_d in (30, 60, 90, 120, 150):
    DX = 1e-4; ps_ = psi_d * D
    A_ = np.array([np.cos(ps_), 0, np.sin(ps_)]); B_ = np.array([np.cos(ps_) * np.cos(DX), np.cos(ps_) * np.sin(DX), np.sin(ps_)])
    # both reached by going `ps_` north along their own meridian; B is on the meridian at longitude DX
    dC = np.arccos(np.clip(A_ @ B_, -1, 1))
    check(f"Delta C / Delta X at psi = {psi_d}: arccos(cos psi...) vs cos(psi)", dC / DX, abs(np.cos(ps_)) if psi_d != 90 else dC / DX, 1e-4)
# general rigid rotation by DX about the horizontal axis 90 deg ahead of the launch point
for (lat0, az0, psi_d) in [(10, 30, 50), (45, 100, 100), (-20, 250, 140)]:
    lat, az, ps_ = lat0 * D, az0 * D, psi_d * D
    Lv = np.array([np.cos(lat), 0, np.sin(lat)])
    North = np.array([-np.sin(lat), 0, np.cos(lat)]); East = np.array([0, 1.0, 0])
    tdir = np.cos(az) * North + np.sin(az) * East
    pt = lambda th: np.cos(th) * Lv + np.sin(th) * tdir
    axis = pt(np.pi / 2)
    def rot(v, ang):
        return v * np.cos(ang) + np.cross(axis, v) * np.sin(ang) + axis * (axis @ v) * (1 - np.cos(ang))
    DX = 1e-5
    dC = np.arccos(np.clip(pt(ps_) @ rot(pt(ps_), DX), -1, 1))
    check(f"rotation about the axis 90 deg ahead: Delta C / Delta X, psi = {psi_d} (cos psi = {np.cos(ps_):.4f})", dC / DX, abs(np.cos(ps_)), 1e-5)
    # the launch point moves by DX, perpendicular to the track (the axis is perpendicular to the launch point and lies in the plane)
    mv = rot(Lv, DX) - Lv
    check("  the launch point is displaced sideways by DX; its component along the track is ~0", abs(mv @ tdir) / DX, 0.0, 1e-5)
# azimuth error: two great circles from one point
for psi_d in (30, 90, 150):
    dA = 1e-4; ps_ = psi_d * D
    Lv = np.array([1.0, 0, 0]); n1 = np.array([0, 1.0, 0]); n2 = np.array([0, np.cos(dA), np.sin(dA)])
    P1 = np.cos(ps_) * Lv + np.sin(ps_) * n1; P2 = np.cos(ps_) * Lv + np.sin(ps_) * n2
    dC = np.arccos(np.clip(P1 @ P2, -1, 1))
    check(f"Delta C / Delta A at psi = {psi_d}: vs sin(psi)", dC / dA, np.sin(ps_), 1e-4)

# ------------------------------------------------------------------ Part 4: the rotating Earth
print("\nPart 4: rotating Earth")
phi_s = 30 * D; V0 = site_speed(phi_s)
check("surface speed at the equator, omega R [km/s]", WE * RE, 0.46510, 5e-6)
check("  at 30 N [km/s]", V0, 0.40279, 5e-6)
beta, ve, ge = 135 * D, 16000 * 0.3048 / 1000, 30 * D
S = -ve * np.cos(ge) * np.cos(beta); E = ve * np.cos(ge) * np.sin(beta) + V0; Z = ve * np.sin(ge)
check("Ex 6: v_rel [km/s]", ve, 4.877, 5e-4); check("  true (inertial) speed [km/s]", np.sqrt(S * S + E * E + Z * Z), 5.1335, 5e-4)
check("  true flight-path angle [deg]", np.degrees(np.arcsin(Z / np.sqrt(S * S + E * E + Z * Z))), 28.4, 0.05)
check("  true azimuth [deg] (the relative one is 135)", np.degrees(np.arctan2(E, -S)) % 360, 131.39, 0.005)
# the minimum-speed example
phi0, lam0, phit, lamt = 60 * D, 160 * D, 0.0, 245 * D
dlam = lamt - lam0


def Lam_of(tA):
    return np.arccos(np.sin(phi0) * np.sin(phit) + np.cos(phi0) * np.cos(phit) * np.cos(dlam + WE * tA))


def t_min(Lam):
    Q = Qmin(Lam); return tff(RE, np.sqrt(Q * MU / RE), (np.pi - Lam) / 4)


Lam_nr = Lam_of(0.0)
check("not rotating: Lambda [deg]", Lam_nr / D, 87.50, 0.005); check("not rotating: t_ff [min]", t_min(Lam_nr) / 60, 31.71, 0.005)
tA = 1800.0; hist = []
for _ in range(10):
    tA = t_min(Lam_of(tA)); hist.append(tA / 60)
print("  iterates [min]:", [round(h_, 4) for h_ in hist])
Lam7 = Lam_of(tA); A7 = launch_azimuth(phi0, phit, Lam7); Q7m = Qmin(Lam7); v7m = np.sqrt(Q7m * MU / RE); g7m = (np.pi - Lam7) / 4
check("rotating: t_A [min]", tA / 60, 32.59, 0.005); check("  Bate's answer for exercise 6.26 [min]", tA / 60, 32.57, 0.03)
check("  Lambda [deg]", Lam7 / D, 91.58, 0.005); check("  omega t_A [deg]", WE * tA / D, 8.17, 0.005)
check("  launch azimuth A [deg]", A7 / D, 87.25, 0.005); check("  Q_min", Q7m, 0.8351, 5e-5); check("  inertial speed [km/s]", v7m, 7.224, 5e-4)
check("  inertial flight-path angle [deg]", g7m / D, 22.10, 0.005)
# velocity seen from the launch site
Nn, Ee, Zz = v7m * np.cos(g7m) * np.cos(A7), v7m * np.cos(g7m) * np.sin(A7), v7m * np.sin(g7m)
Vsite = site_speed(phi0)
Er = Ee - Vsite
vrel = np.sqrt(Nn ** 2 + Er ** 2 + Zz ** 2)
check("  site speed at 60 N [km/s]", Vsite, 0.23255, 5e-6)
check("  speed relative to the site [km/s]", vrel, 7.009, 5e-4); check("  elevation relative to the site [deg]", np.degrees(np.arcsin(Zz / vrel)), 22.82, 0.005); check("  azimuth relative to the site [deg]", np.degrees(np.arctan2(Er, Nn)), 87.16, 0.005)
# Bate's own check of the true speed
check("ground range of the not-rotating case [km]", Lam_nr * RE, 9740.7, 0.05)
# ------------------------------------------------------------------ Part 5: the Try-it boxes, limits and the exercises
print("\nPart 5: Try-it boxes, limits, exercises")
check("Try it: psi at Q = 0.8, gamma = 9.295 deg [deg]", psi_of(0.8, 9.295 * D) / D, 60.0, 0.005)
check("Try it: psi at Q = 0.8, gamma = 50.705 deg [deg]", psi_of(0.8, 50.705 * D) / D, 60.0, 0.005)
rT, vT = 6578.137, 7.3; QT = Q_of(rT, vT)
check("Try it: a = r/(2 - Q) for r = 6578.137 km, v = 7.3 km/s [km]", rT / (2 - QT), 5870, 0.5)
check("Try it: e at gamma = 30 deg for Q = 0.8794", ecc(QT, 30 * D), 0.5108, 5e-5)
# flat-Earth limits: range and flight time
for gdeg in (20.0, 45.0, 70.0):
    rr, vv = RE, 0.05                                                    # v << v_c: a short hop
    gg = MU / rr ** 2
    check(f"flat-Earth limit, gamma = {gdeg}: range / (v^2 sin 2 gamma/g)", psi_of(Q_of(rr, vv), gdeg * D) * rr / (vv ** 2 * np.sin(2 * gdeg * D) / gg), 1.0, 2e-3)
    check(f"  flight time / (2 v sin gamma / g)", tff(rr, vv, gdeg * D) / (2 * vv * np.sin(gdeg * D) / gg), 1.0, 2e-3)
# forward-azimuth formula by atan2 agrees with the arccos form
dl_eff = dlam + WE * tA
A_atan = np.arctan2(np.sin(dl_eff) * np.cos(phit), np.cos(phi0) * np.sin(phit) - np.sin(phi0) * np.cos(phit) * np.cos(dl_eff))
check("azimuth by atan2 [deg]", A_atan / D, A7 / D, 1e-9)
# exercises
check("Ex 1: Q for r = 1.1 R, v = 0.5 v_c", 0.25 * 1.1, 0.275, 1e-12)
check("Ex 1: psi_max [deg]", psi_max(0.275) / D, 18.35, 0.005); check("Ex 1: ground range [km]", psi_max(0.275) * RE, 2042.3, 0.05)
psiE2 = 9334.08 / RE
check("Ex 2: psi for 9334 km [deg]", psiE2 / D, 83.85, 0.005); check("Ex 2: v_bo at r = R [km/s]", np.sqrt(Qmin(psiE2) * MU / RE), 7.0755, 5e-5)
check("Ex 2: Q_min", Qmin(psiE2), 0.8011, 5e-5)
QE5 = 0.83 ** 2 * 1.06
check("Ex 5: Q = 0.83^2 * 1.06", QE5, 0.7302, 5e-5); check("Ex 5: psi_max [deg]", psi_max(QE5) / D, 70.21, 0.005); check("Ex 5: ground range [km]", psi_max(QE5) * RE, 7816, 1.0)
psiE6 = 3333.6 / RE
check("Ex 6: psi [deg]", psiE6 / D, 29.95, 0.005); check("Ex 6: v_bo / v_c", np.sqrt(Qmin(psiE6)), 0.6408, 5e-5); check("Ex 6: v_bo [km/s]", np.sqrt(Qmin(psiE6)) * vc, 5.066, 5e-4)
rE4 = 21.8e6 * 0.3048 / 1000; psiE4 = (7400 - 140 - 60) * NM / RE
check("Ex 4: r_bo [km]", rE4, 6644.6, 0.05); check("Ex 4: psi [deg]", psiE4 / D, 119.78, 0.005)
vE4 = np.sqrt(Qmin(psiE4) * MU / rE4)
check("Ex 4: v_bo [km/s]", vE4, 7.4597, 5e-4); check("Ex 4: gamma [deg]", gamma_max_range(psiE4) / D, 15.05, 0.005)
check("Ex 4: t_ff [min]", tff(rE4, vE4, gamma_max_range(psiE4)) / 60, 40.11, 0.005)
check("Ex 8: eastward gain at 13.7 N [km/s]", site_speed(13.7 * D), 0.4519, 5e-5); check("Ex 8: at 28.5 N [km/s]", site_speed(28.5 * D), 0.4087, 5e-5)
check("Ex 9: delta C = 1.852 km as an angle [arcsec]", 1.852 / RE / D * 3600, 59.9, 0.05)

# the numbers of Examples 5, 2, 3, 6 (low branch) and 8 that the note prints
check("Ex 5: energy [km^2/s^2]", eps1, -38.19, 0.005); check("Ex 5: apogee speed [km/s]", va1, 2.635, 5e-4)
check("Ex 5: tan(psi/2)", np.tan(psi1 / 2), 0.3287, 5e-5); check("Ex 5: e", ecc(Q1, g1), 0.8333, 5e-5); check("Ex 5: a [km]", r1 / (2 - Q1), 5218, 0.5)
check("Ex 2: (2-Q)/Q", (2 - Q2) / Q2, 0.6676, 5e-5); check("Ex 2: (2-Q)/Q sin(psi/2)", (2 - Q2) / Q2 * np.sin(Lam2 / 2), 0.6018, 5e-5)
check("Ex 2: A [deg]", np.degrees(np.arcsin((2 - Q2) / Q2 * np.sin(Lam2 / 2))), 37.00, 0.005)
check("Ex 2: the other root [deg]", 0.5 * (np.degrees(np.arcsin((2 - Q2) / Q2 * np.sin(Lam2 / 2))) - Lam2 / D / 2), -13.67, 0.005)
check("Ex 3: Q/(2-Q) for Ex 3 (Example 3 of the note)", Q4 / (2 - Q4), 0.8932, 5e-5)
check("Ex 4 of the note: r [km]", r5, 7015.2, 0.05); check("  sin(psi/2)", np.sin(psi5 / 2), 0.9174, 5e-5)
check("check yourself: psi_max(0.98) [deg]", psi_max(0.98) / D, 147.8, 0.05)
check("Ex 6: psi_max at Q = 0.9009 [deg]", psi_max(Q6) / D, 110.11, 0.005); check("Ex 6: gamma_max [deg]", gamma_max_range(psi_max(Q6)) / D, 17.47, 0.005)
glow6 = gammas_of(Q6, psi6)[0]
check("Ex 6 low root [deg]", glow6 / D, 9.74, 0.005)
check("Ex 6 low: d psi/d gamma", dpsi_dgamma(psi6, glow6), 3.195, 5e-4); check("Ex 6 low: d psi/d r [1/km]", dpsi_dr(r6, v6, psi6, glow6), 1.1225e-3, 5e-7)
check("Ex 6 low: d psi/d v [s/km]", dpsi_dv(r6, v6, psi6, glow6), 2.2015, 5e-4)
check("Ex 6 low: contribution r [km]", dpsi_dr(r6, v6, psi6, glow6) * dr_ * RE, 22.83, 0.005); check("  v [km]", dpsi_dv(r6, v6, psi6, glow6) * dv_ * RE, -5.55, 0.005)
check("  gamma [km]", dpsi_dgamma(psi6, glow6) * dg_ * RE, -2.04, 0.005)
check("  sum [km]", (dpsi_dr(r6, v6, psi6, glow6) * dr_ + dpsi_dv(r6, v6, psi6, glow6) * dv_ + dpsi_dgamma(psi6, glow6) * dg_) * RE, 15.24, 0.005)
check("Ex 8: v_N [km/s]", Nn, 0.3207, 5e-5); check("Ex 8: v_E [km/s]", Ee, 6.6854, 5e-5); check("Ex 8: v_E relative [km/s]", Er, 6.4528, 5e-5); check("Ex 8: v_Z [km/s]", Zz, 2.7183, 5e-5)
check("Ex 8: extra range angle [km]", (Lam7 - Lam_nr) * RE, 454.4, 0.05); check("Ex 8: extra time [s]", tA - t_min(Lam_nr), 53.06, 0.005)
check("Ex 8: target runs away along the equator [km]", WE * tA * RE, 909.5, 0.05)
slope = (t_min(Lam_of(tA + 1.0)) - t_min(Lam_of(tA - 1.0))) / 2.0
check("Ex 8: slope of the iteration map", slope, 0.027, 5e-4)
check("Ex 8: iterates 1, 2, 3 [min]", hist[0], 32.522, 5e-4); check("   2", hist[1], 32.589, 5e-4); check("   3", hist[2], 32.5906, 5e-5)
check("check yourself: Earth turns in 20 min [deg]", WE * 1200 / D, 5.01, 0.005); check("  shift at 60 deg latitude [km]", WE * 1200 * np.cos(60 * D) * RE, 279, 1.0)
check("Ex 7: v_N", -ve * np.cos(ge) * np.cos(np.pi - beta), 0, 1e9) if False else None
check("Ex 7: v_N [km/s]", ve * np.cos(ge) * np.cos(beta), -2.986, 5e-4); check("Ex 7: v_E [km/s]", E, 3.389, 5e-4); check("Ex 7: v_Z [km/s]", Z, 2.438, 1e-3)
check("Ex 7: speed gain [km/s]", np.sqrt(S * S + E * E + Z * Z) - ve, 0.26, 5e-3)
check("Try it (Q section): perigee distance a(1 - e) at gamma = 30 deg [km]", (rT / (2 - QT)) * (1 - ecc(QT, 30 * D)), 2872, 0.5)
check("Try it (Q section): e^2", ecc(QT, 30 * D) ** 2, 0.2609, 5e-5)
check("Ex 2 of the exercises: gamma = 45 - psi/4 [deg]", gamma_max_range(psiE2) / D, 24.04, 0.005)

# ------------------------------------------------------------------ Part 6: does the flight come down?
print("\nPart 6: does the flight come down? (perigee against the surface)")


def two_body(r, Q, gam):
    """Right-hand side, initial state and period of the ellipse through the burnout point (r, Q, gamma)."""
    v = np.sqrt(Q * MU / r)
    y0 = [r, 0, v * np.sin(gam), v * np.cos(gam)]
    f = lambda t, 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]
    return f, y0, 2 * np.pi * np.sqrt((r / (2 - Q)) ** 3 / MU)


def min_radius(r, Q, gam):
    """Smallest distance from the focus over one period, by numerical integration."""
    f, y0, T = two_body(r, Q, gam)
    sol = solve_ivp(f, [0, T], y0, rtol=1e-12, atol=1e-9, dense_output=True)
    t = np.linspace(0, T, 400001)
    y = sol.sol(t)
    return np.hypot(y[0], y[1]).min()


def to_ground(r, Q, gam, R=RE):
    """(reaches radius R?, angle swept from burnout to the first descending crossing of R) by numerical integration."""
    f, y0, T = two_body(r, Q, gam)
    ev = lambda t, y: np.hypot(y[0], y[1]) - R
    ev.terminal = True; ev.direction = -1
    sol = solve_ivp(f, [0, 1.05 * T], y0, events=ev, rtol=1e-12, atol=1e-9, max_step=T / 2000)
    if len(sol.t_events[0]) == 0:
        return False, None
    y = sol.y_events[0][0]
    return True, np.arctan2(y[1], y[0]) % (2 * np.pi)


worst = 0.0
for Q_, g_, h_ in [(0.97, 10, 1000), (1.0, 45, 1000), (1.2, 3, 200), (1.2, 55, 1000), (1.5, 5, 200), (0.97, 3, 1000), (1.0, 2, 1000), (0.8, 30, 200)]:
    r_ = RE + h_
    worst = max(worst, abs(perigee_radius(Q_, g_ * D, r_) - min_radius(r_, Q_, g_ * D)))
check("perigee radius r(1 - e)/(2 - Q): formula minus integration, worst of 8 cases [km]", worst, 0.0, 1e-3)

# the threshold: bisection on lands() agrees with the closed form
worst = 0.0
for Q_ in (0.97, 1.0, 1.2, 1.5, 1.9):
    for h_ in (200, 500, 1000):
        r_ = RE + h_
        g_cf = gamma_land(Q_, r_)
        lo, hi = 1e-9, np.pi / 2 - 1e-9
        for _ in range(200):
            mid = 0.5 * (lo + hi)
            if lands(Q_, mid, r_): hi = mid
            else: lo = mid
        worst = max(worst, abs(g_cf - hi))
check("gamma_land: closed form minus bisection on lands(), worst of 15 cases [rad]", worst, 0.0, 1e-8)   # bisection stops at its lower bound 1e-9 when every gamma lands

# the integration just above and just below the threshold
ok = True
for Q_, h_ in [(0.97, 1000), (1.0, 200), (1.0, 1000), (1.2, 200), (1.2, 1000), (1.5, 500)]:
    r_ = RE + h_; gs = gamma_land(Q_, r_)
    up, _ = to_ground(r_, Q_, gs + 0.05 * D); dn, _ = to_ground(r_, Q_, gs - 0.05 * D)
    ok = ok and up and not dn
truth("integration: 0.05 deg above gamma* the vehicle reaches the ground, 0.05 deg below it does not (6 cases)", ok)

# the impact angle
worst = 0.0
for Q_, g_, h_ in [(0.97, 10, 1000), (0.97, 65, 1000), (1.0, 12, 1000), (1.0, 45, 1000), (1.2, 20, 1000), (1.2, 55, 1000), (1.2, 10, 200), (1.5, 45, 200)]:
    r_ = RE + h_
    hit, ang = to_ground(r_, Q_, g_ * D)
    worst = max(worst, abs(psi_impact(Q_, g_ * D, r_) - ang)) if hit else 9
check("psi_impact (burnout to ground): formula minus integration, worst of 8 cases [rad]", worst, 0.0, 1e-7)
check("  e.g. Q = 1.2, gamma = 10 deg, burnout 200 km: free flight psi [deg]", psi_of(1.2, 10 * D) / D, 257.20, 0.005)
check("  ... and the angle swept to the ground [deg]", psi_impact(1.2, 10 * D, RE + 200) / D, 268.32, 0.005)

# Q = 1: e = sin(gamma), so the threshold is sin(gamma*) = h/r_bo
check("Q = 1: gamma* = arcsin(h/r), h = 200 km [deg]", gamma_land(1.0, RE + 200) / D, np.degrees(np.arcsin(200 / (RE + 200))), 1e-9)
check("Q = 1: gamma* = arcsin(h/r), h = 1000 km [deg]", gamma_land(1.0, RE + 1000) / D, 7.79, 0.005)
truth("Q = 1: e = sin(gamma)", all(abs(ecc(1.0, g_) - np.sin(g_)) < 1e-12 for g_ in np.linspace(0.01, 1.5, 20)))

# the table of thresholds (note: every gamma lands when gamma_land is 0)
TBL = {0.97: (0.0, 0.0, 1.91, 6.05), 1.0: (0.0, 1.74, 4.17, 7.79), 1.2: (0.0, 5.95, 9.59, 13.87), 1.5: (0.0, 8.25, 12.95, 18.10)}
for Q_, row in TBL.items():
    for h_, stated in zip((0, 200, 500, 1000), row):
        check(f"table: gamma* at Q = {Q_}, burnout {h_} km [deg]", gamma_land(Q_, RE + h_) / D, stated, 0.005)

# when can a flight miss at all?  Q < 1: Q > 2R/(r + R);  1 <= Q < 2: whenever the burnout is above the ground
for h_, stated in zip((200, 500, 1000), (0.9846, 0.9623, 0.9273)):
    check(f"Q < 1: the smallest Q that can stay up, 2R/(r + R), burnout {h_} km", 2 * RE / (2 * RE + h_), stated, 5e-5)
agree = True
for h_ in (200, 500, 1000):
    r_ = RE + h_
    for Q_ in np.linspace(0.90, 1.0, 41)[:-1]:
        stays_up_flat = not lands(Q_, 0.0, r_)           # horizontal burnout is the easiest case to miss
        agree = agree and (stays_up_flat == (Q_ > 2 * RE / (r_ + RE)))
    for Q_ in np.linspace(1.0, 1.95, 20):
        agree = agree and (not lands(Q_, 0.0, r_))        # Q >= 1: a horizontal burnout above the ground never lands
truth("a flat burnout stays up exactly when Q > 2R/(r + R) (Q < 1) and always for Q >= 1 (h > 0)", agree)
truth("burnout at the surface: every gamma > 0 lands (Q = 0.5 ... 1.9, gamma = 0.01 ... 89 deg)",
      all(lands(Q_, g_ * D, RE) for Q_ in np.linspace(0.5, 1.9, 15) for g_ in (0.01, 1, 10, 30, 60, 89)))

# the largest range that still lands for Q >= 1 (psi falls as gamma rises, so the landing flights are those with smaller psi)
check("Q = 1.2, burnout 200 km: gamma* [deg]", gamma_land(1.2, RE + 200) / D, 5.95, 0.005)
check("Q = 1.2, burnout 200 km: largest range that lands [deg]", psi_of(1.2, gamma_land(1.2, RE + 200)) / D, 293.04, 0.005)
check("Q = 1, burnout 200 km: largest range that lands [deg]", psi_of(1.0, gamma_land(1.0, RE + 200)) / D, 176.52, 0.005)

# e depends on cos^2(gamma) only, so a negative gamma has the same perigee: gamma < 0 is rejected because the path starts below r_bo, not because it lands
truth("r_p(Q, -gamma) = r_p(Q, gamma): the sign of gamma does not decide whether the ellipse dips below the surface",
      all(abs(perigee_radius(Q_, -g_ * D, RE + 1000) - perigee_radius(Q_, g_ * D, RE + 1000)) < 1e-9 for Q_ in (0.97, 1.0, 1.2) for g_ in (3, 10, 40)))
check("Q = 1.2, gamma = -3 deg, burnout 1000 km: r_p [km] (stays up, as for +3 deg)", perigee_radius(1.2, -3 * D, RE + 1000), 7318.5, 0.05)

# the figures: perigee radii and landing flags printed in the legends of fig-bal-land-*.png
FIG = [(0.97, 1, 6914.7, False), (0.97, 3, 6731.3, False), (0.97, 10, 5901.5, True), (0.97, 25, 4129.7, True), (0.97, 45, 2095.8, True), (0.97, 65, 670.5, True),
       (1.0, 2, 7120.6, False), (1.0, 5, 6735.1, False), (1.0, 12, 5844.1, True), (1.0, 25, 4260.0, True), (1.0, 45, 2161.0, True), (1.0, 70, 445.0, True),
       (1.2, 3, 7318.5, False), (1.2, 9, 6898.8, False), (1.2, 20, 5623.5, True), (1.2, 35, 3721.2, True), (1.2, 55, 1594.2, True)]
okp = okl = True
for Q_, g_, rp_, hit_ in FIG:
    okp = okp and abs(perigee_radius(Q_, g_ * D, RE + 1000) - rp_) < 0.05
    okl = okl and (lands(Q_, g_ * D, RE + 1000) == hit_)
truth("figures: the 17 perigee radii of the legends (burnout 1000 km)", okp)
truth("figures: the 17 landing flags of the legends agree with lands() and are on the right side of gamma*", okl)


# ------------------------------------------------------------------ Part 7: the example of the range equation (psi below and above 180 deg)
print("\nPart 7: the range-equation example, Q = 1.2")
Qx = 1.2
check("example: v_bo for Q = 1.2 at r = 6578.137 km [km/s]", np.sqrt(Qx * MU / 6578.137), 8.527, 5e-4)
for gx, num_s, den_s, tan_s, half_s, psi_s in ((30, 0.5196, 0.1, 5.196, 79.11, 158.21), (10, 0.2052, -0.1638, -1.2527, 128.60, 257.20)):
    cx, sx = np.cos(gx * D), np.sin(gx * D)
    check(f"example gamma = {gx}: numerator Q sin cos", Qx * sx * cx, num_s, 5e-5)
    check(f"example gamma = {gx}: denominator 1 - Q cos^2", 1 - Qx * cx ** 2, den_s, 5e-5)
    check(f"example gamma = {gx}: tan(psi/2)", Qx * sx * cx / (1 - Qx * cx ** 2), tan_s, 5e-4 if gx == 30 else 5e-5)
    check(f"example gamma = {gx}: psi/2 [deg]", psi_of(Qx, gx * D) / 2 / D, half_s, 0.005)
    check(f"example gamma = {gx}: psi [deg]", psi_of(Qx, gx * D) / D, psi_s, 0.005)
    ig, _ = integrate(6578.137, np.sqrt(Qx * MU / 6578.137), gx * D)
    check(f"example gamma = {gx}: psi from the numerical integration [deg]", ig / D, psi_s, 0.005)
check("example: sin10 cos10", np.sin(10 * D) * np.cos(10 * D), 0.1710, 5e-5); check("example: cos^2 10", np.cos(10 * D) ** 2, 0.9698, 5e-5)
check("example: sin30 cos30", np.sin(30 * D) * np.cos(30 * D), 0.4330, 5e-5)
check("example: plain arctan of -1.2527 [deg]", np.degrees(np.arctan(-1.2527)), -51.40, 0.005)
check("example: 180 - 51.40 [deg]", 180 - 51.40, 128.60, 1e-9)
check("example: e cos(psi/2) = 1 - Q cos^2 gamma at gamma = 10 deg", ecc(Qx, 10 * D) * np.cos(psi_of(Qx, 10 * D) / 2), 1 - Qx * np.cos(10 * D) ** 2, 1e-12)
gx0 = np.arccos(1 / np.sqrt(Qx))
check("example: gamma where psi = 180 deg, arccos(1/sqrt Q) [deg]", gx0 / D, 24.09, 0.005)
check("example: psi there [deg]", psi_of(Qx, gx0) / D, 180.0, 1e-6)
truth("example: Q cos^2(gamma) <= 1 for every Q <= 1 and gamma, so psi <= 180 deg for Q <= 1 (grid)",
      all(Qq * np.cos(gg_) ** 2 <= 1 + 1e-15 and psi_of(Qq, gg_) <= np.pi + 1e-12 for Qq in np.linspace(0.05, 1.0, 20) for gg_ in np.linspace(0.01, 1.55, 30)))

print("\n%d checks, %d failed" % (len(RESULTS), RESULTS.count(False)))
