"""Animation of a low-thrust spiral (lunar-phasing-and-low-thrust.qmd): the craft thrusts along its velocity and its speed falls.

The thrust is EXAGGERATED (kappa_1 = a_T r_1^2/mu = 0.003, about 130 times SMART-1's) so that the spiral has twelve turns instead of seventeen hundred; the shape
depends only on kappa_1.  The orbit is integrated (Cartesian, tangential thrust of constant magnitude); the right panels show the speed, which falls along the
straight line v_c = v_c1 - a_T t, and the radius.  Run from the sfm-slides directory:   python3 codes/anim-lpl-spiral.py
Afterwards, from the project root:  python3 _filters/gif-frames.py
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER
from _gif import save_gif
from _lunar import MU_E
from _lunarlt import R0, integrate_spiral, spiral_revolutions

eps1 = 0.003
aT = eps1*MU_E/R0**2
sol = integrate_spiral(R0, aT, 6*R0)
v1 = np.sqrt(MU_E/R0)
tau = v1/aT
T_END = sol.t[-1]
NF, NH = 170, 22
fr = np.concatenate([np.linspace(0, T_END, NF), np.full(NH, T_END)])

fig = plt.figure(figsize=(11.6, 5.8))
ax = fig.add_axes([0.02, 0.06, 0.46, 0.86])
ay = fig.add_axes([0.58, 0.55, 0.38, 0.36])
az = fig.add_axes([0.58, 0.10, 0.38, 0.33])
ax.set_aspect("equal"); ax.axis("off"); ax.set_xlim(-6.6, 6.6); ax.set_ylim(-6.6, 6.6)
th = np.linspace(0, 2*np.pi, 400)
ax.plot(np.cos(th), np.sin(th), color=BLUE, lw=1.3, ls=(0, (4, 3)))
ax.plot(6*np.cos(th), 6*np.sin(th), color=GREEN, lw=1.3, ls=(0, (4, 3)))
ax.add_patch(Circle((0, 0), 0.5, color=BLUE, alpha=0.4)); ax.plot(0, 0, "o", color=BLUE, ms=5)
trail, = ax.plot([], [], color=MAROON, lw=1.5)
dot, = ax.plot([], [], "o", color=INK, ms=9, mec=CREAM, mew=1.4, zorder=5)
arrow = ax.annotate("", xy=(0, 0), xytext=(0, 0), arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=3.0, mutation_scale=18))
txt = ax.text(-6.4, 6.0, "", fontsize=11, family="monospace", va="top")
ax.set_title("thrust (amber) along the velocity; thrust exaggerated", fontsize=11)

tt = np.linspace(0, T_END/tau, 300)
ay.plot(tt, 1 - tt, color=INK, lw=1.2, ls="--")
ay.set_xlim(0, T_END/tau*1.02); ay.set_ylim(0.35, 1.05)
ay.set_ylabel("$v/v_1$"); ay.set_title("speed falls on the straight line $v_c=v_{c1}-a_Tt$", fontsize=10.5)
ay.tick_params(labelbottom=False)
lv, = ay.plot([], [], color=BLUE, lw=2.6)
dv_, = ay.plot([], [], "o", color=BLUE, ms=8)
az.plot(tt, 1/(1 - tt)**2, color=INK, lw=1.0, ls=":")
az.set_xlim(0, T_END/tau*1.02); az.set_ylim(0.8, 6.4)
az.set_xlabel("time in units of $\\tau=v_1/a_T$"); az.set_ylabel("$r/r_1$")
lr, = az.plot([], [], color=MAROON, lw=2.6)
dr_, = az.plot([], [], "o", color=MAROON, ms=8)
for a_ in (ay, az):
    a_.spines["top"].set_visible(False); a_.spines["right"].set_visible(False)
state = {"th0": 0.0}


def frame(i):
    t = fr[min(i, len(fr) - 1)]
    s = np.linspace(0, t, max(2, int(900*t/T_END) + 2))
    Y = sol.sol(s)
    trail.set_data(Y[0]/R0, Y[1]/R0)
    p = sol.sol(t)
    dot.set_data([p[0]/R0], [p[1]/R0])
    v = p[2:]/np.hypot(*p[2:])
    arrow.xy = (p[0]/R0 + 1.1*v[0], p[1]/R0 + 1.1*v[1]); arrow.set_position((p[0]/R0, p[1]/R0))
    ang = np.unwrap(np.arctan2(Y[1], Y[0]))
    turns = (ang[-1] - ang[0])/(2*np.pi)
    vv = np.hypot(Y[2], Y[3])/v1
    lv.set_data(s/tau, vv); dv_.set_data([t/tau], [vv[-1]])
    rr = np.hypot(Y[0], Y[1])/R0
    lr.set_data(s/tau, rr); dr_.set_data([t/tau], [rr[-1]])
    txt.set_text(f"turns {turns:5.1f}\nr/r1  {rr[-1]:5.2f}\nv/v1  {vv[-1]:5.3f}")


frame(len(fr) - 1)
save_gif(fig, frame, len(fr), 20, "img/anim-lpl-spiral.gif", dpi=70, colors=96)