"""Figure for plane-changes.qmd (img/fig-man-plane-triangles.png): velocity triangles, to scale, for the
LEO -> GEO transfer with a 19.2 deg plane change (from a Sriharikota transfer orbit).

(a) Separate: circularise at GEO, then turn the circular velocity by 19.2 deg (two burns).
(c) Combined: one burn at GEO takes v_t2 straight to the turned v_c2.
(d) Split: 1.64 deg of the turn in the LEO burn, the remaining 17.56 deg in the GEO burn.
In each panel the vectors are drawn in the plane containing them, horizontal = the transfer orbit's
direction of motion; km/s, same scale within a panel.

Run from the sfm-slides directory:   python3 codes/fig-man-plane-triangles.py
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from scipy.optimize import minimize_scalar
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE

MU = 398600.4418
r1, r2 = 6678.137, 42164.0
at = (r1 + r2) / 2
vc1, vc2 = np.sqrt(MU / r1), np.sqrt(MU / r2)
vt1, vt2 = np.sqrt(MU * (2 / r1 - 1 / at)), np.sqrt(MU * (2 / r2 - 1 / at))
di = np.radians(19.2)
lc = lambda x, y, th: np.sqrt(x * x + y * y - 2 * x * y * np.cos(th))
dl = minimize_scalar(lambda d: lc(vc1, vt1, d) + lc(vt2, vc2, di - d), bounds=(0, di), method="bounded",
                     options={"xatol": 1e-12}).x
sep = (vc2 - vt2) + 2 * vc2 * np.sin(di / 2)
comb = lc(vt2, vc2, di)
spl1, spl2 = lc(vc1, vt1, dl), lc(vt2, vc2, di - dl)
print(f"separate {vc2 - vt2:.5f} + {2 * vc2 * np.sin(di / 2):.5f}; combined {comb:.5f}; split delta {np.degrees(dl):.2f}: {spl1:.5f} + {spl2:.5f}")

u = lambda th: np.array([np.cos(th), np.sin(th)])


def arrow(ax, p, q, col, lw=2.4):
    ax.annotate("", xy=q, xytext=p, arrowprops=dict(arrowstyle="-|>", color=col, lw=lw, shrinkA=0, shrinkB=0))


fig, axs = plt.subplots(1, 4, figsize=(16, 4.6), gridspec_kw=dict(width_ratios=[1, 1, 1.25, 1]))
O = np.zeros(2)
# (a) separate, at GEO
ax = axs[0]
arrow(ax, O, vt2 * u(0), BLUE); arrow(ax, O, vc2 * u(0), GREEN, 1.4); arrow(ax, O, vc2 * u(di), GREEN)
arrow(ax, vt2 * u(0), vc2 * u(0), MAROON, 3); arrow(ax, vc2 * u(0), vc2 * u(di), MAROON, 3)
ax.text(vt2 / 2, -0.32, r"$v_{t2}=%.3f$" % vt2, color=BLUE, ha="center")
ax.text(vc2 * 0.5, vc2 * 0.28, r"$v_{c2}$, turned", color=GREEN, ha="right")
ax.text(vc2 + 0.08, 0.42, "%.3f" % (2 * vc2 * np.sin(di / 2)), color=MAROON)
ax.text((vt2 + vc2) / 2, -0.32, "%.3f" % (vc2 - vt2), color=MAROON, ha="center")
ax.set_title(f"(a) separate, at GEO:\n{sep:.3f} km/s", fontsize=11.5)
# (c) combined, at GEO
ax = axs[1]
arrow(ax, O, vt2 * u(0), BLUE); arrow(ax, O, vc2 * u(di), GREEN); arrow(ax, vt2 * u(0), vc2 * u(di), MAROON, 3)
ax.text(vt2 / 2, -0.25, r"$v_{t2}$", color=BLUE, ha="center")
ax.text(vc2 * 0.5, vc2 * 0.28, r"$v_{c2}$, turned $19.2^\circ$", color=GREEN, ha="right")
ax.text(2.42, 0.32, "%.3f" % comb, color=MAROON)
ax.set_title(f"(c) combined, at GEO:\n{comb:.3f} km/s", fontsize=11.5)
# (d) split: the LEO burn, turned by delta
ax = axs[2]
arrow(ax, O, vc1 * u(0), BLUE); arrow(ax, O, vt1 * u(dl), GREEN); arrow(ax, vc1 * u(0), vt1 * u(dl), MAROON, 3)
ax.text(vc1 / 2, -0.75, r"$v_{c1}=%.3f$" % vc1, color=BLUE, ha="center")
ax.text(vt1 * 0.55, 0.65, r"$v_{t1}=%.3f$, turned $%.2f^\circ$" % (vt1, np.degrees(dl)), color=GREEN, ha="center")
ax.text(vc1 + 0.4, -0.75, "%.3f" % spl1, color=MAROON)
ax.set_title(f"(d) split, at LEO: {spl1:.3f} km/s", fontsize=11.5)
# (d) split: the GEO burn
ax = axs[3]
arrow(ax, O, vt2 * u(0), BLUE); arrow(ax, O, vc2 * u(di - dl), GREEN); arrow(ax, vt2 * u(0), vc2 * u(di - dl), MAROON, 3)
ax.text(vt2 / 2, -0.25, r"$v_{t2}$", color=BLUE, ha="center")
ax.text(vc2 * 0.5, vc2 * 0.26, r"$v_{c2}$, turned $%.2f^\circ$" % np.degrees(di - dl), color=GREEN, ha="right")
ax.text(2.42, 0.3, "%.3f" % spl2, color=MAROON)
ax.set_title(f"... and at GEO: {spl2:.3f} km/s\ntotal {spl1 + spl2:.3f} km/s", fontsize=11.5)
for ax in axs:
    ax.set_aspect("equal"); ax.axis("off")
for ax in (axs[0], axs[1], axs[3]):
    ax.set_xlim(-0.6, 3.4); ax.set_ylim(-0.5, 1.4)
axs[2].set_xlim(-0.5, 11.2); axs[2].set_ylim(-1.3, 1.4)
fig.suptitle("LEO to GEO with a $19.2^\\circ$ plane change: the burns as velocity triangles (km/s, to scale within each panel;"
             "\n(a) and (c) also pay the tangential LEO burn, $2.426$ km/s)", fontsize=11.5, y=0.97)
fig.tight_layout()
fig.savefig("img/fig-man-plane-triangles.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-man-plane-triangles.png")