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": True, "grid.alpha": 0.25, "grid.color": INK,
})

cases = [(0.0, "$e = 0$ (circle)", GREEN),
         (0.5, "$e = 0.5$ (ellipse)", BLUE),
         (1.0, "$e = 1$ (parabola)", MAROON),
         (1.5, "$e = 1.5$ (hyperbola)", PURPLE)]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13.2, 5.2))

for e, lab, col in cases:
    f = np.linspace(-180.0, 180.0, 2001)
    fr = np.deg2rad(f)
    keep = 1.0 + e*np.cos(fr) > 1e-3          # only where the conic exists
    v = np.sqrt(np.clip(1.0 + 2*e*np.cos(fr) + e**2, 0.0, None))
    g = np.rad2deg(np.arctan2(e*np.sin(fr), 1.0 + e*np.cos(fr)))
    ax1.plot(f[keep], v[keep], color=col, lw=2.2, label=lab)
    ax2.plot(f[keep], g[keep], color=col, lw=2.2, label=lab)
    if e > 1.0:                                # speed at infinity
        vinf = np.sqrt(e**2 - 1.0)
        ax1.axhline(vinf, color=col, lw=1.2, ls=(0, (5, 4)), alpha=0.8)
        ax1.text(-172, vinf + 0.13, r"$v_\infty = \sqrt{\mu/p}\,\sqrt{e^2-1}$",
                 color=col, fontsize=11, ha="left", va="bottom")
    if 0.0 < e < 1.0:                          # maximum flight path angle
        fm = np.rad2deg(np.arccos(-e))
        ax2.plot([fm], [np.rad2deg(np.arcsin(e))], "o", color=col, ms=7, zorder=5)
        ax2.text(fm - 8, np.rad2deg(np.arcsin(e)) + 4.0,
                 r"$\gamma_{\max} = \arcsin e$   at   $\cos f = -e$", color=col,
                 fontsize=11, ha="right", va="bottom")

ax1.set_xlabel("true anomaly $f$ (deg)")
ax1.set_ylabel(r"$v \,/\, \sqrt{\mu/p}$")
ax1.set_title(r"speed: $v = \sqrt{\mu/p}\;\sqrt{1 + 2e\cos f + e^2}$", fontsize=13, pad=10)
ax1.set_xlim(-180, 180)
ax1.set_ylim(0.0, 2.6)
ax1.set_xticks(range(-180, 181, 60))
ax1.legend(frameon=False, fontsize=11, loc="upper right")

ax2.axhline(0.0, color=INK, lw=1.0, alpha=0.5)
ax2.set_xlabel("true anomaly $f$ (deg)")
ax2.set_ylabel(r"flight path angle $\gamma$ (deg)")
ax2.set_title(r"$\tan\gamma = \dfrac{e\sin f}{1 + e\cos f}$", fontsize=13, pad=10)
ax2.set_xlim(-180, 180)
ax2.set_ylim(-95, 95)
ax2.set_xticks(range(-180, 181, 60))
ax2.set_yticks(range(-90, 91, 30))

plt.tight_layout()
fig.savefig("img/fig-speed-fpa.png", dpi=160, bbox_inches="tight")
plt.close(fig)