"""The rotating basis of the orbital plane, and why  d(rhat)/dt = fdot thetahat.

Left  -- the pair (rhat, thetahat) carried by the body, shown at two nearby true
         anomalies, together with the fixed perifocal pair (Phat, Qhat) it is
         built from.
Right -- the same two rhat's redrawn from a common origin: their tips lie on the
         unit circle, so the increment Delta(rhat) is a chord of that circle,
         of length 2 sin(Delta f / 2) -> Delta f, pointing along thetahat.
"""

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 = 2.2, 0.45
p = a*(1.0 - e**2)
f0, df = np.deg2rad(32.0), np.deg2rad(60.0)


def rvec(f):
    return p/(1.0 + e*np.cos(f))*np.array([np.cos(f), np.sin(f)])


def basis(f):
    return (np.array([np.cos(f), np.sin(f)]),          # rhat
            np.array([-np.sin(f), np.cos(f)]))         # thetahat


def arrow(ax, tail, vec, color, lw=2.2, alpha=1.0, ls="-"):
    ax.annotate("", xy=tail + vec, xytext=tail,
                arrowprops=dict(arrowstyle="-|>", color=color, lw=lw,
                                alpha=alpha, linestyle=ls,
                                shrinkA=0, shrinkB=0))


fig, (axL, axR) = plt.subplots(1, 2, figsize=(13.0, 5.4),
                              gridspec_kw=dict(width_ratios=[1.45, 1.0]))

# ── left: the basis carried along the orbit ─────────────────────────────────
th = np.linspace(0.0, 2*np.pi, 800)
orb = (p/(1.0 + e*np.cos(th)))*np.array([np.cos(th), np.sin(th)])
axL.plot(orb[0], orb[1], color=INK, lw=2.2)
axL.plot([0], [0], "o", color=AMBER, ms=11, zorder=6)
axL.text(0.0, -0.22, "focus", color=AMBER, fontsize=11, ha="center", va="top")

# fixed perifocal basis
arrow(axL, np.zeros(2), np.array([0.72, 0.0]), GREEN, lw=1.8)
arrow(axL, np.zeros(2), np.array([0.0, 0.72]), GREEN, lw=1.8)
axL.text(0.78, -0.04, r"$\hat{\mathbf{P}}$", color=GREEN, fontsize=15, va="top")
axL.text(-0.07, 0.78, r"$\hat{\mathbf{Q}}$", color=GREEN, fontsize=15,
         ha="right")
axL.text(-1.45, -0.95, "fixed perifocal basis\n($\\hat{\\mathbf{P}}$ towards"
         " periapsis)", color=GREEN, fontsize=10.5, ha="center")

L = 0.80
for f, alpha in ((f0, 1.0), (f0 + df, 0.40)):
    P = rvec(f)
    rh, th_ = basis(f)
    axL.plot([0, P[0]], [0, P[1]], color=INK, lw=1.0, alpha=0.5*alpha)
    axL.plot([P[0]], [P[1]], "o", color=MAROON, ms=8, alpha=alpha, zorder=6)
    arrow(axL, P, L*rh, MAROON, alpha=alpha)
    arrow(axL, P, L*th_, BLUE, alpha=alpha)
    axL.text(*(P + (L + 0.14)*rh), r"$\hat{\mathbf{r}}$", color=MAROON,
             fontsize=15, alpha=alpha, ha="left", va="center")
    axL.text(*(P + (L + 0.16)*th_), r"$\hat{\boldsymbol{\theta}}$", color=BLUE,
             fontsize=15, alpha=alpha, ha="center", va="center")

# the swept angle at the focus
arc = np.linspace(f0, f0 + df, 60)
axL.plot(0.62*np.cos(arc), 0.62*np.sin(arc), color=INK, lw=1.4)
axL.text(0.85*np.cos(f0 + 0.5*df), 0.85*np.sin(f0 + 0.5*df), r"$\Delta f$",
         color=INK, fontsize=14, ha="center", va="center")

axL.set_xlim(-3.6, 3.2)
axL.set_ylim(-2.35, 3.15)
axL.set_aspect("equal")
axL.axis("off")
axL.set_title("the basis is carried by the body", fontsize=12.5, pad=2)

# ── right: the unit circle of directions ────────────────────────────────────
c = np.linspace(0.0, 2*np.pi, 400)
axR.plot(np.cos(c), np.sin(c), color=INK, lw=1.2, ls=(0, (5, 4)), alpha=0.55)
axR.plot([0], [0], "o", color=INK, ms=5)

rh0, th0 = basis(f0)
rh1, th1 = basis(f0 + df)
arrow(axR, np.zeros(2), rh0, MAROON)
arrow(axR, np.zeros(2), rh1, MAROON, alpha=0.42)
arrow(axR, rh0, 0.72*th0, BLUE, lw=2.0, ls=(0, (4, 2)))
arrow(axR, rh0, rh1 - rh0, PURPLE, lw=2.6)          # the increment itself

axR.text(*(1.08*rh0 + np.array([0.05, -0.03])), r"$\hat{\mathbf{r}}(f)$",
         color=MAROON, fontsize=14, ha="left", va="top")
axR.text(*(1.10*rh1 + np.array([-0.06, 0.03])), r"$\hat{\mathbf{r}}(f+\Delta f)$",
         color=MAROON, fontsize=14, alpha=0.8, ha="right", va="bottom")
chord = rh1 - rh0
axR.text(*(0.78*(rh0 + 0.5*chord)), r"$\Delta\hat{\mathbf{r}}$", color=PURPLE,
         fontsize=15, ha="center", va="center")
axR.text(*(rh0 + 0.80*th0), r"$\hat{\boldsymbol{\theta}}$", color=BLUE,
         fontsize=14, ha="left", va="bottom")

arc2 = np.linspace(f0, f0 + df, 60)
axR.plot(0.26*np.cos(arc2), 0.26*np.sin(arc2), color=INK, lw=1.3)
axR.text(0.40*np.cos(f0 + 0.5*df), 0.40*np.sin(f0 + 0.5*df), r"$\Delta f$",
         color=INK, fontsize=13, ha="center", va="center")

axR.text(0.0, -1.25,
         r"$|\Delta\hat{\mathbf{r}}| = 2\sin\dfrac{\Delta f}{2}"
         r"\;\longrightarrow\; \Delta f$"
         "\n"
         r"$\Delta\hat{\mathbf{r}} \parallel \hat{\boldsymbol{\theta}}$"
         r" as $\Delta f \to 0$",
         fontsize=13, ha="center", va="top")

axR.set_xlim(-1.35, 1.55)
axR.set_ylim(-2.15, 1.45)
axR.set_aspect("equal")
axR.axis("off")
axR.set_title(r"the same two $\hat{\mathbf{r}}$, from one origin", fontsize=12.5,
              pad=2)

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