"""Three number-free sketches for the apse-rotation section: the two orbits, the velocity reflection
at a crossing, and the two-burn route through apoapsis. Run from sfm-slides/."""
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

# ---- the two orbits, with no result drawn on them --------------------------
fig, ax = plt.subplots(figsize=(6.6, 6.2))
p0, e0, eta0 = 0.84, 0.4, np.radians(60.0)
th0 = np.linspace(0, 2*np.pi, 800)
r0 = p0/(1 + e0*np.cos(th0))
ax.plot(r0*np.cos(th0), r0*np.sin(th0), color=BLUE, lw=2.4)
ax.plot(r0*np.cos(th0 + eta0), r0*np.sin(th0 + eta0), color=GREEN, lw=2.4)
ax.plot([0, 0.92], [0, 0], color=BLUE, lw=1.0, ls=":")
ax.plot([0, 0.92*np.cos(eta0)], [0, 0.92*np.sin(eta0)], color=GREEN, lw=1.0, ls=":")
arc0 = np.linspace(0, eta0, 60)
ax.plot(0.3*np.cos(arc0), 0.3*np.sin(arc0), color=INK, lw=1.3)
ax.text(0.37*np.cos(eta0/2), 0.37*np.sin(eta0/2), r"$\eta$", fontsize=15, ha="center", va="center")
ax.plot(0, 0, "o", color=INK, ms=8)
ax.text(0.03, -0.05, "focus", fontsize=10.5, va="top")
ax.text(-1.5, 0.92, "initial orbit", color=BLUE, fontsize=12.5)
ax.text(0.15, -1.32, "final orbit", color=GREEN, fontsize=12.5)
ax.set_aspect("equal"); ax.axis("off")
fig.tight_layout()
fig.savefig("img/fig-man-apse-orbits.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-man-apse-orbits.png")

# ---- the reflection of the velocity in the local horizontal line ----------
fig, ax = plt.subplots(figsize=(6.4, 4.8))
gam = np.radians(26.0)                        # exaggerated for the picture
vx, vy = np.cos(gam), np.sin(gam)


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


ax.plot([-0.12, 1.32], [0, 0], color=INK, lw=1.0, ls="--")
ax.text(1.33, -0.025, "local horizontal", ha="right", va="top", fontsize=10.5, color=INK)
arrow(ax, (0, 0), (0, 0.62), INK, lw=1.3)
ax.text(-0.03, 0.62, r"$\hat{\mathbf{r}}$", ha="right", va="top", fontsize=13)
arrow(ax, (0, 0), (vx, vy), BLUE)
arrow(ax, (0, 0), (vx, -vy), GREEN)
arrow(ax, (vx, vy), (vx, -vy), AMBER)
ax.text(0.5*vx, 0.5*vy + 0.05, r"$\mathbf{v}$", color=BLUE, fontsize=14, ha="right", va="bottom")
ax.text(0.5*vx, -0.5*vy - 0.05, r"$\mathbf{v}'$", color=GREEN, fontsize=14, ha="right", va="top")
ax.text(vx + 0.03, 0.13, r"$\Delta\mathbf{v}=-2v_r\,\hat{\mathbf{r}}$", color=AMBER, fontsize=13, va="center")
for s, lab in ((1, r"$\gamma$"), (-1, r"$-\gamma$")):
    th = np.linspace(0, s*gam, 30)
    ax.plot(0.42*np.cos(th), 0.42*np.sin(th), color=INK, lw=1.1)
    ax.text(0.47*np.cos(s*gam/2), 0.47*np.sin(s*gam/2), lab, fontsize=12, va="center")
ax.plot(0, 0, "o", color=INK, ms=7)
ax.set_xlim(-0.2, 1.42); ax.set_ylim(-0.58, 0.72)
ax.set_aspect("equal"); ax.axis("off")
fig.tight_layout()
fig.savefig("img/fig-man-apse-reflection.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-man-apse-reflection.png")

# ---- the two-burn route through apoapsis ----------------------------------
a, e = 1.0, 0.4
p = a*(1 - e**2)
ra = a*(1 + e)
eta = np.radians(75.0)                        # any angle: nothing below depends on it
fig, ax = plt.subplots(figsize=(7.4, 7.0))


def ellipse(th, th0=0.0):
    r = p/(1 + e*np.cos(th - th0))
    return r*np.cos(th), r*np.sin(th)


th = np.linspace(0, 2*np.pi, 800)
ax.plot(*ellipse(th), color=BLUE, lw=1.0, ls="--", alpha=0.55)
ax.plot(*ellipse(th, eta), color=GREEN, lw=1.0, ls="--", alpha=0.55)
ax.plot(ra*np.cos(th), ra*np.sin(th), color=MAROON, lw=1.0, ls="--", alpha=0.55)

t1 = np.linspace(0, np.pi, 300)                       # initial orbit, periapsis to apoapsis
t2 = np.linspace(np.pi, np.pi + eta, 200)             # the coast
t3 = np.linspace(np.pi + eta, 2*np.pi + eta, 300)     # final orbit, apoapsis to periapsis
ax.plot(*ellipse(t1), color=BLUE, lw=2.8)
ax.plot(ra*np.cos(t2), ra*np.sin(t2), color=MAROON, lw=2.8)
ax.plot(*ellipse(t3, eta), color=GREEN, lw=2.8)


def along(x, y, col, i):
    ax.annotate("", xy=(x[i + 3], y[i + 3]), xytext=(x[i], y[i]),
                arrowprops=dict(arrowstyle="-|>", color=col, lw=2.2, shrinkA=0, shrinkB=0, mutation_scale=18))


along(*ellipse(t1), BLUE, 130)
along(ra*np.cos(t2), ra*np.sin(t2), MAROON, 90)
along(*ellipse(t3, eta), GREEN, 170)

# apse lines and the angle between them
for ang, col in ((np.pi, BLUE), (np.pi + eta, GREEN)):
    ax.plot([0, ra*np.cos(ang)], [0, ra*np.sin(ang)], color=col, lw=1.0, ls=":")
arc = np.linspace(np.pi, np.pi + eta, 60)
ax.plot(0.5*np.cos(arc), 0.5*np.sin(arc), color=INK, lw=1.3)
mid = np.pi + eta/2
ax.text(0.6*np.cos(mid), 0.6*np.sin(mid), r"$\eta$", fontsize=15, ha="center", va="center")
ax.plot(0, 0, "o", color=INK, ms=8)
ax.text(0.04, -0.09, "focus", fontsize=10.5, va="top")

# the burns
L = 0.42
for ang, sgn, lab, xy, ha in ((np.pi, 1, "burn 1\nspeed $v_a$ to $v_c$", (-1.5, 0.12), "right"),
                              (np.pi + eta, -1, "burn 2\nspeed $v_c$ to $v_a$", (0.05, -1.68), "left")):
    x, y = ra*np.cos(ang), ra*np.sin(ang)
    ux, uy = -np.sin(ang), np.cos(ang)
    ax.annotate("", xy=(x + sgn*L*ux, y + sgn*L*uy), xytext=(x, y),
                arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=3, shrinkA=0, shrinkB=0, mutation_scale=20))
    ax.plot(x, y, "*", color=AMBER, ms=17, mec=INK, mew=0.6, zorder=6)
    ax.text(*xy, lab, color=AMBER, fontsize=11.5, ha=ha, va="center")

ax.text(0.15, 1.02, "initial orbit", color=BLUE, fontsize=12, ha="center")
ax.text(-1.26, -1.12, "circle of\nradius $r_a$", color=MAROON, fontsize=12, ha="right", va="center")
ax.text(1.12, -1.22, "final orbit", color=GREEN, fontsize=12, ha="left")
ax.set_aspect("equal"); ax.axis("off")
ax.set_xlim(-2.15, 1.75); ax.set_ylim(-1.9, 1.5)
fig.tight_layout()
fig.savefig("img/fig-man-apse-twoburn-geom.png", dpi=160, bbox_inches="tight")
plt.close(fig)
print("wrote img/fig-man-apse-twoburn-geom.png")