"""Lambert transfers between the same two points: short way, long way, and one extra revolution (both ways).

All four panels use r2/r1 = 11/7 and the same time of flight, dt = 2 T1 (T1 = period of the circular orbit of
radius r1).  The craft moves at its true Keplerian speed (fast near periapsis, slow near apoapsis), only slowed down.
Zero revolutions: one solution.  One revolution: two solutions, drawn in two colours.

Run from the sfm-slides directory:   python3 codes/anim-lam-types.py
Writes img/anim-lam-types.gif (the animation) and img/fig-lam-types.png (its last frame, with direction arrows).
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE
from _gif import save_gif
from _lambertp import mu, p_window, tof_of_p, conic_through, kepler_E

RA, RB = 7000.0, 11000.0
T1 = 2*np.pi*np.sqrt(RA**3/mu)
DT = 2.0*T1
panels = [("(a) short way, $\\Delta\\theta=100^\\circ$, no extra revolution", 100.0, 0),
          ("(b) long way, $\\Delta\\theta=260^\\circ$, no extra revolution", 260.0, 0),
          ("(c) short way, one extra revolution ($N=1$): two solutions", 100.0, 1),
          ("(d) long way, one extra revolution ($N=1$): two solutions", 260.0, 1)]
COLS = {0: [MAROON], 1: [BLUE, GREEN]}


def solutions(deg, N):
    dth = np.radians(deg)
    lo, hi = p_window(RA, RB, dth)
    ps = np.linspace(lo*(1 + 1e-9), hi*(1 - 1e-9), 30001)
    ts = np.array([tof_of_p(q, RA, RB, dth, N) for q in ps]) - DT
    out = []
    for i in range(len(ps) - 1):
        if np.isfinite(ts[i]) and np.isfinite(ts[i + 1]) and ts[i]*ts[i + 1] < 0:
            out.append(brentq(lambda x: tof_of_p(x, RA, RB, dth, N) - DT, ps[i], ps[i + 1], xtol=1e-9))
    return out


def path(p, deg, s):
    """Position (units of r1) at the fraction s of the flight time; r1 on the +x axis."""
    dth = np.radians(deg)
    e, f1 = conic_through(p, RA, RB, dth)
    a = p/(1 - e*e); n = np.sqrt(mu/a**3)
    E1 = 2*np.arctan2(np.sqrt(1 - e)*np.sin(f1/2), np.sqrt(1 + e)*np.cos(f1/2))
    M = E1 - e*np.sin(E1) + n*DT*np.asarray(s)
    E = kepler_E(np.mod(M, 2*np.pi), e)
    f = 2*np.arctan2(np.sqrt(1 + e)*np.sin(E/2), np.sqrt(1 - e)*np.cos(E/2))
    r = p/(1 + e*np.cos(f))
    ang = f - f1
    return np.stack([r*np.cos(ang), r*np.sin(ang)], axis=-1)/RA, e, f1


def outline(p, e, f1, n=900):
    f = np.linspace(-np.pi, np.pi, n)
    r = p/(1 + e*np.cos(f))
    return np.stack([r*np.cos(f - f1), r*np.sin(f - f1)], axis=-1)/RA


sols = [solutions(deg, N) for _, deg, N in panels]
for (title, deg, N), ss in zip(panels, sols):
    for q in ss:
        e_, f1_ = conic_through(q, RA, RB, np.radians(deg))
        end = path(q, deg, 1.0)[0]
        tgt = RB/RA*np.array([np.cos(np.radians(deg)), np.sin(np.radians(deg))])
        assert np.allclose(end, tgt, atol=1e-6), (deg, N, q, end, tgt)
        print(f"{title[:3]} dtheta={deg:5.1f} N={N}: p={q:8.2f} km, e={e_:.4f}")

ss_ = np.linspace(0, 1, 400)
boxes = []
for (_, deg, _), s in zip(panels, sols):
    pts = np.concatenate([path(q, deg, ss_)[0] for q in s] + [np.zeros((1, 2))])
    boxes.append((pts.min(axis=0), pts.max(axis=0)))
HALF = max(np.max(hi - lo) for lo, hi in boxes)/2 + 0.45             # one scale for all four panels

fig, axs = plt.subplots(2, 2, figsize=(10.2, 10.0))
state = []
for ax, (title, deg, N), ss, (blo, bhi) in zip(axs.flat, panels, sols, boxes):
    dth = np.radians(deg)
    P1, P2 = np.array([1.0, 0.0]), (RB/RA)*np.array([np.cos(dth), np.sin(dth)])
    ax.plot(0, 0, "o", color=INK, ms=9, zorder=6)
    ax.plot([0, P1[0]], [0, 0], color=INK, lw=1.0, alpha=0.6)
    ax.plot([0, P2[0]], [0, P2[1]], color=INK, lw=1.0, alpha=0.6)
    ax.plot(*P1, "o", color=INK, ms=9, zorder=7); ax.plot(*P2, "s", color=INK, ms=9, zorder=7)
    ax.text(P1[0] + 0.12, P1[1] - 0.30, r"$\mathbf{r}_1$", fontsize=14)
    u2 = P2/np.linalg.norm(P2)
    ax.text(*(P2 + 0.34*u2), r"$\mathbf{r}_2$", fontsize=14, ha="center", va="center")
    arc = np.linspace(0, dth, 60)
    ax.plot(0.45*np.cos(arc), 0.45*np.sin(arc), color=INK, lw=1)
    ax.text(0.62*np.cos(dth/2), 0.62*np.sin(dth/2), r"$\Delta\theta$", fontsize=12)
    for q, col in zip(ss, COLS[N]):
        e_, f1_ = conic_through(q, RA, RB, dth)
        o = outline(q, e_, f1_)
        ax.plot(o[:, 0], o[:, 1], color=col, lw=1.0, ls=(0, (4, 4)), alpha=0.45, zorder=2)
        trail, = ax.plot([], [], color=col, lw=3.0, zorder=5,
                         label=f"$p={q:.0f}$ km,  $e={e_:.2f}$")
        dot, = ax.plot([], [], "o", color=col, ms=11, zorder=8, mec=INK)
        arrows = []
        for sf in (0.18, 0.42, 0.66, 0.90):
            a0, a1 = path(q, deg, [sf - 0.012, sf + 0.012])[0]
            ar = ax.annotate("", xy=a1, xytext=a0, zorder=6, annotation_clip=False,
                             arrowprops=dict(arrowstyle="-|>", color=col, lw=2.4, mutation_scale=17))
            ar.set_visible(False)
            arrows.append((sf, ar))
        state.append((q, deg, trail, dot, arrows))
    cx, cy = 0.5*(blo + bhi)
    ax.set_xlim(cx - HALF, cx + HALF); ax.set_ylim(cy - HALF, cy + HALF)
    ax.set_aspect("equal"); ax.axis("off")
    ax.set_title(title, fontsize=11.5)
    ax.legend(frameon=False, fontsize=10.5, loc="upper center", bbox_to_anchor=(0.5, 0.02), ncol=1)
clock = fig.text(0.5, 0.012, "", ha="center", fontsize=13, family="monospace", color=INK)
fig.suptitle(r"Same $\mathbf{r}_1$, $\mathbf{r}_2$ and time of flight $\Delta t=2\,T_1$: different transfers", fontsize=13.5, y=0.995)
fig.tight_layout(rect=(0, 0.035, 1, 0.975))

N_MOVE, N_HOLD = 100, 30


def frame(i):
    s = min(i, N_MOVE)/N_MOVE
    for q, deg, trail, dot, arrows in state:
        pts = path(q, deg, np.linspace(0, s, max(2, int(400*s) + 2)))[0]
        trail.set_data(pts[:, 0], pts[:, 1]); dot.set_data([pts[-1, 0]], [pts[-1, 1]])
        for sf, ar in arrows:
            ar.set_visible(s >= sf + 0.012)
    clock.set_text(f"t = {s*DT/T1:4.2f} T1   (end: {DT/T1:.2f} T1)")


frame(N_MOVE)
fig.savefig("img/fig-lam-types.png", dpi=150, facecolor=CREAM)
print("wrote img/fig-lam-types.png")
save_gif(fig, frame, N_MOVE + N_HOLD + 1, 20, "img/anim-lam-types.gif", dpi=100, colors=64)