"""CW against the exact two-body relative motion, for growing separations."""
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lambertp import mu, RE, cw_stm, two_body_rhs

a = RE + 420.0
n = np.sqrt(mu/a**3); T = 2*np.pi/n
t = np.linspace(0, 3*T, 900)

def exact(X0):
    """Integrate target and chaser inertially, express the offset in the LVLH frame."""
    x, y, xd, yd = X0
    rt, vt = np.array([a, 0]), np.array([0, a*n])
    rc = rt + np.array([x, y])
    vc = vt + np.array([xd, yd]) + n*np.array([-y, x])      # + omega x rho
    st = solve_ivp(two_body_rhs, [0, t[-1]], [*rt, *vt], t_eval=t, rtol=1e-12, atol=1e-10)
    sc = solve_ivp(two_body_rhs, [0, t[-1]], [*rc, *vc], t_eval=t, rtol=1e-12, atol=1e-10)
    out = []
    for k, tt in enumerate(t):
        c, s = np.cos(n*tt), np.sin(n*tt)
        d = sc.y[:2, k] - st.y[:2, k]
        out.append([c*d[0] + s*d[1], -s*d[0] + c*d[1]])
    return np.array(out)

fig, (axL, axR) = plt.subplots(1, 2, figsize=(12.6, 5.0))
for d, col in ((1, GREEN), (10, BLUE), (100, PURPLE), (1000, MAROON)):
    X0 = np.array([0, 0, n*d/4, 0])                         # radial kick: 2:1 ellipse, extent d
    cw = np.array([cw_stm(n, tt) @ X0 for tt in t])[:, :2]
    ex = exact(X0)
    axL.plot(ex[:, 1]/d, ex[:, 0]/d, color=col, lw=2.2, label=f"exact, extent {d} km")
    err = np.linalg.norm(ex - cw, axis=1)/d
    axR.semilogy(t/T, np.maximum(err, 1e-12), color=col, lw=2.4, label=f"{d} km")
cw1 = np.array([cw_stm(n, tt) @ [0, 0, n/4, 0] for tt in t])
axL.plot(cw1[:, 1], cw1[:, 0], color=INK, lw=1.3, ls=(0, (5, 3)), label="CW (any size)")
axL.set_xlabel(r"$y/d$"); axL.set_ylabel(r"$x/d$")
axL.set_title("Three orbits after a radial kick, scaled by the size $d$", fontsize=11.5)
axL.set_aspect("equal", adjustable="datalim")
axL.set_ylim(-0.75, 0.45)
axL.legend(frameon=False, fontsize=9.5, loc="lower center", ncol=2)
axR.set_xlabel("time  [orbits]"); axR.set_ylabel(r"$|\rho_{\rm exact}-\rho_{\rm CW}|\,/\,d$")
axR.set_title(r"Relative error grows like $d/a$: ten times farther, ten times worse", fontsize=11.5)
axR.legend(frameon=False, fontsize=10, title="extent $d$")
axR.set_ylim(1e-5, 1)
for ax in (axL, axR):
    despine(ax)
fig.tight_layout()
fig.savefig("img/fig-cw-validity.png", dpi=160, bbox_inches="tight")
print("wrote img/fig-cw-validity.png")