import os
import numpy as np
import matplotlib.pyplot as plt

os.makedirs("img", exist_ok=True)

CREAM, INK = "#fcfbf7", "#2f2a24"
MAROON, BLUE, GREEN, AMBER, PURPLE = "#6b1f1f", "#2f5ea8", "#2d6a2d", "#b45309", "#7e22ce"

plt.rcParams.update({
    "figure.facecolor": CREAM, "savefig.facecolor": CREAM, "axes.facecolor": CREAM,
    "axes.edgecolor": INK, "axes.labelcolor": INK, "text.color": INK,
    "xtick.color": INK, "ytick.color": INK, "font.size": 12,
    "axes.grid": False,
})

e = 1.6
A = 2.5                            # |a| = p/(e^2-1), centre-to-vertex distance
b = A*np.sqrt(e**2 - 1.0)          # semi-minor axis = aiming radius
c = A*e                            # centre-to-focus distance
th_a = np.arctan2(b, A)            # asymptote tilt, cos(th_a) = 1/e
f_inf = np.pi - th_a               # true anomaly of the asymptote
delta = 2.0*f_inf - np.pi          # turning angle, sin(delta/2) = 1/e

# The attracting focus sits at (-c, 0) so that periapsis lies in the +x
# direction from it, as in the elliptic figure; the physical branch then
# opens to the left, and the body runs in along the lower asymptote.
FOC = np.array([-c, 0.0])
EIN = np.array([np.cos(th_a), np.sin(th_a)])        # incoming velocity direction
EOUT = np.array([-np.cos(th_a), np.sin(th_a)])      # outgoing velocity direction

fig, (axL, axR) = plt.subplots(
    1, 2, figsize=(16.4, 6.4), gridspec_kw=dict(width_ratios=[1.0, 0.92]))

# ======================================================================
# Left: the conic itself
# ======================================================================
u = np.linspace(-2.05, 2.05, 600)
xb, yb = A*np.cosh(u), b*np.sinh(u)
axL.plot(-xb, yb, color=INK, lw=2.6, zorder=5)
uo = np.linspace(-1.5, 1.5, 400)
axL.plot(A*np.cosh(uo), b*np.sinh(uo), color=MAROON, lw=1.8, ls=(0, (5, 4)),
         alpha=0.32, zorder=3)
axL.annotate("the *other* branch: the same $h$ and\n"
             r"$\varepsilon$ under a *repulsive* $1/r^{2}$ force",
             xy=(A*np.cosh(-0.62), b*np.sinh(-0.62)), xytext=(6.3, -2.3),
             color=MAROON, alpha=0.75, fontsize=10.5, ha="center", va="center",
             arrowprops=dict(arrowstyle="->", color=MAROON, alpha=0.55, lw=1.1))

t = np.linspace(-4.6, 4.6, 2)
for s in (+1.0, -1.0):
    axL.plot(t*A, s*t*b, ls=(0, (7, 5)), color=INK, lw=1.1, alpha=0.55, zorder=1)
axL.plot([-A, A, A, -A, -A], [-b, -b, b, b, -b],
         ls=(0, (2, 4)), color=INK, lw=0.9, alpha=0.35, zorder=1)
axL.text(3.25, 4.20, "asymptote", color=INK, alpha=0.75, fontsize=10.5,
         rotation=np.rad2deg(th_a), rotation_mode="anchor", ha="center", va="bottom")

arc = np.linspace(0.0, th_a, 80)
axL.plot(1.55*np.cos(arc), 1.55*np.sin(arc), color=PURPLE, lw=1.5, zorder=6)
axL.text(1.90, 1.24, r"$\arctan\dfrac{b}{|a|}$", color=PURPLE, fontsize=12,
         ha="left", va="center", zorder=8)

axL.plot([FOC[0]], [FOC[1]], "o", color=MAROON, ms=12, zorder=7)
axL.plot([0.0], [0.0], "o", color=INK, ms=5, zorder=7)
axL.plot([-A], [0.0], "o", color=INK, ms=5, zorder=7)
axL.annotate("focus $F$\n(attracting body)", xy=(FOC[0], 0.0), xytext=(-7.4, 2.1),
             color=MAROON, fontsize=11, ha="center", va="bottom",
             arrowprops=dict(arrowstyle="->", color=MAROON, lw=1.3))
axL.text(0.30, -0.30, "centre $O$", color=INK, fontsize=10.5, ha="left", va="top")
axL.annotate("vertex\n(periapsis)", xy=(-A, 0.0), xytext=(-5.0, 3.5),
             color=INK, fontsize=11, ha="center", va="bottom",
             arrowprops=dict(arrowstyle="->", color=INK, lw=1.2))

axL.annotate("", xy=(-A, 0.0), xytext=(0.0, 0.0),
             arrowprops=dict(arrowstyle="<->", color=AMBER, lw=2.0))
axL.text(-A/2, 0.16, "$|a|$", color=AMBER, fontsize=14, ha="center", va="bottom")

axL.annotate("", xy=(-A, 0.0), xytext=(FOC[0], 0.0),
             arrowprops=dict(arrowstyle="<->", color=BLUE, lw=2.0))
axL.text(-(c + A)/2, 0.16, "$r_p$", color=BLUE, fontsize=14,
         ha="center", va="bottom")

yd = -1.95
for x0 in (0.0, FOC[0]):
    axL.plot([x0, x0], [0.0, yd], ls=(0, (2, 3)), color=PURPLE, lw=0.9,
             alpha=0.6, zorder=1)
axL.annotate("", xy=(FOC[0], yd), xytext=(0.0, yd),
             arrowprops=dict(arrowstyle="<->", color=PURPLE, lw=2.0), zorder=1)
axL.text(-c/2 + 0.75, yd - 0.18, "$|a|e$", color=PURPLE, fontsize=14,
         ha="center", va="top", zorder=6)

# the conjugate axis: the segment 2b through the centre, met by no point
# of the curve -- which is why b has to be defined rather than measured
axL.plot([0.0, 0.0], [-b, b], color=AMBER, lw=1.6, alpha=0.45, zorder=2)
axL.text(0.75, -b - 0.18, "conjugate axis $2b$", color=AMBER, alpha=0.95,
         fontsize=10.5, ha="center", va="top")
axL.annotate("", xy=(0.0, b), xytext=(0.0, 0.0),
             arrowprops=dict(arrowstyle="<->", color=AMBER, lw=2.0))
axL.text(0.0, b + 0.20, r"$b = |a|\sqrt{e^{2}-1}$", color=AMBER, fontsize=12.5,
         ha="center", va="bottom")

ray = FOC + 6.4*np.array([np.cos(-f_inf), np.sin(-f_inf)])
axL.plot([FOC[0], ray[0]], [FOC[1], ray[1]], ls=(0, (1, 3)), color=GREEN,
         lw=1.5, zorder=4)
arc = np.linspace(0.0, -f_inf, 120)
axL.plot(FOC[0] + 1.15*np.cos(arc), 1.15*np.sin(arc), color=GREEN, lw=1.6, zorder=6)
axL.text(FOC[0] + 1.45*np.cos(-0.45*f_inf), 1.45*np.sin(-0.45*f_inf),
         r"$f_\infty$", color=GREEN, fontsize=14, ha="center", va="center")
axL.text(0.0, -4.4, r"the body lives only on $|f| < f_\infty$,"
         "\n" r"$\cos f_\infty = -1/e$", color=GREEN, fontsize=11,
         ha="center", va="top")

axL.set_xlim(-9.6, 9.2)
axL.set_ylim(-5.6, 5.6)
axL.set_aspect("equal")
axL.axis("off")

# ======================================================================
# Right: the flyby -- turning angle and aiming radius
# ======================================================================
u = np.linspace(-1.42, 1.42, 500)
axR.plot(-A*np.cosh(u), b*np.sinh(u), color=INK, lw=2.6, zorder=5)

s = np.linspace(-7.6, 7.6, 2)
for d in (EIN, EOUT):
    axR.plot(s*d[0], s*d[1], ls=(0, (7, 5)), color=INK, lw=1.1, alpha=0.5, zorder=1)
axR.plot([FOC[0]], [FOC[1]], "o", color=MAROON, ms=12, zorder=7)
axR.plot([0.0], [0.0], "o", color=INK, ms=5, zorder=7)
axR.text(FOC[0] - 0.15, 0.28, "$F$", color=MAROON, fontsize=13,
         ha="right", va="bottom")
axR.text(0.30, -0.30, "$O$", color=INK, fontsize=12, ha="left", va="top")

# incoming and outgoing velocity at infinity
P1, P2 = 6.7*(-EIN), 5.2*(-EIN)           # along the lower-left asymptote
axR.annotate("", xy=P2, xytext=P1,
             arrowprops=dict(arrowstyle="-|>", color=GREEN, lw=2.4,
                             mutation_scale=18))
axR.text(P1[0] + 1.05, P1[1] + 0.25, r"$\mathbf{v}_\infty$ in", color=GREEN,
         fontsize=12.5, ha="left", va="center")
P3, P4 = 5.2*EOUT, 6.7*EOUT               # along the upper-left asymptote
axR.annotate("", xy=P4, xytext=P3,
             arrowprops=dict(arrowstyle="-|>", color=GREEN, lw=2.4,
                             mutation_scale=18))
axR.text(P4[0] + 1.05, P4[1] - 0.25, r"$\mathbf{v}_\infty$ out", color=GREEN,
         fontsize=12.5, ha="left", va="center")

# the turning angle, read at O between the two asymptote directions
axR.annotate("", xy=2.9*EIN, xytext=(0.0, 0.0),
             arrowprops=dict(arrowstyle="-|>", color=GREEN, lw=1.8, alpha=0.9,
                             mutation_scale=15))
axR.text(3.05*EIN[0], 3.05*EIN[1] - 0.12,
         "incoming direction,\ncarried to $O$", color=GREEN, fontsize=10.5,
         alpha=0.9, ha="left", va="top")
arc = np.linspace(th_a, f_inf, 140)
axR.plot(2.05*np.cos(arc), 2.05*np.sin(arc), color=GREEN, lw=1.8, zorder=6)
axR.text(0.0, 2.40, r"$\delta$", color=GREEN, fontsize=15, ha="center", va="bottom")

# the aiming radius: perpendicular from the focus to the incoming asymptote
foot = np.dot(FOC, -EIN)*(-EIN)
axR.annotate("", xy=foot, xytext=tuple(FOC),
             arrowprops=dict(arrowstyle="<->", color=PURPLE, lw=2.0), zorder=6)
n = np.array([EIN[1], -EIN[0]])                      # unit normal, F -> foot
tick = 0.34
axR.plot([foot[0] - tick*n[0], foot[0] - tick*n[0] - tick*EIN[0],
          foot[0] - tick*EIN[0]],
         [foot[1] - tick*n[1], foot[1] - tick*n[1] - tick*EIN[1],
          foot[1] - tick*EIN[1]], color=PURPLE, lw=1.0, alpha=0.9, zorder=6)
mid = 0.5*(FOC + foot)
axR.text(mid[0] - 0.55, mid[1] - 0.30, r"$\Delta$", color=PURPLE, fontsize=15,
         ha="right", va="top")

axR.text(2.5, -1.15,
         r"$\Delta = b = |a|\sqrt{e^{2}-1}$" "\n"
         r"$h = \Delta\,v_\infty$" "\n"
         r"$\sin\dfrac{\delta}{2} = \dfrac{1}{e}$",
         color=INK, fontsize=12, ha="left", va="top", linespacing=1.9,
         bbox=dict(facecolor=CREAM, edgecolor="none", alpha=0.85, pad=4.0))

axR.set_xlim(-9.6, 8.2)
axR.set_ylim(-6.3, 6.3)
axR.set_aspect("equal")
axR.axis("off")

plt.tight_layout(rect=(0, 0, 1, 0.95))
fig.text(0.262, 0.968, "the conic", color=INK, fontsize=13.5,
         ha="center", va="center")
fig.text(0.760, 0.968, "the flyby: turning angle and aiming radius",
         color=INK, fontsize=13.5, ha="center", va="center")
fig.savefig("img/fig-hyperbola-geometry.png", dpi=160, bbox_inches="tight")
plt.close(fig)