"""CW: four natural motions about a target in a 420 km circular orbit."""
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
dv = 1e-3                                       # 1 m/s in km/s
cases = [("(a) prograde 1 m/s at the target", [0, 0, 0, dv], 2, MAROON),
         ("(b) radial (outward) 1 m/s at the target", [0, 0, dv, 0], 1, BLUE),
         ("(c) 1 km below, co-circular: drifts ahead", [-1, 0, 0, 1.5*n], 2, GREEN),
         ("(d) 1 km above with $\\dot y_0=-2nx_0$: closed", [1, 0, 0, -2*n], 1, PURPLE)]
fig, axs = plt.subplots(2, 2, figsize=(12.4, 7.4))
for ax, (title, X0, revs, col) in zip(axs.flat, cases):
    t = np.linspace(0, revs*T, 800)
    X = np.array([cw_stm(n, tt) @ X0 for tt in t])
    ax.plot(X[:, 1], X[:, 0], color=col, lw=2.6)
    for k in range(revs + 1):
        Xk = cw_stm(n, k*T) @ X0
        lab = "0T = 1T" if (revs == 1 and k == 0) else f"{k}T"
        if revs == 1 and k == 1:
            continue
        ax.plot(Xk[1], Xk[0], "o", color=col, ms=7, zorder=5)
        ax.annotate(lab, (Xk[1], Xk[0]), xytext=(6, -14), textcoords="offset points",
                    fontsize=9.5, color=col)
    ax.plot(0, 0, "*", color=INK, ms=14, zorder=6)
    ax.axhline(0, color=INK, lw=0.6); ax.axvline(0, color=INK, lw=0.6)
    ax.set_title(title, fontsize=11)
    ax.set_xlabel(r"along-track $y$ [km]  (direction of motion $\rightarrow$)")
    ax.set_ylabel(r"radial $x$ [km]")
    ax.set_aspect("equal", adjustable="datalim")
    despine(ax)
fig.suptitle(f"Target at 420 km: $n={n:.4e}$ rad/s, $T={T/60:.2f}$ min", fontsize=12)
fig.tight_layout()
fig.savefig("img/fig-cw-motions.png", dpi=160, bbox_inches="tight")
print("wrote img/fig-cw-motions.png")