"""Figures for ballistic-family.qmd: the one-parameter family of ellipses through two points at the same radius.

fig-bal-fam-sketch : several ellipses of the family, with their empty foci on the ray towards the apogee (no numbers).
fig-bal-fam-circle : the energy fixes the distance of the empty focus from the burnout point, a circle that meets the ray 0, 1 or 2 times (no numbers).
fig-bal-fam-links  : burnout speed, apogee height and flight time against gamma, for burnout at 200 km and psi = 60 deg.
fig-bal-fam-arcs   : four members of the same family to scale, with the descent from the re-entry radius to the ground.

Run from the sfm-slides directory:   python3 codes/fig-bal-family.py
"""
import numpy as np
from numpy import sin, cos, sqrt, pi, radians as rad, degrees as deg
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 MU, RE

PSI = rad(60.0)
S, C = sin(PSI / 2), cos(PSI / 2)
COL = [BLUE, GREEN, AMBER, MAROON, PURPLE]


def curve(r, psi, e, f0=0.0, f1=2 * pi, n=600):
    """Points of the ellipse with p = r(1 - e cos(psi/2)), apogee on +y, F at the origin. f is the true anomaly."""
    p = r * (1 - e * cos(psi / 2))
    f = np.linspace(f0, f1, n)
    rr = p / (1 + e * np.cos(f))
    th = pi - f
    return rr * np.sin(th), rr * np.cos(th)


def arc_of(r, psi, e):
    return curve(r, psi, e, pi - psi / 2, pi + psi / 2)


# ---------------- the family (no numbers)
fig, ax = plt.subplots(figsize=(7.2, 7.4))
r = 1.0
ax.add_patch(Circle((0, 0), 0.92 * r, color="#dfe9f3", ec=BLUE, lw=1.0, zorder=1))
ax.text(0.0, -0.45, "Earth", color=BLUE, ha="center", fontsize=11)
P = np.array([r * S, r * C]); Q = np.array([-r * S, r * C])
for e, col in zip((0.0, 0.3, 0.55, 0.72, 0.82), COL):
    x, y = curve(r, PSI, e); ax.plot(x, y, color=col, lw=0.7, ls=(0, (2, 3)), zorder=2)
    x, y = arc_of(r, PSI, e); ax.plot(x, y, color=col, lw=2.6, zorder=4)
    p = r * (1 - e * C); a = p / (1 - e * e)
    if e > 0:
        ax.scatter([0], [2 * a * e], color=col, s=42, zorder=6, ec="white", lw=0.6)
ax.plot([0, 0], [0, 2.05], color=INK, lw=0.9, ls=":", zorder=3)
ax.annotate("", xy=(0, 2.12), xytext=(0, 1.95), arrowprops=dict(arrowstyle="->", color=INK))
ax.text(0.05, 1.98, "towards the apogee", fontsize=11, va="center")
ax.annotate("empty foci $F'$", xy=(0.02, 1.451), xytext=(0.62, 1.60), fontsize=11, va="center", arrowprops=dict(arrowstyle="-", color=INK, lw=0.7))
ax.scatter([0], [0], color=INK, s=26, zorder=7); ax.text(0.04, -0.13, "$F$", fontsize=13)
for X, lab, ha, dx in ((P, "burnout $P$", "left", 0.05), (Q, "re-entry", "right", -0.05)):
    ax.plot([0, X[0]], [0, X[1]], color=INK, lw=1.0, zorder=3)
    ax.scatter(*X, color=INK, s=48, zorder=8)
    ax.text(X[0] + dx, X[1] - 0.05, lab, ha=ha, fontsize=12)
ax.add_patch(Arc((0, 0), 0.72, 0.72, theta1=90 - 30, theta2=90 + 30, color=INK, lw=1.2))
ax.text(-0.10, 0.39, r"$\psi$", fontsize=14, ha="center")
ax.set_aspect("equal"); ax.set_xlim(-1.45, 1.45); ax.set_ylim(-1.15, 2.2); ax.axis("off")
fig.savefig("img/fig-bal-fam-sketch.png", dpi=150, bbox_inches="tight"); plt.close(fig)

# ---------------- the circle about P (no numbers)
fig, axs = plt.subplots(1, 4, figsize=(15.5, 6.9))
cases = ((0.68, "no crossing\nno ellipse"), (0.75, "the circle touches\none ellipse"), (0.90, "two crossings\ntwo ellipses"), (1.25, "one crossing on the ray\none arc re-enters"))
for ax, (a, tt) in zip(axs, cases):
    rad_c = 2 * a - r
    ax.add_patch(Circle((0, 0), 0.92 * r, color="#dfe9f3", ec=BLUE, lw=1.0, zorder=1))
    ax.text(-0.42, -0.70 if a > r else -0.5, "Earth", color=BLUE, fontsize=10, ha="center")
    ax.plot([0, 0], [0, 2.75], color=INK, lw=1.3, zorder=3)
    ax.plot([0, 0], [-1.65, 0], color=INK, lw=1.0, ls=(0, (2, 3)), zorder=3)
    th = np.linspace(0, 2 * pi, 500)
    ax.plot(P[0] + rad_c * np.cos(th), P[1] + rad_c * np.sin(th), color=AMBER, lw=1.7, zorder=4)
    disc = rad_c ** 2 - (r * S) ** 2
    ys = [] if disc < -1e-9 else sorted({round(r * C - sqrt(max(disc, 0)), 9), round(r * C + sqrt(max(disc, 0)), 9)})
    n = 0
    for y in ys:
        if y >= 0:
            e = y / (2 * a)
            cc = COL[n]
            x_, y_ = curve(r, PSI, e); ax.plot(x_, y_, color=cc, lw=1.3, ls=(0, (4, 2.5)), zorder=2)
            x_, y_ = arc_of(r, PSI, e); ax.plot(x_, y_, color=cc, lw=2.8, zorder=5)
            ax.scatter([0], [y], color=cc, s=56, zorder=7, ec="white", lw=0.6)
            n += 1
        else:
            # the crossing behind F: the negative root e, whose apogee points the other way. Its arc above the chord passes through the perigee.
            x_, y_ = curve(r, PSI, y / (2 * a)); ax.plot(x_, y_, color="#8a8378", lw=1.5, zorder=2)
            ax.scatter([0], [y], facecolor="white", edgecolor=INK, s=56, zorder=7)
    ax.scatter(*P, color=INK, s=44, zorder=8); ax.scatter(*Q, color=INK, s=44, zorder=8)
    ax.scatter([0], [0], color=INK, s=26, zorder=8)
    ax.text(0.07, -0.1, "$F$", fontsize=12); ax.text(P[0] + 0.05, P[1] - 0.1, "$P$", fontsize=12)
    ax.set_title(tt, fontsize=11)
    ax.set_aspect("equal"); ax.set_xlim(-1.35, 2.1); ax.set_ylim(-1.7, 2.8); ax.axis("off")
fig.tight_layout()
fig.savefig("img/fig-bal-fam-circle.png", dpi=150, bbox_inches="tight"); plt.close(fig)

# ---------------- the links, burnout at 200 km and psi = 60 deg
rb = RE + 200.0
top = pi / 2 - PSI / 4
g = np.linspace(rad(0.05), top - rad(0.3), 900)
def v_g(g): return sqrt(MU / rb * S / (cos(g) * sin(g + PSI / 2)))
def alt_g(g): return rb * cos(g) * cos(PSI / 4) / cos(g + PSI / 4) - RE
def t_g(g):
    e = sin(g) / sin(g + PSI / 2)
    a = rb * (1 - e * C) / (1 - e * e)
    E = np.arccos((1 - rb / a) / e)
    return 2 * sqrt(a ** 3 / MU) * (pi - E + e * sin(E)) / 60
marks = rad(np.array([9.2952, 22.0312, 30.0, 36.59, 50.7048]))
fig, axs = plt.subplots(1, 3, figsize=(15.5, 4.8))
axs[0].plot(deg(g), v_g(g), color=MAROON, lw=2.2)
axs[0].axhline(sqrt(0.8 * MU / rb), color=INK, lw=0.9, ls=(0, (4, 3)))
axs[0].text(30, sqrt(0.8 * MU / rb) + 0.06, "a chosen speed", fontsize=10, ha="center")
axs[0].axhline(sqrt(MU / rb), color=BLUE, lw=0.9, ls=":")
axs[0].text(30, sqrt(MU / rb) + 0.06, "circular speed", color=BLUE, fontsize=10, ha="center")
axs[0].scatter(deg(rad(np.array([9.2952, 50.7048]))), [sqrt(0.8 * MU / rb)] * 2, color=INK, s=34, zorder=5)
axs[0].scatter([30], [v_g(rad(30.0))], color=AMBER, s=48, zorder=5)
axs[0].set_ylim(6.0, 10.5); axs[0].set_ylabel("burnout speed $v$ (km/s)")
axs[1].semilogy(deg(g), alt_g(g), color=GREEN, lw=2.2)
axs[1].scatter(deg(marks), alt_g(marks), color=INK, s=30, zorder=5)
axs[1].set_ylabel("apogee altitude (km)")
axs[2].semilogy(deg(g), t_g(g), color=PURPLE, lw=2.2)
axs[2].axhline(PSI * sqrt(rb ** 3 / MU) / 60, color=INK, lw=0.9, ls=(0, (4, 3)))
axs[2].text(40, PSI * sqrt(rb ** 3 / MU) / 60 * 0.8, "circular-arc floor", fontsize=10, ha="center", va="top")
axs[2].scatter(deg(marks), t_g(marks), color=INK, s=30, zorder=5)
axs[2].set_ylim(6.5, 3000); axs[2].set_ylabel(r"flight time $\Delta t$ (min)")
for ax, lab in zip(axs, "abc"):
    ax.axvline(deg(top), color=INK, lw=0.8, ls=":")
    ax.set_xlim(0, 78); ax.set_xlabel(r"flight-path angle at burnout $\gamma$ (deg)"); despine(ax)
    ax.text(0.02, 0.98, f"({lab})", transform=ax.transAxes, va="top", fontsize=12)
fig.tight_layout()
fig.savefig("img/fig-bal-fam-links.png", dpi=150, bbox_inches="tight"); plt.close(fig)

# ---------------- to scale
fig, ax = plt.subplots(figsize=(8.5, 7.4))
ax.add_patch(Circle((0, 0), RE, color="#dfe9f3", ec=BLUE, lw=1.0, zorder=1))
th = np.linspace(0, 2 * pi, 400); ax.plot(rb * np.cos(th), rb * np.sin(th), color=INK, lw=0.6, ls=":", zorder=2)
rows = ((0.14627, BLUE, "low, 7.30 km/s"), (0.25504, GREEN, "low, 6.96 km/s"), (sin(rad(30)) / sin(rad(60)), AMBER, "least energy, 6.36 km/s"), (0.78419, MAROON, "high, 6.96 km/s"))
for e, col, lab in rows:
    x, y = arc_of(rb, PSI, e); ax.plot(x, y, color=col, lw=2.4, zorder=5, label=lab)
    # continue from re-entry down to the surface
    p = rb * (1 - e * C)
    fg = np.linspace(pi + PSI / 2, 2 * pi, 800); rr = p / (1 + e * np.cos(fg)); keep = rr >= RE
    xk, yk = rr[keep] * np.sin(pi - fg[keep]), rr[keep] * np.cos(pi - fg[keep])
    ax.plot(xk, yk, color=col, lw=1.2, ls=(0, (3, 2)), zorder=4)
    ax.scatter([xk[-1]], [yk[-1]], marker="x", color=col, s=46, lw=1.6, zorder=6)
ax.scatter([rb * S, -rb * S], [rb * C, rb * C], color=INK, s=38, zorder=7)
ax.text(rb * S + 200, rb * C + 120, "burnout", fontsize=11); ax.text(-rb * S - 200, rb * C + 120, "re-entry", fontsize=11, ha="right")
ax.text(0, 0, "Earth, to scale", color=BLUE, fontsize=11, ha="center")
ax.legend(frameon=False, fontsize=10, loc="upper left", bbox_to_anchor=(0.60, 1.0))
ax.set_aspect("equal"); ax.set_xlim(-7400, 11500); ax.set_ylim(-1800, 10300); ax.axis("off")
fig.savefig("img/fig-bal-fam-arcs.png", dpi=150, bbox_inches="tight"); plt.close(fig)
print("wrote img/fig-bal-fam-{sketch,circle,links,arcs}.png")