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,
})

a, e = 4.0, 0.6
b = a*np.sqrt(1.0 - e**2)          # semi-minor axis,  b^2 = a p
c = a*e                            # centre-to-focus distance
p = a*(1.0 - e**2)                 # semi-latus rectum

fig, ax = plt.subplots(figsize=(8.8, 5.6))

th = np.linspace(0.0, 2*np.pi, 800)
ax.plot(a*np.cos(th), b*np.sin(th), color=INK, lw=2.4, zorder=3)

# axes of the ellipse
ax.plot([-a, a], [0, 0], ls=(0, (6, 4)), color=INK, lw=1.0, alpha=0.55, zorder=1)
ax.plot([0, 0], [-b, b], ls=(0, (6, 4)), color=INK, lw=1.0, alpha=0.55, zorder=1)

# occupied focus, empty focus, centre
ax.plot([c], [0], "o", color=MAROON, ms=12, zorder=6)
ax.plot([-c], [0], "o", mfc=CREAM, mec=INK, mew=1.5, ms=9, zorder=6)
ax.plot([0], [0], "o", color=INK, ms=5, zorder=6)
ax.annotate("focus $F$\n(attracting body)", xy=(c, 0), xytext=(c - 1.3, -2.45),
            color=MAROON, fontsize=11, ha="right", va="center",
            arrowprops=dict(arrowstyle="->", color=MAROON, lw=1.3))
ax.text(-c, -0.30, "empty focus", color=INK, fontsize=10, ha="center", va="top")
ax.text(-0.10, 0.22, "centre $O$", color=INK, fontsize=10, ha="right", va="bottom")

# semi-major axis a  (centre -> apoapsis side)
ax.annotate("", xy=(-a, 0), xytext=(0, 0),
            arrowprops=dict(arrowstyle="<->", color=AMBER, lw=2.0))
ax.text(-a/2, 0.20, "$a$", color=AMBER, fontsize=14, ha="center", va="bottom")

# centre-to-focus distance ae
ax.annotate("", xy=(c, 0), xytext=(0, 0),
            arrowprops=dict(arrowstyle="<->", color=PURPLE, lw=2.0))
ax.text(c/2, 0.20, "$ae$", color=PURPLE, fontsize=13, ha="center", va="bottom")

# semi-minor axis b
ax.annotate("", xy=(0, b), xytext=(0, 0),
            arrowprops=dict(arrowstyle="<->", color=AMBER, lw=2.0))
ax.text(-0.15, b/2, "$b$", color=AMBER, fontsize=14, ha="right", va="center")

# semi-latus rectum p, drawn downwards from the focus
ax.annotate("", xy=(c, -p), xytext=(c, 0),
            arrowprops=dict(arrowstyle="<->", color=BLUE, lw=2.0))
ax.text(c + 0.18, -p/2, "$p$", color=BLUE, fontsize=14, ha="left", va="center")

# radius vector and true anomaly
f = np.deg2rad(118.0)
rf = p/(1.0 + e*np.cos(f))
P = (c + rf*np.cos(f), rf*np.sin(f))
ax.annotate("", xy=P, xytext=(c, 0),
            arrowprops=dict(arrowstyle="->", color=GREEN, lw=2.2))
ax.plot([P[0]], [P[1]], "o", color=GREEN, ms=7, zorder=6)
ax.text(0.5*(c + P[0]) - 0.42, 0.5*P[1] + 0.10, "$r$", color=GREEN, fontsize=14,
        ha="right", va="bottom")
arc = np.linspace(0.0, f, 80)
ax.plot(c + 0.95*np.cos(arc), 0.95*np.sin(arc), color=GREEN, lw=1.5)
ax.text(c + 1.20*np.cos(f/2), 1.20*np.sin(f/2), "$f$", color=GREEN, fontsize=14,
        ha="center", va="center")

# apsides
ax.plot([a, -a], [0, 0], "o", color=INK, ms=5, zorder=6)
ax.text(a + 0.20, 0.30, "periapsis\n$r_p = a(1-e)$", color=INK, fontsize=11,
        ha="left", va="bottom")
ax.text(-a - 0.20, 0.30, "apoapsis\n$r_a = a(1+e)$", color=INK, fontsize=11,
        ha="right", va="bottom")

ax.set_xlim(-7.4, 7.4)
ax.set_ylim(-4.2, 4.6)
ax.set_aspect("equal")
ax.axis("off")
plt.tight_layout()
fig.savefig("img/fig-conic-geometry.png", dpi=160, bbox_inches="tight")
plt.close(fig)