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": 11,
    "axes.grid": False,
})

ecc    = [0.0, 0.2, 0.4, 0.6, 0.8, 0.95]
colors = [INK, BLUE, GREEN, AMBER, PURPLE, MAROON]

E = np.linspace(0.0, 2*np.pi, 1600)

def true_from_ecc(E, e):
    """f(E), continuous on [0, 2pi]: the half-angle relation in arctan2 form."""
    return 2.0*np.arctan2(np.sqrt(1.0 + e)*np.sin(0.5*E),
                          np.sqrt(1.0 - e)*np.cos(0.5*E))

def ecc_from_mean(M, e, tol=1e-13, itmax=60):
    """E(M) by Newton on Kepler's equation; the denominator 1 - e cos E = r/a
    never vanishes for e < 1, so E0 = M is a safe start at these eccentricities."""
    E = np.array(M, dtype=float)
    for _ in range(itmax):
        dE = -(E - e*np.sin(E) - M)/(1.0 - e*np.cos(E))
        E = E + dE
        if np.max(np.abs(dE)) < tol:
            break
    return E

fig, (axA, axB, axC) = plt.subplots(1, 3, figsize=(15.4, 4.8))

# ---- panel A : f versus E -------------------------------------------------
for e, col in zip(ecc, colors):
    f = true_from_ecc(E, e)
    axA.plot(np.rad2deg(E), np.rad2deg(f), color=col, lw=2.2 if e else 1.4,
             ls="-" if e else (0, (4, 3)), label=f"$e = {e:g}$")

axA.set_xlabel("eccentric anomaly $E$  [deg]")
axA.set_ylabel("true anomaly $f$  [deg]")
axA.set_xlim(0, 360); axA.set_ylim(0, 360)
axA.set_xticks(range(0, 361, 90)); axA.set_yticks(range(0, 361, 90))
axA.set_aspect("equal")
axA.axhline(180, color=INK, lw=0.7, alpha=0.3)
axA.axvline(180, color=INK, lw=0.7, alpha=0.3)
axA.set_title("$f$ against $E$", fontsize=12.5, pad=8)
axA.legend(frameon=False, fontsize=9.5, loc="lower right", ncol=2,
           handlelength=1.6, labelspacing=0.3, columnspacing=1.0)
for sp in ("top", "right"):
    axA.spines[sp].set_visible(False)

axA.text(184, 96, "all curves pass\nthrough $(180^\\circ,180^\\circ)$",
         fontsize=9, color=INK, alpha=0.75, ha="left", va="center")

# ---- panel B : the lead f - E --------------------------------------------
locus_E, locus_d = [], []
for e, col in zip(ecc, colors):
    f = true_from_ecc(E, e)
    axB.plot(np.rad2deg(E), np.rad2deg(f - E), color=col, lw=2.2 if e else 1.4,
             ls="-" if e else (0, (4, 3)))
    if e > 0:
        cE = (1.0 - np.sqrt(1.0 - e*e))/e          # d(f-E)/dE = 0  <=>  r = b
        Em = np.arccos(cE)
        fm = true_from_ecc(np.array([Em]), e)[0]
        locus_E.append(np.rad2deg(Em)); locus_d.append(np.rad2deg(fm - Em))
        axB.plot([np.rad2deg(Em)], [np.rad2deg(fm - Em)], "o", color=col, ms=6, zorder=6)

axB.plot(locus_E, locus_d, color=INK, lw=1.0, ls=":", alpha=0.7, zorder=5)
axB.annotate("maxima: $r = b$,\ni.e. $\\cos E = \\dfrac{1 - \\sqrt{1-e^{2}}}{e}$",
             xy=(locus_E[-2], locus_d[-2]), xytext=(150, 62),
             fontsize=9.5, color=INK, ha="left", va="center",
             arrowprops=dict(arrowstyle="->", color=INK, lw=1.1, alpha=0.8))

axB.axhline(0.0, color=INK, lw=0.8, alpha=0.45)
axB.set_xlabel("eccentric anomaly $E$  [deg]")
axB.set_ylabel("$f - E$  [deg]")
axB.set_xlim(0, 360); axB.set_ylim(-95, 95)
axB.set_xticks(range(0, 361, 90)); axB.set_yticks(range(-90, 91, 45))
axB.set_title("the lead of $f$ over $E$", fontsize=12.5, pad=8)
for sp in ("top", "right"):
    axB.spines[sp].set_visible(False)

axB.text(90, -70, "$f > E$ outbound\n$f < E$ inbound", fontsize=9.5,
         color=INK, alpha=0.8, ha="center", va="center")

# ---- panel C : f versus M -------------------------------------------------
# A different kind of relationship from panel A: E <-> f is pure geometry, read
# off a single frozen picture of the ellipse, while M carries the *time*.  The
# extra bow here is Kepler's second law on top of the focus offset.
M = np.linspace(0.0, 2*np.pi, 1600)

for e, col in zip(ecc, colors):
    f = true_from_ecc(ecc_from_mean(M, e), e)
    axC.plot(np.rad2deg(M), np.rad2deg(f), color=col, lw=2.2 if e else 1.4,
             ls="-" if e else (0, (4, 3)), label=f"$e = {e:g}$")

# the intermediate step, for one eccentricity, to show the two bows compounding
E_mid = ecc_from_mean(M, 0.6)
axC.plot(np.rad2deg(M), np.rad2deg(E_mid), color=AMBER, lw=1.3, ls=(0, (1.5, 2.5)))
_Mtag = 70.0
axC.annotate("$E$ against $M$ at $e = 0.6$:\nKepler's equation alone",
             xy=(_Mtag, np.rad2deg(ecc_from_mean(np.deg2rad(_Mtag), 0.6))),
             xytext=(132, 42), fontsize=9, color=INK, ha="left", va="center",
             arrowprops=dict(arrowstyle="->", color=INK, lw=1.0, alpha=0.8))

axC.set_xlabel("mean anomaly $M$  [deg]")
axC.set_ylabel("true anomaly $f$  [deg]")
axC.set_xlim(0, 360); axC.set_ylim(0, 360)
axC.set_xticks(range(0, 361, 90)); axC.set_yticks(range(0, 361, 90))
axC.set_aspect("equal")
axC.axhline(180, color=INK, lw=0.7, alpha=0.3)
axC.axvline(180, color=INK, lw=0.7, alpha=0.3)
axC.set_title("$f$ against $M$", fontsize=12.5, pad=8)
for sp in ("top", "right"):
    axC.spines[sp].set_visible(False)

axC.text(14, 332, "$M \\leq E \\leq f$ outbound", fontsize=9, color=INK,
         alpha=0.8, ha="left", va="center")

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