"""Q1: what a magnitude error and a pointing error do to a Hohmann departure burn.

Left: the velocity triangle at the burn point.  Right: the transfer ellipses that
result (errors exaggerated about ten times so that they can be seen).

Run from the sfm-slides directory:   python3 codes/fig-mq-burn-errors.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _mqcommon import mu, Re, vis, elements_2d, conic_xy

r1, r2 = Re + 300.0, 42164.0
a_t = (r1 + r2)/2
vc = np.sqrt(mu/r1)
dv = vis(r1, a_t) - vc                       # nominal tangential burn, 2.42573 km/s

eps, alpha = 0.08, np.radians(16.0)          # exaggerated
rvec = np.array([r1, 0.0])

def orbit(dv_t, dv_r):
    return elements_2d(rvec, np.array([dv_r, vc + dv_t]))

nom = orbit(dv, 0.0)
mag = orbit(dv*(1 + eps), 0.0)
ptr = orbit(dv*np.cos(alpha), dv*np.sin(alpha))

fig, (axL, axR) = plt.subplots(1, 2, figsize=(14.4, 6.2), gridspec_kw=dict(width_ratios=[0.8, 1.2]))

# ---- A : the velocity change at the burn point ----------------------------------
def arrow(ax, p, q, col, lw=2.8):
    ax.annotate("", xy=q, xytext=p, arrowprops=dict(arrowstyle="-|>", color=col, lw=lw, shrinkA=0, shrinkB=0))

O = np.array([0.0, 0.0])
arrow(axL, O, np.array([dv, 0]), GREEN)
axL.text(dv/2, -0.30, r"commanded: $\Delta v=2.426$ km/s, tangential" "\n"
         r"(the circular speed $v_c=7.726$ km/s is along the same line)", color=GREEN, ha="center", va="top", fontsize=11)
act = dv*(1 + eps)*np.array([np.cos(alpha), np.sin(alpha)])
arrow(axL, O, act, MAROON)
axL.text(act[0] + 0.1, act[1] + 0.08, r"executed: $(1+\varepsilon)\,\Delta v$," "\n" r"turned by $\alpha$ from tangential",
         color=MAROON, fontsize=11, va="bottom")
axL.plot([0, act[0]], [0, 0], color=MAROON, lw=1.0, ls=":")
axL.plot([act[0]]*2, [0, act[1]], color=MAROON, lw=1.0, ls=":")
axL.text(act[0] + 0.1, act[1]/2, r"$\Delta v\sin\alpha$ (radial):" "\n" "rotates the apse line", color=MAROON, fontsize=10, va="center")
axL.text(act[0]/2, 0.14, r"$\Delta v\cos\alpha$ (tangential): changes $a$", color=MAROON, fontsize=10, ha="center", va="bottom", alpha=0.0)
th = np.linspace(0, alpha, 40)
axL.plot(1.3*np.cos(th), 1.3*np.sin(th), color=MAROON, lw=1.4)
axL.text(1.55, 0.2, r"$\alpha$", color=MAROON, fontsize=13)
axL.set_xlim(-0.4, 6.6); axL.set_ylim(-1.5, 1.8)
axL.set_aspect("equal"); axL.axis("off")
axL.set_title(r"The velocity change at the burn point ($r_1=6678$ km)", fontsize=12)

# ---- B : the resulting orbits ---------------------------------------------------
ang = np.linspace(0, 2*np.pi, 400)
axR.fill(Re*np.cos(ang), Re*np.sin(ang), color=BLUE, alpha=0.18, zorder=1)
axR.plot(r2*np.cos(ang), r2*np.sin(ang), color=INK, lw=1.0, ls=(0, (5, 4)), alpha=0.7)
axR.text(r2*0.72, r2*0.80, "GEO\n42 164 km", color=INK, ha="left", fontsize=10, alpha=0.8)
for (a, e, w), col, lab in ((nom, BLUE, "nominal"),
                            (mag, AMBER, r"magnitude error $\varepsilon$: apoapsis moves out"),
                            (ptr, MAROON, r"pointing error $\alpha$: the apse line rotates")):
    x, y = conic_xy(a, e, w)
    axR.plot(x, y, color=col, lw=2.3, label=lab)
    ra = a*(1 + e)
    axR.plot(ra*np.cos(w + np.pi), ra*np.sin(w + np.pi), "o", color=col, ms=8, zorder=5)
axR.plot(r1, 0, "o", color=INK, ms=8, zorder=6)
axR.annotate("burn point", (r1, 0), (r1 + 3000, 7000), fontsize=10, color=INK, arrowprops=dict(arrowstyle="-", color=INK, lw=0.8))
axR.set_aspect("equal"); axR.axis("off")
axR.set_xlim(-84000, 46000); axR.set_ylim(-46000, 46000)
axR.legend(frameon=False, fontsize=10.5, loc="lower left")
axR.set_title("Resulting transfer ellipses (errors drawn about ten times too large)", fontsize=12)
axR.text(-83000, 45000,
         r"to first order:  $\delta r_a = 2\,(2a^2v/\mu)\,\delta(\Delta v)\ \approx\ 61$ km per m/s" "\n"
         r"$\delta\omega \approx -\alpha\,(\Delta v/v_t)(1+e)/e\ \approx\ -0.57\,\alpha$",
         fontsize=10.5, color=INK, va="top")

plt.tight_layout()
fig.savefig("img/fig-mq-burn-errors.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-mq-burn-errors.png")