"""Figure for complications-to-impulsive-transfers-chobotov.qmd (img/fig-ci-gravity.png): gravity losses and what to do about them.

(a) A vertical ascent with constant gravity and constant thrust: the speed actually gained against the thrust-to-weight ratio, mass ratio 3.
    The loss is g t_b = u (1 - m_f/m_0) / (T/W_0): it falls as 1/(T/W_0).
(b) Raising the apogee of a circular 278 km orbit to the geosynchronous radius with a finite-thrust, prograde burn, split into N burns at successive
    perigee passages: the loss against the best impulsive transfer to the same final orbit, for three thrust levels.
(c) The price of splitting: the elapsed time. The numbers are cached in codes/_cache/ci-burns.json (written by this script or by
    complications-verify.py, whichever runs first).

Run from the sfm-slides directory:   python3 codes/fig-ci-gravity.py
"""
import json
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _complic import *

import os
CACHE = "codes/_cache/ci-burns.json"
if os.path.exists(CACHE):
    data = json.load(open(CACHE))
else:                                    # computed here if complications-verify.py has not been run yet (about ten seconds)
    r1_, rg_ = RE + 278, RE + 35786
    data = {f"{tw}_{N}": list(perigee_burns(r1_, rg_, N, tw, centred=True)) for tw in (0.1, 0.03, 0.01) for N in (1, 2, 3, 4, 6, 8)}
    os.makedirs("codes/_cache", exist_ok=True); json.dump(data, open(CACHE, "w"))
fig, ax = plt.subplots(1, 3, figsize=(15.6, 5.0))
A = ax[0]
tw = np.logspace(-0.1, 1.2, 200); u = 300 * 9.80665; R = 3.0
loss = u * (1 - 1 / R) / tw; ideal = u * np.log(R)
A.plot(tw, ideal - loss, color=BLUE, lw=2.4, label="speed actually gained")
A.axhline(ideal, color=INK, lw=1.0, ls="--"); A.text(0.9, ideal + 60, r"ideal: $u\ln(m_0/m_f)=%.0f$ m/s" % ideal, fontsize=10)
A.fill_between(tw, ideal - loss, ideal, color="#f3dccc"); A.text(2.0, ideal - 300, "gravity loss", color=AMBER, fontsize=11)
A.set_xscale("log"); A.set_xlabel(r"initial thrust-to-weight ratio $T/W_0$"); A.set_ylabel("speed gained [m/s]"); A.set_ylim(0, 3500); A.set_xlim(0.8, 16)
A.legend(loc="lower right", frameon=False, fontsize=10); A.set_title("(a)  vertical ascent: the loss is $g\\,t_b$", fontsize=11.5, loc="left"); despine(A)
Ns = [1, 2, 3, 4, 6, 8]
cols = {0.1: BLUE, 0.03: GREEN, 0.01: MAROON}
B = ax[1]; C = ax[2]
for twv, c in cols.items():
    L = [(data[f"{twv}_{N}"][0] - data[f"{twv}_{N}"][1]) * 1000 for N in Ns]
    T = [data[f"{twv}_{N}"][4] / 3600 for N in Ns]
    B.plot(Ns, L, "o-", color=c, lw=2.2, label=f"$T/W_0={twv}$"); C.plot(Ns, T, "o-", color=c, lw=2.2, label=f"$T/W_0={twv}$")
B.set_yscale("log"); B.set_xlabel("number of perigee burns"); B.set_ylabel("loss against the impulsive transfer [m/s]"); B.set_ylim(1, 1000)
B.legend(frameon=False, fontsize=10); B.set_title("(b)  finite-burn loss falls as the burn is split", fontsize=11.5, loc="left"); despine(B)
C.set_xlabel("number of perigee burns"); C.set_ylabel("elapsed time [h]"); C.legend(frameon=False, fontsize=10)
C.set_title("(c)  and the elapsed time grows", fontsize=11.5, loc="left"); despine(C)
fig.tight_layout(); fig.savefig("img/fig-ci-gravity.png", dpi=150, bbox_inches="tight")
print("ok")