"""Figures of the low-thrust part of lunar-phasing-and-low-thrust.qmd.

fig-lpl-spiral       the shape of a spiral under tangential thrust (thrust exaggerated: kappa_1 = 0.003), and the straight-line fall of the speed
fig-lpl-spiral-real  a thrust like SMART-1's (a_T = 2.08e-7 km/s^2): integration against r(t) = mu/(v_c1 - a_T t)^2, the osculating eccentricity against 4 eps,
                     and where the Earth-only model stops
fig-lpl-isp-trade    propellant fraction and trip time against specific impulse at a fixed power

Run from the sfm-slides directory:   python3 codes/fig-lpl-lowthrust.py     (about 30 s)
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from scipy.optimize import brentq
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lunar import MU_E, MU_M, D, R_SOI
from _lunarlt import *

MOONC = "#6b655b"

# ------------------------------------------------------------------ the shape of the spiral and the fall of the speed (kappa_1 = 0.003)
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
ts = np.linspace(0, sol.t[-1], 6000)
Y = sol.sol(ts)
r = np.hypot(Y[0], Y[1])/R0
fig, (ax, ay) = plt.subplots(1, 2, figsize=(13.2, 6.0), gridspec_kw=dict(width_ratios=[1.0, 1.05]))
ax.plot(Y[0]/R0, Y[1]/R0, color=MAROON, lw=1.1)
th = np.linspace(0, 2*np.pi, 400)
ax.plot(np.cos(th), np.sin(th), color=BLUE, lw=1.6, ls=(0, (4, 3)))
ax.plot(6*np.cos(th), 6*np.sin(th), color=GREEN, lw=1.6, ls=(0, (4, 3)))
ax.plot(0, 0, "o", color=BLUE, ms=8)
for tf in (0.15, 0.4, 0.62, 0.85):
    i = int(tf*len(ts)) - 1
    p, v = Y[:2, i]/R0, Y[2:, i]/np.hypot(*Y[2:, i])
    ax.annotate("", xy=p + 0.9*v, xytext=p, arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.4))
ax.annotate("start: circle, $r_1$", xy=(0.72, -0.72), xytext=(-6.3, -5.6), color=BLUE, fontsize=10.5, arrowprops=dict(arrowstyle="-", color=BLUE, lw=0.9))
ax.text(3.9, -5.5, "end: $r=6\\,r_1$", color=GREEN, fontsize=10.5)
ax.text(-5.8, 5.2, "thrust $a_T$ along\nthe velocity (arrows)", color=AMBER, fontsize=10.5, va="top")
ax.set_aspect("equal"); ax.set_xlim(-6.6, 6.6); ax.set_ylim(-6.6, 6.6)
ax.set_xlabel("[$r_1$]")
ax.set_title(f"A spiral under tangential thrust ($\\kappa_1={eps1}$, thrust exaggerated: {spiral_revolutions(R0, 6*R0, aT):.1f} turns)", fontsize=11.5)
despine(ax)
vv = np.hypot(Y[2], Y[3])/v1
ay.plot(ts/tau, vv, color=BLUE, lw=2.4, label="speed $v/v_1$ (integrated)")
tt = np.linspace(0, ts[-1]/tau, 300)
ay.plot(tt, 1 - tt, color=INK, lw=1.2, ls="--", label="$v_c/v_1=1-a_Tt/v_1$ (the theorem)")
ay.set_xlabel("time in units of $\\tau=v_1/a_T$"); ay.set_ylabel("$v/v_1$", color=BLUE)
ay.set_ylim(0.35, 1.05); ay.set_xlim(0, ts[-1]/tau*1.02)
az = ay.twinx()
az.plot(ts/tau, r, color=MAROON, lw=2.4, label="radius $r/r_1$ (integrated)")
az.plot(tt, 1/(1 - tt)**2, color=INK, lw=1.0, ls=":", label="$r/r_1=(1-a_Tt/v_1)^{-2}$")
az.set_ylabel("$r/r_1$", color=MAROON); az.set_ylim(0.8, 6.4)
h1, l1 = ay.get_legend_handles_labels(); h2, l2 = az.get_legend_handles_labels()
ay.legend(h1 + h2, l1 + l2, frameon=False, fontsize=10, loc="center left")
ay.set_title("Thrusting forward slows the craft down: $dv_c/dt=-a_T$", fontsize=11.5)
ay.spines["top"].set_visible(False); az.spines["top"].set_visible(False)
fig.tight_layout()
fig.savefig("img/fig-lpl-spiral.png", dpi=150, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-lpl-spiral.png")

# ------------------------------------------------------------------ a thrust like SMART-1's: integration against theory, and where the model stops
aT = 2.08e-7
r_end = 1.0e5
sol = integrate_spiral(R0, aT, r_end)
ts = np.linspace(0, sol.t[-1], 120001)
Y = sol.sol(ts)
rn = np.hypot(Y[0], Y[1])
tan = np.linspace(0, spiral_time(R0, D*0.999, aT), 800)
fig, (ax, ay) = plt.subplots(1, 2, figsize=(13.4, 5.9))
ax.plot(tan/DAY, spiral_r_of_t(tan, R0, aT)/1e3, color=INK, lw=1.4, ls="--", label="$r=\\mu/(v_{c1}-a_Tt)^2$")
ax.plot(ts/DAY, rn/1e3, color=MAROON, lw=2.8, label="integration (Earth only)")
r_lim = np.sqrt(0.01*MU_E/aT)
t_lim = spiral_time(R0, r_lim, aT)
ax.axhline(D/1e3, color=MOONC, lw=1.0, ls=":"); ax.text(150, D/1e3 + 6, "the Moon's orbit", color=MOONC, fontsize=10)
ax.axhline(r_lim/1e3, color=AMBER, lw=1.0, ls="--")
ax.text(8, r_lim/1e3 + 6, f"$\\kappa=0.01$ at {r_lim/1e3:,.0f} thousand km", color=AMBER, fontsize=10)
ax.axhline(1e2, color=GREEN, lw=1.0, ls="--")
ax.text(8, 100 + 6, f"end of the integration, 100 thousand km (day {sol.t[-1]/DAY:.0f}); the Moon's tide is {tidal_ratio(1e5)/aT:.2f} of the thrust here", color=GREEN, fontsize=9.5)
ax.axvspan(sol.t[-1]/DAY, tan[-1]/DAY, color=MOONC, alpha=0.08)
ax.text((sol.t[-1]/DAY + tan[-1]/DAY)/2, 250, "Earth-only\nmodel fails:\nthe Moon\nenters", color=MOONC, fontsize=10, ha="center")
ax.set_xlabel("time [days]"); ax.set_ylabel("radius of the spiral [1000 km]")
ax.set_xlim(0, tan[-1]/DAY*1.01); ax.set_ylim(0, 420)
ax.legend(frameon=False, fontsize=10, loc="upper left", bbox_to_anchor=(0.0, 0.93))
ax.set_title("$a_T=2.08\\times10^{-7}$ km/s$^2$ (SMART-1's, from ESA's figures)", fontsize=11.5)
despine(ax)
idx = np.arange(0, len(ts), 25)
e_os = np.array([osculating_e(Y[:, i]) for i in idx])
rr = rn[idx]
ay.loglog(rr/1e3, e_os, ".", color=BLUE, ms=1.6, alpha=0.5, label="osculating $e$ (integrated)")
rg = np.logspace(np.log10(R0), np.log10(r_end), 200)
ay.loglog(rg/1e3, 4*aT*rg**2/MU_E, color=MAROON, lw=2.2, label="$4\\kappa$")
ay.loglog(rg/1e3, 2*aT*rg**2/MU_E, color=GREEN, lw=1.6, ls="--", label="$2\\kappa$")
k1_ = aT*R0**2/MU_E
ay.loglog(rg/1e3, 2*aT*rg**2/MU_E + 2*k1_, color=GREEN, lw=0.9, ls=":", label="$2\\kappa\\pm2\\kappa_1$")
ay.loglog(rg[1:]/1e3, np.maximum(2*aT*rg[1:]**2/MU_E - 2*k1_, 1e-9), color=GREEN, lw=0.9, ls=":")
ay.set_xlabel("radius [1000 km]"); ay.set_ylabel("eccentricity")
ay.set_xlim(R0/1e3*0.95, r_end/1e3*1.05); ay.set_ylim(2e-5, 0.1)
ay.legend(frameon=False, fontsize=10, loc="upper left")
ay.set_title("Eccentricity: on $2\\kappa=2a_Tr^2/\\mu$ within $2\\kappa_1$, never above $4\\kappa$", fontsize=11.5)
despine(ay)
fig.tight_layout()
fig.savefig("img/fig-lpl-spiral-real.png", dpi=150, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-lpl-spiral-real.png")

# ------------------------------------------------------------------ the trade against specific impulse, at fixed power
P_w, eta, m0, dv = 5000.0, 0.6, 1000.0, spiral_dv(R0, 1e5)
isp = np.linspace(300, 6000, 400)
frac = 100*propellant_fraction(dv, isp)
tdays = np.array([trip_time_constant_thrust(m0, thrust_from_power(P_w, i, eta), i, dv)/DAY for i in isp])
fig, ax = plt.subplots(figsize=(10.2, 5.8))
ax.plot(isp, frac, color=BLUE, lw=2.8)
ax.set_xlabel("specific impulse $I_{sp}$ [s]"); ax.set_ylabel("propellant, per cent of the initial mass", color=BLUE)
ax.set_ylim(0, 100); ax.set_xlim(0, 6000)
az = ax.twinx()
az.plot(isp, tdays, color=MAROON, lw=2.8)
az.set_ylabel("trip time [days]", color=MAROON); az.set_ylim(0, 700)
i18 = 1800
ax.plot(i18, 100*propellant_fraction(dv, i18), "o", color=BLUE, ms=9)
az.plot(i18, trip_time_constant_thrust(m0, thrust_from_power(P_w, i18, eta), i18, dv)/DAY, "o", color=MAROON, ms=9)
ax.annotate(f"$I_{{sp}}=1800$ s: {100*propellant_fraction(dv, i18):.1f}% propellant,\n{trip_time_constant_thrust(m0, thrust_from_power(P_w, i18, eta), i18, dv)/DAY:.0f} d, thrust {thrust_from_power(P_w, i18, eta):.2f} N",
            xy=(i18, 100*propellant_fraction(dv, i18)), xytext=(2400, 62), fontsize=10.5, arrowprops=dict(arrowstyle="-|>", color=INK, lw=1.2))
ax.text(300, 93, f"a chemical stage ($I_{{sp}}=320$ s) needs {100*propellant_fraction(spiral_dv(R0, 1e5)*0 + 4.1427, 320):.0f}% for the Hohmann transfer to the same circle,\nbut takes 17 h", fontsize=9.5, color=INK, va="top")
ax.set_title("LEO to a circle at 100 000 km, 5 kW to the jet, efficiency 0.6, 1000 kg: raise $I_{sp}$ and the propellant falls, the trip lengthens", fontsize=10.5)
despine(ax); az.spines["top"].set_visible(False)
fig.tight_layout()
fig.savefig("img/fig-lpl-isp-trade.png", dpi=150, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-lpl-isp-trade.png")