"""Figures for ballistic-trajectories-bates.qmd (img/fig-bal-geometry.png and img/fig-bal-focal.png).

fig-bal-geometry: the geometry of a symmetric free flight: burnout and re-entry at the same radius, the arc through apogee, the range angle psi.
fig-bal-focal   : the focal triangle F-P-F' behind the flight-path-angle equation: the normal at the burnout point bisects the angle between the focal radii.
Q = 0.8 and psi = 60 deg, high trajectory (gamma = 50.7 deg).

Run from the sfm-slides directory:   python3 codes/fig-bal-geometry.py
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from matplotlib.patches import Arc, Circle
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _ballistic import gammas_of, ecc, RE

Q, psi = 0.8, np.radians(60.0)
glo, ghi = gammas_of(Q, psi)
g = ghi
e = ecc(Q, g)
r = 1.0                       # work in units of r_bo for the drawing only
a = r / (2 - Q)
p = a * (1 - e * e)
Rearth = 0.92 * r             # the Earth is a little smaller than the burnout radius, so the burnout point is above the ground

# apogee along +y; the point with true anomaly f sits at polar angle (pi - f) from +y
def pt(f):
    rr = p / (1 + e * np.cos(f))
    th = np.pi - f
    return np.array([rr * np.sin(th), rr * np.cos(th)])

fbo = np.pi - psi / 2
P = pt(fbo)                                     # burnout (right)
Pre = np.array([-P[0], P[1]])                   # re-entry (left)
F = np.zeros(2)
Fp = np.array([0.0, 2 * a * e])                 # the empty focus lies on the apogee side
f_all = np.linspace(0, 2 * np.pi, 800)
ell = np.array([pt(f) for f in f_all])
arc_f = np.linspace(fbo, 2 * np.pi - fbo, 300)
arc = np.array([pt(f) for f in arc_f])

# ---------------- symmetric free flight
figA, A = plt.subplots(figsize=(6.6, 6.0))
A.add_patch(Circle((0, 0), Rearth, color="#dfe9f3", ec=BLUE, lw=1.4, zorder=1))
A.plot(ell[:, 0], ell[:, 1], color=INK, lw=0.8, ls=(0, (4, 3)), zorder=2)
A.plot(arc[:, 0], arc[:, 1], color=MAROON, lw=2.8, zorder=4)
A.scatter([0], [0], color=INK, s=22, zorder=6); A.text(-0.30, -0.13, "$F$", fontsize=12); A.text(-1.42, -0.55, "Earth", color=BLUE, fontsize=11)
for X, lab, dx in ((P, "burnout", 0.05), (Pre, "re-entry", -0.05)):
    A.scatter(*X, color=AMBER, s=48, zorder=7)
    A.plot([0, X[0]], [0, X[1]], color=AMBER, lw=1.2, zorder=3)
    A.text(X[0] + dx, X[1] + 0.03, lab, color=AMBER, fontsize=11, ha="left" if dx > 0 else "right")
apo = pt(np.pi)
A.scatter(*apo, color=PURPLE, s=40, zorder=7); A.text(apo[0] + 0.04, apo[1] + 0.02, "apogee", color=PURPLE, fontsize=11)
A.plot([0, 0], [0, apo[1]], color=INK, lw=0.9, ls=":", zorder=2)
A.text(0.02, apo[1] * 0.55, "major axis", color=INK, fontsize=10, rotation=90, va="center")
# range angle psi and its halves
A.add_patch(Arc((0, 0), 0.5, 0.5, theta1=90 - np.degrees(psi / 2), theta2=90 + np.degrees(psi / 2), color=MAROON, lw=1.8))
A.text(-0.02, 0.33, r"$\psi$", color=MAROON, fontsize=14, ha="center")
A.add_patch(Arc((0, 0), 0.72, 0.72, theta1=90 - np.degrees(psi / 2), theta2=90, color=BLUE, lw=1.4, ls="--"))
A.text(0.17, 0.46, r"$\psi/2$", color=BLUE, fontsize=11)
# true anomaly of burnout: measured from the perigee direction (down)
A.add_patch(Arc((0, 0), 0.62, 0.62, theta1=-90, theta2=-90 + np.degrees(fbo), color=GREEN, lw=1.4, ls="--"))
A.text(0.36, -0.30, r"$f_{bo}=180^\circ-\psi/2$", color=GREEN, fontsize=10.5)
A.plot([0, 0], [0, -0.8], color=GREEN, lw=1.1, ls=":", zorder=2)
A.text(0.04, -0.78, "to perigee", color=GREEN, fontsize=10)
A.text(-1.5, 1.72, "symmetric free flight", fontsize=12, color=INK, weight="bold")
A.set_aspect("equal"); A.set_xlim(-1.55, 1.55); A.set_ylim(-0.95, 1.85); A.axis("off")
figA.tight_layout()
figA.savefig("img/fig-bal-geometry.png", dpi=150, bbox_inches="tight")
# ---------------- the focal triangle
figB, B = plt.subplots(figsize=(7.4, 6.0))
Bc = np.array([0.0, a * e])
B.plot(ell[:, 0], ell[:, 1], color=INK, lw=1.0, zorder=2)
B.plot(arc[:, 0], arc[:, 1], color=MAROON, lw=2.8, zorder=3)
B.scatter(*F, color=INK, s=28, zorder=6); B.scatter(*Fp, color=INK, s=28, zorder=6)
B.text(-0.12, -0.12, "$F$", fontsize=12); B.text(0.06, Fp[1] - 0.02, "$F'$", fontsize=12)
B.scatter(*P, color=AMBER, s=48, zorder=7); B.text(P[0] + 0.07, P[1] - 0.12, "$P$", fontsize=12, color=AMBER)
B.plot([F[0], P[0]], [F[1], P[1]], color=BLUE, lw=1.6, zorder=4); B.plot([Fp[0], P[0]], [Fp[1], P[1]], color=GREEN, lw=1.6, zorder=4)
B.plot([F[0], Fp[0]], [F[1], Fp[1]], color=INK, lw=1.0, ls=":", zorder=3)
rp = np.hypot(*(P - F)); rpp = np.hypot(*(P - Fp))
B.text(P[0] * 0.38 + 0.06, P[1] * 0.4 - 0.1, r"$r_{bo}=a(2-Q)$", color=BLUE, fontsize=10.5, rotation=np.degrees(np.arctan2(P[1], P[0])))
mid = 0.5 * (P + Fp)
B.text(mid[0] + 0.02, mid[1] + 0.05, r"$r'=aQ$", color=GREEN, fontsize=10.5, rotation=np.degrees(np.arctan2(P[1] - Fp[1], P[0] - Fp[0])))
# velocity, local horizontal, normal at P
rh = P / np.linalg.norm(P)
hor = np.array([rh[1], -rh[0]])                      # local horizontal, pointing along increasing f (right-handed about +z...) -> check sign below
# the velocity direction is the tangent in the direction of motion (towards apogee, i.e. leftwards-up here)
tvec = arc[np.argmin(np.hypot(arc[:, 0] - P[0], arc[:, 1] - P[1])) + 3] - P
tvec /= np.linalg.norm(tvec)
nrm = np.array([-tvec[1], tvec[0]]); nrm = nrm if nrm @ rh > 0 else -nrm
L = 0.55
B.annotate("", xy=P + 0.62 * tvec, xytext=P, arrowprops=dict(arrowstyle="-|>", color=MAROON, lw=2.2), zorder=8)
B.text(P[0] + 0.62 * tvec[0] - 0.02, P[1] + 0.62 * tvec[1] + 0.04, r"$\mathbf{v}_{bo}$", color=MAROON, fontsize=12)
hdir = hor if hor @ tvec > 0 else -hor
B.plot([P[0] - 0.5 * hdir[0], P[0] + 0.5 * hdir[0]], [P[1] - 0.5 * hdir[1], P[1] + 0.5 * hdir[1]], color=INK, lw=1.0, ls="--", zorder=4)
B.text(P[0] - 0.5 * hdir[0] + 0.02, P[1] - 0.5 * hdir[1] - 0.1, "local horizontal", fontsize=9.5, color=INK, va="top")
B.plot([P[0] - L * nrm[0], P[0] + L * nrm[0]], [P[1] - L * nrm[1], P[1] + L * nrm[1]], color=PURPLE, lw=1.6, ls="-.", zorder=5)
B.text(P[0] + L * nrm[0] + 0.02, P[1] + L * nrm[1] + 0.02, "normal", color=PURPLE, fontsize=10.5)
# angle marks: angle between r_hat and normal = gamma; angle FPF' = 2 gamma; angle at F = psi/2
def ang(v): return np.degrees(np.arctan2(v[1], v[0]))
a_r, a_n, a_Fp = ang(-(P - F)) , ang(nrm), ang(Fp - P)
B.add_patch(Arc(P, 0.42, 0.42, theta1=ang(Fp - P), theta2=ang(F - P) % 360, color=PURPLE, lw=1.6))
B.text(P[0] - 0.33, P[1] - 0.05, r"$2\gamma$", color=PURPLE, fontsize=12)
B.add_patch(Arc(F, 0.5, 0.5, theta1=ang(P - F), theta2=90, color=BLUE, lw=1.6))
B.text(0.07, 0.31, r"$\psi/2$", color=BLUE, fontsize=11.5)
B.add_patch(Arc(Fp, 0.5, 0.5, theta1=270, theta2=ang(P - Fp) % 360, color=GREEN, lw=1.6))
B.annotate("$180^\\circ-2\\gamma-\\psi/2$", xy=(0.14, Fp[1] - 0.2), xytext=(-1.72, Fp[1] - 0.45), color=GREEN, fontsize=10, arrowprops=dict(arrowstyle="-", color=GREEN, lw=0.9))
B.text(-1.72, 1.78, "the focal triangle", fontsize=12, color=INK, weight="bold")
B.set_aspect("equal"); B.set_xlim(-1.75, 1.55); B.set_ylim(-0.55, 1.9); B.axis("off")
figB.tight_layout()
figB.savefig("img/fig-bal-focal.png", dpi=150, bbox_inches="tight")
print(f"Q={Q}, psi=60 deg: gamma_low={np.degrees(glo):.2f}, gamma_high={np.degrees(ghi):.2f}, e={e:.4f}, a/r={a:.4f}, "
      f"|PF|={rp:.4f} (a(2-Q)={a*(2-Q):.4f}), |PF'|={rpp:.4f} (aQ={a*Q:.4f})")