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.75
b = a*np.sqrt(1.0 - e**2)          # semi-minor axis
c = a*e                            # centre-to-focus distance

E = np.deg2rad(65.0)               # eccentric anomaly of the sample point
Q = (a*np.cos(E), a*np.sin(E))     # on the auxiliary circle
P = (a*np.cos(E), b*np.sin(E))     # on the ellipse: the circle squashed by b/a
N = (a*np.cos(E), 0.0)             # foot of the perpendicular
F = (c, 0.0)                       # occupied focus
r = a*(1.0 - e*np.cos(E))
f = np.arctan2(P[1], P[0] - c)     # true anomaly, measured at F

fig, ax = plt.subplots(figsize=(9.0, 6.4))

th = np.linspace(0.0, 2*np.pi, 900)
ax.plot(a*np.cos(th), a*np.sin(th), color=BLUE, lw=1.6, ls=(0, (7, 4)), alpha=0.9, zorder=2)
ax.plot(a*np.cos(th), b*np.sin(th), color=INK, lw=2.4, zorder=3)
ax.plot([-a - 0.7, a + 0.7], [0, 0], color=INK, lw=1.0, alpha=0.55, zorder=1)

# the perpendicular Q -> P -> N : the projection that defines E
ax.plot([Q[0], N[0]], [Q[1], N[1]], color=INK, lw=1.3, ls=(0, (3, 3)), alpha=0.85, zorder=4)

# radius of the auxiliary circle: E is measured at the CENTRE
ax.plot([0, Q[0]], [0, Q[1]], color=PURPLE, lw=2.2, zorder=5)
ax.text(0.52*Q[0] - 0.28, 0.52*Q[1] + 0.06, "$a$", color=PURPLE, fontsize=14,
        ha="right", va="bottom")

# radius vector: f is measured at the FOCUS
ax.annotate("", xy=P, xytext=F, arrowprops=dict(arrowstyle="->", color=GREEN, lw=2.3))
ax.text(0.5*(F[0] + P[0]) + 0.22, 0.5*P[1] + 0.05, "$r$", color=GREEN, fontsize=14,
        ha="left", va="bottom")

# angle arcs
arcE = np.linspace(0.0, E, 90)
ax.plot(1.15*np.cos(arcE), 1.15*np.sin(arcE), color=PURPLE, lw=1.6, zorder=5)
ax.text(1.46*np.cos(E/2), 1.46*np.sin(E/2), "$E$", color=PURPLE, fontsize=15,
        ha="center", va="center")

arcf = np.linspace(0.0, f, 90)
ax.plot(c + 0.62*np.cos(arcf), 0.62*np.sin(arcf), color=GREEN, lw=1.6, zorder=5)
ax.text(c + 0.95*np.cos(f/2), 0.95*np.sin(f/2), "$f$", color=GREEN, fontsize=15,
        ha="center", va="center")

# points
ax.plot([F[0]], [F[1]], "o", color=MAROON, ms=12, zorder=7)
ax.plot([0], [0], "o", color=INK, ms=5, zorder=7)
ax.plot([P[0]], [P[1]], "o", color=GREEN, ms=8, zorder=7)
ax.plot([Q[0]], [Q[1]], "o", color=BLUE, ms=8, zorder=7)
ax.plot([N[0]], [N[1]], "o", color=INK, ms=4, zorder=7)
ax.plot([a], [0], "o", color=INK, ms=5, zorder=7)

ax.text(Q[0] + 0.22, Q[1] + 0.10, "$Q$", color=BLUE, fontsize=13, ha="left", va="bottom")
ax.text(P[0] + 0.24, P[1] + 0.02, "$P$ (the body)", color=GREEN, fontsize=12,
        ha="left", va="bottom")
ax.text(-0.12, 0.20, "centre $O$", color=INK, fontsize=11, ha="right", va="bottom")
ax.text(F[0] - 0.20, 0.20, "focus $F$", color=MAROON, fontsize=11, ha="right", va="bottom")
ax.text(a + 0.28, 0.22, "periapsis", color=INK, fontsize=11, ha="left", va="bottom")
ax.text(a + 0.28, -0.26, "$E = f = 0$", color=INK, fontsize=10, ha="left", va="top",
        alpha=0.8)

# measurements below the axis, inside the ellipse
ax.plot([N[0], N[0]], [0, -0.60], color=PURPLE, lw=0.8, ls=":", alpha=0.6, zorder=1)
ax.plot([c, c], [0, -1.35], color=MAROON, lw=0.8, ls=":", alpha=0.6, zorder=1)

ax.annotate("", xy=(N[0], -0.60), xytext=(0, -0.60),
            arrowprops=dict(arrowstyle="<->", color=PURPLE, lw=1.8))
ax.text(N[0]/2, -0.78, r"$a\cos E$", color=PURPLE, fontsize=13, ha="center", va="top")

ax.annotate("", xy=(c, -1.35), xytext=(0, -1.35),
            arrowprops=dict(arrowstyle="<->", color=MAROON, lw=1.8))
ax.text(c/2, -1.53, "$ae$", color=MAROON, fontsize=13, ha="center", va="top")

ax.text(-3.9, 2.85, "auxiliary circle,\nradius $a$", color=BLUE, fontsize=11,
        ha="left", va="center")

ax.set_xlim(-4.9, 5.6)
ax.set_ylim(-4.6, 4.6)
ax.set_aspect("equal")
ax.axis("off")
plt.tight_layout()
fig.savefig("img/fig-anomalies.png", dpi=160, bbox_inches="tight")
plt.close(fig)