"""Number-free sketches for the Hohmann, bi-elliptic and phasing notes. Each one shows the geometry of a manoeuvre and
carries no result, so it can stand before the derivation or the example.
Run from sfm-slides/:   python3 codes/fig-man-transfer-sketches.py"""
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

TH = np.linspace(0, 2*np.pi, 700)


def conic(p, e, w, f):
    r = p/(1 + e*np.cos(f))
    return r*np.cos(f + w), r*np.sin(f + w)


def ellipse(ax, rp, ra, w=0.0, **kw):
    a, e = (rp + ra)/2, (ra - rp)/(ra + rp)
    ax.plot(*conic(a*(1 - e*e), e, w, TH), **kw)


def half(ax, ra_start, rb_end, w_start, **kw):
    """Half an ellipse from radius ra_start at angle w_start to radius rb_end on the opposite side, counter-clockwise."""
    rp, ra = min(ra_start, rb_end), max(ra_start, rb_end)
    a, e = (rp + ra)/2, (ra - rp)/(ra + rp)
    if ra_start <= rb_end:
        f, w = np.linspace(0, np.pi, 300), w_start
    else:
        f, w = np.linspace(np.pi, 2*np.pi, 300), w_start - np.pi
    ax.plot(*conic(a*(1 - e*e), e, w, f), **kw)


def tick(ax, x, y, s, up=True, col=AMBER):
    """A short tangential burn arrow at an apsis on the x axis, labelled s."""
    d = 0.11*ax._span*(1 if up else -1)
    ax.annotate("", xy=(x, y + d), xytext=(x, y), arrowprops=dict(arrowstyle="-|>", color=col, lw=2.2, shrinkA=0, shrinkB=0))
    ax.plot([x], [y], "o", color=col, ms=5, zorder=6)


def new(span_x, span_y, size=(7.6, 5.4)):
    fig, ax = plt.subplots(figsize=size)
    ax.set_aspect("equal"); ax.axis("off")
    ax.set_xlim(*span_x); ax.set_ylim(*span_y)
    ax._span = span_y[1] - span_y[0]
    return fig, ax


def save(fig, name):
    fig.tight_layout()
    fig.savefig(f"img/{name}.png", dpi=160, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote img/{name}.png")


def focus(ax, s=9):
    ax.plot(0, 0, "o", color=INK, ms=s, zorder=7)


# ---- coaxial ellipses, aligned -------------------------------------------------------------------------------
rp1, ra1, rp2, ra2 = 8400.0, 15600.0, 21000.0, 63000.0
fig, ax = new((-70000, 30000), (-46000, 50000))
ellipse(ax, rp1, ra1, 0.0, color=BLUE, lw=1.8); ellipse(ax, rp2, ra2, 0.0, color=GREEN, lw=1.8)
ax.plot([-ra2, rp2], [0, 0], color=INK, lw=0.7, alpha=0.4)
half(ax, rp1, ra2, 0.0, color=MAROON, lw=2.6)                   # A1: periapsis of 1 (right) to apoapsis of 2 (left), upper half
half(ax, ra1, rp2, np.pi, color=PURPLE, lw=2.6)                 # A2: apoapsis of 1 (left) to periapsis of 2 (right), lower half
focus(ax)
ax.text(-27000, 18000, "A1", color=MAROON, fontsize=15); ax.text(1000, -28000, "A2", color=PURPLE, fontsize=15)
ax.text(-3600, 3500, "orbit 1", color=BLUE, fontsize=13, ha="center"); ax.text(-22000, -43000, "orbit 2", color=GREEN, fontsize=13, ha="center")
for x, y in ((rp1, 0), (-ra2, 0), (-ra1, 0), (rp2, 0)):
    ax.plot([x], [y], "o", color=AMBER, ms=8, zorder=6)
save(fig, "fig-man-sk-coax-aligned")

# ---- coaxial ellipses, anti-aligned --------------------------------------------------------------------------
fig, ax = new((-30000, 70000), (-46000, 50000))
ellipse(ax, rp1, ra1, 0.0, color=BLUE, lw=1.8); ellipse(ax, rp2, ra2, np.pi, color=GREEN, lw=1.8)
ax.plot([-rp2, ra2], [0, 0], color=INK, lw=0.7, alpha=0.4)
half(ax, rp1, rp2, 0.0, color=MAROON, lw=2.6)                   # T2: the periapses, upper half
half(ax, ra1, ra2, np.pi, color=PURPLE, lw=2.6)                 # T1: the apoapses, lower half
focus(ax)
ax.text(-3000, 15000, "T2", color=MAROON, fontsize=15, ha="center"); ax.text(22000, -24000, "T1", color=PURPLE, fontsize=15)
ax.text(-3600, -6000, "orbit 1", color=BLUE, fontsize=13, ha="center"); ax.text(32000, 42000, "orbit 2", color=GREEN, fontsize=13, ha="center")
for x, y in ((rp1, 0), (-rp2, 0), (-ra1, 0), (ra2, 0)):
    ax.plot([x], [y], "o", color=AMBER, ms=8, zorder=6)
save(fig, "fig-man-sk-coax-anti")

# ---- crossing orbits -----------------------------------------------------------------------------------------
Rp1, Ra1, Rp2, Ra2 = 8000.0, 30000.0, 12000.0, 24000.0
fig, ax = new((-36000, 18000), (-24000, 24000), size=(6.8, 5.8))
ellipse(ax, Rp1, Ra1, 0.0, color=BLUE, lw=2.0); ellipse(ax, Rp2, Ra2, 0.0, color=GREEN, lw=2.0)
ax.plot([-Ra1, Rp2], [0, 0], color=INK, lw=0.7, alpha=0.4)
focus(ax)
p1, e1, p2, e2 = 2*Rp1*Ra1/(Rp1 + Ra1), (Ra1 - Rp1)/(Ra1 + Rp1), 2*Rp2*Ra2/(Rp2 + Ra2), (Ra2 - Rp2)/(Ra2 + Rp2)
fs = np.arccos((p2 - p1)/(p1*e2 - p2*e1)); rs = p1/(1 + e1*np.cos(fs))
for s in (1, -1):
    ax.plot([rs*np.cos(fs)], [s*rs*np.sin(fs)], "o", color=AMBER, ms=10, zorder=7)
ax.text(-27000, 14500, "orbit 1", color=BLUE, fontsize=13, ha="center"); ax.text(9000, 14000, "orbit 2", color=GREEN, fontsize=13, ha="center")
ax.text(rs*np.cos(fs) + 1800, rs*np.sin(fs) + 1800, "a crossing", color=AMBER, fontsize=13)
save(fig, "fig-man-sk-coax-cross")

# ---- stopover ------------------------------------------------------------------------------------------------
r1, rb, r2 = 1.0, 3.5, 5.0
fig, ax = new((-6.4, 5.4), (-5.9, 5.9), size=(6.2, 6.0))
ax.plot(r1*np.cos(TH), r1*np.sin(TH), color=BLUE, lw=1.8); ax.plot(r2*np.cos(TH), r2*np.sin(TH), color=GREEN, lw=1.8)
ax.plot(rb*np.cos(TH), rb*np.sin(TH), color=INK, lw=0.8, ls=":", alpha=0.7)
half(ax, r1, rb, 0.0, color=MAROON, lw=2.6)                     # ellipse 1, upper half, from r1 out to r_b on the left
half(ax, rb, r2, np.pi, color=PURPLE, lw=2.6)                   # ellipse 2, lower half, from r_b out to r_2 on the right
ax.plot([-rb, r2], [0, 0], color=INK, lw=0.7, alpha=0.4); focus(ax, 7)
for x, s, lx, ha in ((r1, "burn 1", 1.15, "left"), (-rb, "burn 2", -3.65, "right"), (r2, "burn 3", 4.85, "right")):
    ax.plot([x], [0], "o", color=AMBER, ms=9, zorder=7)
    ax.text(lx, -0.48 if x > 0 else 0.25, s, color=AMBER, fontsize=12, ha=ha)
ax.text(-1.2, 2.45, "ellipse 1", color=MAROON, fontsize=12, ha="center"); ax.text(0.4, -3.55, "ellipse 2", color=PURPLE, fontsize=12, ha="center")
ax.text(0.45, 0.14, "$r_1$", color=BLUE, fontsize=13); ax.text(-2.0, 0.14, "$r_b$", color=INK, fontsize=13); ax.text(2.6, 0.14, "$r_2$", color=GREEN, fontsize=13)
save(fig, "fig-man-sk-stopover")

# ---- phasing -------------------------------------------------------------------------------------------------
r = 1.0; dth = np.radians(55.0); aph = r*(1 - dth/(2*np.pi))**(2/3); rp = 2*aph - r; e = (r - rp)/(r + rp)
fig, ax = new((-1.45, 1.95), (-1.35, 1.35), size=(7.0, 5.4))
ax.plot(r*np.cos(TH), r*np.sin(TH), color=BLUE, lw=2.0)
f = TH; R = aph*(1 - e*e)/(1 + e*np.cos(f))
ax.plot(-R*np.cos(f), -R*np.sin(f), color=MAROON, lw=2.4)
ax.plot(0, 0, "o", color=INK, ms=8, zorder=7)
ax.plot([0, r], [0, 0], color=INK, lw=0.8, ls=":"); ax.plot([0, r*np.cos(dth)], [0, r*np.sin(dth)], color=INK, lw=0.8, ls=":")
t = np.linspace(0, dth, 50); ax.plot(0.42*np.cos(t), 0.42*np.sin(t), color=INK, lw=1.2)
ax.text(0.47*np.cos(dth/2), 0.47*np.sin(dth/2), r"$\Delta\theta$", fontsize=14, va="bottom")
ax.plot([r], [0], "o", color=MAROON, ms=11, zorder=7); ax.text(r + 0.06, -0.08, "chaser", color=MAROON, fontsize=13, va="top")
ax.plot([r*np.cos(dth)], [r*np.sin(dth)], "o", color=GREEN, ms=11, zorder=7)
ax.text(r*np.cos(dth) + 0.06, r*np.sin(dth) + 0.03, "target", color=GREEN, fontsize=13)
ax.text(-0.15, 1.1, "circular orbit, period $T$", color=BLUE, fontsize=12, ha="center")
ax.text(-0.1, -0.42, "phasing ellipse,\nperiod $T_{ph}$", color=MAROON, fontsize=12, ha="center")
ax.annotate("", xy=(-0.72, 0.76), xytext=(-0.62, 0.86), arrowprops=dict(arrowstyle="-|>", color=INK, lw=1.4))
ax.text(-0.82, 0.86, "motion", fontsize=12, ha="right")
save(fig, "fig-man-sk-phasing")