"""The perifocal frame, and the rotating basis carried along with the body.

Left panel   the conic drawn in its own plane, with the perifocal basis
             {phat, qhat, what} pinned to the *orbit* (phat at periapsis) and
             the polar basis {rhat, thetahat} pinned to the *body*.

Right panel  the only relation between them.  Both pairs are orthonormal and
             both lie in the orbital plane, so one is the other turned through
             the true anomaly f:

                 rhat     =  cos f  phat + sin f  qhat,
                 thetahat = -sin f  phat + cos f  qhat.

That is the whole content of the perifocal decomposition: the perifocal frame
does not turn, the polar frame does, and f is the angle between them.

Run from the deck directory:   python codes/fig-perifocal-frame.py
"""

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

# colours shared with fig-orbital-elements-3d.png: periapsis purple, radius and
# true anomaly maroon, so the two figures can be read side by side.
C_P, C_Q, C_R, C_TH = PURPLE, AMBER, MAROON, BLUE

a, e = 4.0, 0.6
p = a * (1.0 - e**2)
f0 = np.deg2rad(52.0)                       # where the body is drawn
U = 1.35                                    # drawn length of a unit vector


def conic(f):
    r = p / (1.0 + e * np.cos(f))
    return r * np.cos(f), r * np.sin(f)


def arrow(ax, tail, head, color, lw=2.4):
    ax.annotate("", xy=head, xytext=tail,
                arrowprops=dict(arrowstyle="-|>", color=color, lw=lw,
                                shrinkA=0, shrinkB=0, mutation_scale=15))


fig, (axL, axR) = plt.subplots(1, 2, figsize=(14.4, 6.3),
                               gridspec_kw=dict(width_ratios=[1.42, 1.0]))

# ══ left: the two frames, in place on the orbit ═════════════════════════════
th = np.linspace(0.0, 2 * np.pi, 900)
axL.plot(*conic(th), color=INK, lw=2.4, zorder=3)
axL.plot([-a * (1 + e) - 0.35, a * (1 - e) + 0.9], [0, 0], ls=(0, (6, 4)),
         color=INK, lw=1.0, alpha=0.45, zorder=1)

# the focus, with what-hat coming out of the page
axL.plot([0], [0], "o", mfc=CREAM, mec=GREEN, mew=1.7, ms=16, zorder=8)
axL.plot([0], [0], "o", color=GREEN, ms=4.5, zorder=9)
axL.text(-0.26, 0.14, "focus", color=INK, fontsize=11, ha="right",
         va="bottom")
axL.text(-0.26, -0.16, r"$\hat{\mathbf{w}}$ out of the page", color=GREEN,
         fontsize=11, ha="right", va="top")

# perifocal basis: fixed to the orbit, so drawn at the focus
arrow(axL, (0, 0), (U, 0), C_P, lw=2.7)
axL.text(U - 0.05, 0.12, r"$\hat{\mathbf{p}}$", color=C_P, fontsize=18,
         ha="center", va="bottom")
arrow(axL, (0, 0), (0, U), C_Q, lw=2.7)
axL.text(-0.12, U + 0.06, r"$\hat{\mathbf{q}}$", color=C_Q, fontsize=18,
         ha="right", va="bottom")

# apsides
axL.plot([a * (1 - e), -a * (1 + e)], [0, 0], "o", color=INK, ms=5, zorder=6)
axL.text(a * (1 - e) + 0.16, -0.22, "periapsis", color=INK, fontsize=10.5,
         ha="left", va="top")

# the body, its radius, and the true anomaly
bx, by = conic(f0)
axL.plot([0, bx], [0, by], color=C_R, lw=2.0, zorder=4)
axL.plot([bx], [by], "o", color=C_R, ms=10, zorder=9)
axL.text(0.44, 0.86, r"$\mathbf{r}$", color=C_R, fontsize=16, ha="right",
         va="center")

arc = np.linspace(0.0, f0, 90)
axL.plot(0.62 * np.cos(arc), 0.62 * np.sin(arc), color=C_R, lw=1.6, zorder=5)
axL.text(0.80 * np.cos(f0 / 2), 0.80 * np.sin(f0 / 2), "$f$", color=C_R,
         fontsize=16, ha="center", va="center")

# polar basis: carried by the body, so drawn at the body
rhat = np.array([np.cos(f0), np.sin(f0)])
that = np.array([-np.sin(f0), np.cos(f0)])
arrow(axL, (bx, by), (bx + U * rhat[0], by + U * rhat[1]), C_R, lw=2.7)
axL.text(bx + (U + 0.14) * rhat[0], by + (U + 0.14) * rhat[1],
         r"$\hat{\mathbf{r}}$", color=C_R, fontsize=18, ha="left", va="bottom")
arrow(axL, (bx, by), (bx + U * that[0], by + U * that[1]), C_TH, lw=2.7)
axL.text(bx + (U + 0.16) * that[0], by + (U + 0.16) * that[1],
         r"$\hat{\boldsymbol{\theta}}$", color=C_TH, fontsize=18, ha="right",
         va="bottom")

# right angle between rhat and thetahat
sq = 0.26
axL.plot(*np.array([[bx + sq * rhat[0], by + sq * rhat[1]],
                    [bx + sq * (rhat[0] + that[0]), by + sq * (rhat[1] + that[1])],
                    [bx + sq * that[0], by + sq * that[1]]]).T,
         color=INK, lw=0.9, alpha=0.6)

# direction of motion, parked in empty space below periapsis
arrow(axL, conic(np.deg2rad(-52.0)), conic(np.deg2rad(-40.0)), INK, lw=1.7)
axL.text(1.62, -1.34, "motion", color=INK, fontsize=10.5, alpha=0.85,
         ha="left", va="center")

axL.set_title(r"$\{\hat{\mathbf{p}},\hat{\mathbf{q}},\hat{\mathbf{w}}\}$ is fixed to "
              r"the orbit;   $\{\hat{\mathbf{r}},\hat{\boldsymbol{\theta}}\}$ is "
              r"carried by the body", fontsize=13, pad=14)
axL.set_xlim(-7.15, 4.05)
axL.set_ylim(-3.85, 4.05)
axL.set_aspect("equal")
axL.axis("off")

# ══ right: the decomposition, on the unit circle ════════════════════════════
c, s = np.cos(f0), np.sin(f0)

axR.plot(np.cos(th), np.sin(th), color=INK, lw=1.0, alpha=0.30, zorder=1)
axR.plot([-1.38, 1.38], [0, 0], color=INK, lw=0.7, alpha=0.28)
axR.plot([0, 0], [-1.38, 1.38], color=INK, lw=0.7, alpha=0.28)

arrow(axR, (0, 0), (1.0, 0), C_P, lw=2.7)
axR.text(1.06, -0.05, r"$\hat{\mathbf{p}}$", color=C_P, fontsize=18, ha="left",
         va="top")
arrow(axR, (0, 0), (0, 1.0), C_Q, lw=2.7)
axR.text(0.06, 1.06, r"$\hat{\mathbf{q}}$", color=C_Q, fontsize=18, ha="left",
         va="bottom")

arrow(axR, (0, 0), (c, s), C_R, lw=2.7)
axR.text(c + 0.06, s + 0.06, r"$\hat{\mathbf{r}}$", color=C_R, fontsize=18,
         ha="left", va="bottom")
arrow(axR, (0, 0), (-s, c), C_TH, lw=2.7)
axR.text(-s - 0.06, c + 0.06, r"$\hat{\boldsymbol{\theta}}$", color=C_TH,
         fontsize=18, ha="right", va="bottom")

# components of rhat: down to the phat axis, across to the qhat axis
axR.plot([c, c], [0, s], ls=(0, (3, 3)), color=C_R, lw=1.3, alpha=0.9)
axR.plot([0, c], [s, s], ls=(0, (3, 3)), color=C_R, lw=1.3, alpha=0.9)
axR.text(c, -0.07, r"$\cos f$", color=C_R, fontsize=13, ha="center", va="top")
axR.text(-0.06, s, r"$\sin f$", color=C_R, fontsize=13, ha="right", va="center")

# components of thetahat
axR.plot([-s, -s], [0, c], ls=(0, (3, 3)), color=C_TH, lw=1.3, alpha=0.9)
axR.plot([0, -s], [c, c], ls=(0, (3, 3)), color=C_TH, lw=1.3, alpha=0.9)
axR.text(-s, -0.07, r"$-\sin f$", color=C_TH, fontsize=13, ha="center",
         va="top")
axR.text(0.06, c, r"$\cos f$", color=C_TH, fontsize=13, ha="left", va="center")

# the two angles
arc = np.linspace(0.0, f0, 80)
axR.plot(0.28 * np.cos(arc), 0.28 * np.sin(arc), color=C_R, lw=1.6)
axR.text(0.40 * np.cos(f0 / 2), 0.40 * np.sin(f0 / 2), "$f$", color=C_R,
         fontsize=15, ha="center", va="center")
rh = np.array([c, s])
tht = np.array([-s, c])
q = 0.20
axR.plot(*np.array([q * rh, q * (rh + tht), q * tht]).T, color=C_TH, lw=1.3,
         alpha=0.9)

axR.text(0.0, -1.62,
         r"$\hat{\mathbf{r}} = \cos f\,\hat{\mathbf{p}} + \sin f\,\hat{\mathbf{q}}$"
         "\n"
         r"$\hat{\boldsymbol{\theta}} = -\sin f\,\hat{\mathbf{p}} + \cos f\,\hat{\mathbf{q}}$",
         fontsize=15.5, ha="center", va="top", linespacing=2.0)

axR.set_title("one frame is the other turned through $f$", fontsize=13, pad=14)
axR.set_xlim(-1.55, 1.55)
axR.set_ylim(-2.45, 1.55)
axR.set_aspect("equal")
axR.axis("off")

plt.tight_layout()
fig.savefig("img/fig-perifocal-frame.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-perifocal-frame.png")
