"""CW: every natural motion is a guiding centre that drifts, carrying a 2:1 epicycle."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lambertp import mu, RE, cw_stm

a = RE + 420.0
n = np.sqrt(mu/a**3); T = 2*np.pi/n

# x0 [km], xd0 [km/s], yd0 [km/s]; y0 = 0 throughout
cases = [("(a) wavy path:  $|x_c|>\\frac{4}{3} A$", (2.0, 0.0, -2.5*n), BLUE),
         ("(b) cusped path:  $|x_c|=\\frac{4}{3} A$", (1.0, 0.0, 0.0), MAROON),
         ("(c) looped path:  $|x_c|<\\frac{4}{3} A$", (-1.0, 0.0, 2.5*n), GREEN)]
revs = 2
fig, axs = plt.subplots(3, 1, figsize=(10.8, 8.6))
for ax, (title, (x0, xd0, yd0), col) in zip(axs, cases):
    X0 = np.array([x0, 0.0, xd0, yd0])
    xc = 4*x0 + 2*yd0/n
    al, be = 3*x0 + 2*yd0/n, xd0/n
    A = np.hypot(al, be)
    yc = 0.0 - 2*be
    t = np.linspace(0, revs*T, 1200)
    P = np.array([cw_stm(n, tt) @ X0 for tt in t])
    ax.plot(P[:, 1], P[:, 0], color=col, lw=2.6, zorder=4)
    # guiding centre
    ax.plot(yc - 1.5*n*xc*t, np.full_like(t, xc), color=INK, lw=1.0, ls=(0, (4, 3)), zorder=2)
    # the epicycle ellipse, drawn about the guiding centre at t = 0, T, 2T
    u = np.linspace(0, 2*np.pi, 200)
    for k in range(revs + 1):
        cy = yc - 1.5*n*xc*k*T
        ax.plot(cy + 2*A*np.cos(u), xc + A*np.sin(u), color=col, lw=0.9, alpha=0.55, zorder=3)
        ax.plot(cy, xc, "+", color=INK, ms=9, zorder=5)
    ax.plot(P[0, 1], P[0, 0], "o", color=col, ms=8, zorder=6)
    ax.plot(0, 0, "*", color=INK, ms=13, zorder=6)
    ax.axhline(0, color=INK, lw=0.5)
    drift = -3*np.pi*xc
    ax.set_title(f"{title}      $x_c={xc:.0f}$ km,  $A={A:.0f}$ km,  guiding centre drifts "
                 f"${drift:+.1f}$ km per orbit", fontsize=11, loc="left")
    ax.set_ylabel(r"radial $x$ [km]")
    ax.set_aspect("equal", adjustable="datalim")
    despine(ax)
axs[-1].set_xlabel(r"along-track $y$ [km]  (direction of motion $\rightarrow$);   dot: start, $+$: guiding centre at $t=0,T,2T$, star: target")
fig.tight_layout()
fig.savefig("img/fig-cw-epicycle.png", dpi=160, bbox_inches="tight")
print("wrote img/fig-cw-epicycle.png")