"""Plane change: the velocity triangle, and how to split it across two burns."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

mu, RE = 398600.4418, 6378.137

figL, axL = plt.subplots(figsize=(7.4, 5.8))
figM, axM = plt.subplots(figsize=(9.0, 5.8))
figR, axR = plt.subplots(figsize=(9.0, 5.8))

# ---- panel A : the isosceles velocity triangle ----------------------------
th = np.radians(40.0)
v = 1.0
axL.annotate("", xy=(v, 0), xytext=(0, 0),
             arrowprops=dict(arrowstyle="-|>", color=BLUE, lw=2.8))
axL.annotate("", xy=(v*np.cos(th), v*np.sin(th)), xytext=(0, 0),
             arrowprops=dict(arrowstyle="-|>", color=GREEN, lw=2.8))
axL.annotate("", xy=(v*np.cos(th), v*np.sin(th)), xytext=(v, 0),
             arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.8))
axL.text(v*0.52, -0.11, r"$\mathbf{v}_1$", color=BLUE, fontsize=15)
axL.text(v*0.44*np.cos(th) - 0.10, v*0.52*np.sin(th) + 0.03, r"$\mathbf{v}_2$", color=GREEN, fontsize=15)
axL.text(1.02, 0.20, r"$\Delta\mathbf{v}$", color=AMBER, fontsize=15)
arc = np.linspace(0, th, 50)
axL.plot(0.24*np.cos(arc), 0.24*np.sin(arc), color=INK, lw=1.4)
axL.text(0.30*np.cos(th/2), 0.30*np.sin(th/2) - 0.02, r"$\theta$", color=INK, fontsize=14)
axL.text(-0.18, 0.86, r"$|\mathbf{v}_1|=|\mathbf{v}_2|=v \;\Rightarrow\; "
                      r"\Delta v = 2v\sin\dfrac{\theta}{2}$",
         color=INK, fontsize=13, ha="left")
axL.text(0.5, -0.30, r"$\theta=60^\circ$ already costs a full $v$;"
                     "\n" r"$\theta=90^\circ$ costs $\sqrt{2}\,v$.",
         color=MAROON, fontsize=10.5, ha="center")
axL.set_aspect("equal"); axL.axis("off")
axL.set_xlim(-0.25, 1.45); axL.set_ylim(-0.45, 1.12)
axL.set_title("A pure rotation of $\\mathbf{v}$ is isosceles", fontsize=12)

# ---- panel B : cost of a plane change versus where you do it -------------
rr = np.linspace(6678.137, 60000, 500)
for d, col in ((28.5, MAROON), (10.0, BLUE), (5.0, GREEN)):
    axM.plot(rr, 2*np.sqrt(mu/rr)*np.sin(np.radians(d)/2), color=col, lw=2.4,
             label=fr"$\Delta i={d}^\circ$")
for r0, lab in ((6678.137, "LEO"), (42164.0, "GEO")):
    axM.axvline(r0, color=INK, lw=0.9, ls=":")
    axM.text(r0 + 700, 3.55, lab, color=INK, fontsize=10)
a_gto = 0.5*(6678.137 + 42164.0)
v_apo = np.sqrt(mu*(2/42164.0 - 1/a_gto))
axM.plot(42164.0, 2*v_apo*np.sin(np.radians(28.5)/2), "o", color=AMBER, ms=9)
axM.annotate("at GTO apoapsis\n(v = 1.608 km/s):  0.792 km/s",
             xy=(42164.0, 2*v_apo*np.sin(np.radians(28.5)/2)),
             xytext=(13500, 1.30), color=AMBER, fontsize=10,
             arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=1.3))
axM.set_xlabel("radius at which the plane change is made  [km]")
axM.set_ylabel(r"$\Delta v = 2v\sin(\Delta i/2)$   [km/s]")
axM.set_xlim(6000, 60000); axM.set_ylim(0, 4.0)
axM.legend(frameon=False, fontsize=10.5)
axM.set_title("Turn where you are slow", fontsize=12)
despine(axM)

# ---- panel C : optimal split of the plane change, LEO -> GEO -------------
r1, r2, di = 6678.137, 42164.0, np.radians(28.5)
at = 0.5*(r1 + r2)
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))
d = np.linspace(0, di, 700)
A = np.sqrt(vc1**2 + vt1**2 - 2*vc1*vt1*np.cos(d))
B = np.sqrt(vt2**2 + vc2**2 - 2*vt2*vc2*np.cos(di - d))
axR.plot(np.degrees(d), A, color=BLUE, lw=2.0, label=r"$\Delta v_1$ (at LEO)")
axR.plot(np.degrees(d), B, color=GREEN, lw=2.0, label=r"$\Delta v_2$ (at GEO)")
axR.plot(np.degrees(d), A + B, color=MAROON, lw=2.8, label="total")
i0 = int(np.argmin(A + B))
axR.plot(np.degrees(d[i0]), (A + B)[i0], "o", color=MAROON, ms=9)
axR.annotate(f"optimum {np.degrees(d[i0]):.2f}°:  {(A+B)[i0]:.4f} km/s",
             xy=(np.degrees(d[i0]), (A + B)[i0]), xytext=(8.2, 5.75),
             color=MAROON, fontsize=10.5,
             arrowprops=dict(arrowstyle="-|>", color=MAROON, lw=1.3))
axR.plot(0, (A + B)[0], "s", color=INK, ms=7)
axR.text(1.2, 3.42, f"all at GEO:\n{(A+B)[0]:.4f} km/s", color=INK, fontsize=10)
axR.set_xlabel(r"plane change taken at the first burn,  $\delta$  [deg]")
axR.set_ylabel(r"$\Delta v$   [km/s]")
axR.set_xlim(0, 28.5); axR.set_ylim(1.2, 6.9)
axR.legend(frameon=False, fontsize=10, loc="center right")
axR.set_title("LEO$\\to$GEO with $\\Delta i=28.5^\\circ$: split it", fontsize=12)
despine(axR)

for f, name in ((figL, "fig-man-plane-change-geom"),
                (figM, "fig-man-plane-change-cost"),
                (figR, "fig-man-plane-change-split")):
    f.tight_layout()
    f.savefig(f"img/{name}.png", dpi=160, bbox_inches="tight")
    plt.close(f)
    print(f"wrote img/{name}.png")