"""Figures for cr3bp.qmd.

img/fig-cr3bp-hill.png   : zero-velocity curves of the Earth-Moon problem for five Jacobi constants, and the
                           pseudo-potential along the x axis with the levels of the three collinear points.
img/fig-cr3bp-points.png : the five equilibrium points, and the stability of the triangular points against mu.
img/fig-cr3bp-orbits.png : planar Lyapunov orbits about Earth-Moon L1 and L2, and a Sun-Earth L1 halo orbit
                           (Aditya-L1's, A_z = 120 000 km) in two projections.

Run from the sfm-slides directory:   python3 codes/fig-cr3bp.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _cr3bp import units, collinear, triangular, jacobi, omega, DAY, orbit_history, lyapunov, halo, linear_eigenvalues

mu, L, T, V = units("Earth-Moon")
L1, L2, L3 = collinear(mu)
L4, L5 = triangular(mu)


def fz(x, y):
    r1 = np.sqrt((x + mu) ** 2 + y * y); r2 = np.sqrt((x - 1 + mu) ** 2 + y * y)
    return x * x + y * y + 2 * (1 - mu) / r1 + 2 * mu / r2


C1, C2, C3, C4 = (jacobi([p, 0, 0, 0, 0, 0], mu) for p in (L1, L2, L3, L4[0]))
C4 = jacobi([L4[0], L4[1], 0, 0, 0, 0], mu)

# ------------------------------------------------------------------ Hill regions
fig, axs = plt.subplots(2, 3, figsize=(15, 9.4))
xs = np.linspace(-1.7, 1.7, 1200)
X, Y = np.meshgrid(xs, np.linspace(-1.7, 1.7, 1200))
F = fz(X, Y)
cases = ((3.25, r"$C>C_{L_1}$: two separate regions"), (3.18, r"$C_{L_2}<C<C_{L_1}$: the $L_1$ neck is open"),
         (3.10, r"$C_{L_3}<C<C_{L_2}$: the $L_2$ neck is open"), (3.00, r"$C_{L_4}<C<C_{L_3}$: the $L_3$ neck is open"),
         (2.95, r"$C<C_{L_4}$: no forbidden region left"))
for ax, (C, ttl) in zip(axs.flat, cases):
    ax.contourf(X, Y, F, levels=[-1e9, C, 1e9], colors=["#d9d4c8", CREAM])
    ax.contour(X, Y, F, levels=[C], colors=[INK], linewidths=1.2)
    ax.scatter([-mu], [0], s=120, color=BLUE, zorder=5); ax.scatter([1 - mu], [0], s=45, color="#7a756b", zorder=5)
    for xp, lab in ((L1, "L1"), (L2, "L2"), (L3, "L3")):
        ax.scatter([xp], [0], s=18, color=MAROON, zorder=6); ax.text(xp, 0.07, lab, color=MAROON, fontsize=9, ha="center")
    ax.scatter(L4[0], L4[1], s=18, color=MAROON); ax.scatter(L5[0], L5[1], s=18, color=MAROON)
    ax.text(L4[0], L4[1] + 0.07, "L4", color=MAROON, fontsize=9, ha="center"); ax.text(L5[0], L5[1] - 0.14, "L5", color=MAROON, fontsize=9, ha="center")
    ax.set_aspect("equal"); ax.set_xlim(-1.7, 1.7); ax.set_ylim(-1.7, 1.7); ax.set_xticks([]); ax.set_yticks([])
    ax.set_title(f"$C={C:.2f}$\n{ttl}", fontsize=11)
ax = axs[1, 2]
xx = np.linspace(-1.6, 1.6, 3000)
fx = fz(xx, 0 * xx)
ax.plot(xx, np.clip(fx, 2.5, 8), color=INK, lw=1.6)
for Cc, lab, col, dy in ((C1, "$C_{L_1}$", GREEN, 0.05), (C2, "$C_{L_2}$", AMBER, -0.05), (C3, "$C_{L_3}$", PURPLE, 0.0)):
    ax.axhline(Cc, color=col, lw=1, ls="--"); ax.text(1.62, Cc + dy, lab, color=col, fontsize=10, va="center")
for xp in (L1, L2, L3):
    ax.scatter([xp], [fz(xp, 0.0)], color=MAROON, s=30, zorder=5)
ax.set_ylim(2.7, 4.6); ax.set_xlim(-1.6, 1.9)
ax.set_xlabel("$x$ (along the line of the primaries)"); ax.set_ylabel(r"$x^2+2(1-\mu)/r_1+2\mu/r_2$")
ax.set_title("A cut along the $x$ axis: the saddles are $L_1$, $L_2$, $L_3$", fontsize=11); despine(ax)
fig.suptitle(f"Earth--Moon zero-velocity curves: the grey region is forbidden to a spacecraft of Jacobi constant $C$ ($\\mu={mu:.5f}$)", fontsize=12.5, y=0.995)
fig.tight_layout(); fig.savefig("img/fig-cr3bp-hill.png", dpi=140, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ the five points and Routh's criterion
fig, (aL, aR) = plt.subplots(1, 2, figsize=(14, 5.8))
mu_d = 0.1
aL.set_title(r"The five points (drawn for $\mu=0.1$)", fontsize=12)
Ld = collinear(mu_d); L4d, L5d = triangular(mu_d)
xs_ = np.linspace(-1.7, 1.7, 600); Xd, Yd = np.meshgrid(xs_, xs_)
r1 = np.sqrt((Xd + mu_d) ** 2 + Yd ** 2); r2 = np.sqrt((Xd - 1 + mu_d) ** 2 + Yd ** 2)
Fd = Xd ** 2 + Yd ** 2 + 2 * (1 - mu_d) / r1 + 2 * mu_d / r2
aL.contour(Xd, Yd, Fd, levels=[2.7, 2.85, 3.0, 3.1, 3.25, 3.4, 3.7, 4.2, 5], colors="#b9b2a4", linewidths=0.8)
aL.scatter([-mu_d], [0], s=150, color=BLUE, zorder=5); aL.scatter([1 - mu_d], [0], s=70, color="#7a756b", zorder=5)
aL.text(-mu_d, -0.16, r"$m_1$", ha="center", color=BLUE); aL.text(1 - mu_d, -0.16, r"$m_2$", ha="center", color="#6f685e")
for xp, lab in zip(Ld, ("L1", "L2", "L3")):
    aL.scatter([xp], [0], s=35, color=MAROON, zorder=6); aL.text(xp, 0.09, lab, color=MAROON, ha="center", fontsize=11)
for P, lab, dy in ((L4d, "L4", 0.07), (L5d, "L5", -0.14)):
    aL.scatter(P[0], P[1], s=35, color=MAROON, zorder=6); aL.text(P[0], P[1] + dy, lab, color=MAROON, ha="center", fontsize=11)
for P in (L4d, L5d):
    aL.plot([-mu_d, 1 - mu_d, P[0], -mu_d], [0, 0, P[1], 0], color=GREEN, lw=1, ls=":")
aL.text(0.5 - mu_d, 0.45, "equilateral", color=GREEN, fontsize=9, ha="center")
aL.set_aspect("equal"); aL.set_xlim(-1.6, 1.6); aL.set_ylim(-1.35, 1.35); aL.axis("off")
mus = np.linspace(1e-5, 0.5, 4000)
disc = 1 - 27 * mus * (1 - mus)
mu_R = (1 - np.sqrt(23 / 27)) / 2
st = mus <= mu_R
w1 = np.sqrt((1 + np.sqrt(np.clip(disc, 0, None))) / 2); w2 = np.sqrt((1 - np.sqrt(np.clip(disc, 0, None))) / 2)
aR.plot(mus[st], w1[st], color=BLUE, lw=2.2, label=r"short-period frequency $\omega_1$")
aR.plot(mus[st], w2[st], color=GREEN, lw=2.2, label=r"long-period frequency $\omega_2$")
re = np.array([np.max(np.real(np.sqrt(np.array([(-1 + np.sqrt(complex(d))) / 2, (-1 - np.sqrt(complex(d))) / 2], dtype=complex)))) for d in disc])
aR.plot(mus[~st], re[~st], color=MAROON, lw=2.2, label=r"growth rate $\max\,\mathrm{Re}\,\lambda$ (unstable)")
aR.axvline(mu_R, color=INK, lw=0.9, ls="--"); aR.text(mu_R + 0.005, 0.9, f"Routh: $\\mu={mu_R:.5f}$", fontsize=10)
for m_, lab in ((9.5388e-4, "Sun--Jupiter"), (mu, "Earth--Moon")):
    aR.axvline(m_, color="#9a9388", lw=0.8); aR.text(m_ + 0.003, 0.05 if lab[0] == "S" else 0.2, lab, fontsize=9.5, rotation=90, va="bottom")
aR.set_xlabel(r"mass parameter $\mu$"); aR.set_ylabel(r"$|\lambda|$ at $L_4$, $L_5$ (in units of the primaries' rate)")
aR.set_xlim(0, 0.2); aR.set_ylim(0, 1.05); aR.legend(frameon=False, fontsize=9.5, loc="center right")
aR.set_title(r"Linear stability of $L_4$, $L_5$: stable only for $\mu<0.0385$", fontsize=12); despine(aR)
fig.tight_layout(); fig.savefig("img/fig-cr3bp-points.png", dpi=140, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ periodic orbits
fig = plt.figure(figsize=(15, 5.4))
ax = fig.add_subplot(1, 3, 1)
from _cr3bp import collinear_linear
for (Lp, sgn, dl, col, tag) in ((L1, +1, -1, GREEN, "L1"), (L2, -1, +1, AMBER, "L2")):
    c2, lam, nu, nz, k = collinear_linear(mu, Lp)
    out = []
    for amp in np.arange(0.004, 0.0441, 0.004):          # continuation in the amplitude
        if not out:
            guess = sgn * k * nu * amp
        elif len(out) == 1:
            guess = out[-1][1] * amp / out[-1][0]
        else:
            guess = out[-1][1] + (out[-1][1] - out[-2][1]) / (out[-1][0] - out[-2][0]) * (amp - out[-1][0])
        vy, th = lyapunov(mu, Lp + dl * amp, guess)
        out.append((amp, vy, th))
        if round(amp / 0.004) % 2 == 1:
            ts, Yo = orbit_history(mu, [Lp + dl * amp, 0, 0, 0, vy, 0], 2 * th, 800)
            ax.plot(Yo[0], Yo[1], color=col, lw=1.3)
    ax.scatter([Lp], [0], color=MAROON, s=20, zorder=5); ax.text(Lp, -0.025, tag, color=MAROON, ha="center", va="top", fontsize=10)
ax.scatter([1 - mu], [0], s=80, color="#7a756b", zorder=5); ax.text(1 - mu, 0.03, "Moon", ha="center", fontsize=9.5)
ax.scatter([-mu + 0.0], [0], s=0)
ax.set_xlim(0.75, 1.25); ax.set_ylim(-0.3, 0.3); ax.set_aspect("equal")
ax.set_xlabel("$x$"); ax.set_ylabel("$y$"); ax.set_title("Planar Lyapunov orbits about Earth--Moon $L_1$, $L_2$", fontsize=11.5); despine(ax)
# Sun-Earth halo
muS, LS, TS, VS = units("Sun-(Earth+Moon)")
L1S = collinear(muS)[0]
z0 = 120000.0 / LS
x0h, vyh, thh = halo(muS, L1S - 0.0008, z0, 0.006)
ts, Yh = orbit_history(muS, [x0h, 0, z0, 0, vyh, 0], 2 * thh, 3000)
gamma = (1 - muS - L1S) * LS
xh = (Yh[0] - L1S) * LS / 1e3; yh = Yh[1] * LS / 1e3; zh = Yh[2] * LS / 1e3
b = fig.add_subplot(1, 3, 2)
b.plot(yh, xh, color=PURPLE, lw=1.6); b.scatter([0], [0], color=MAROON, s=25, zorder=5); b.text(30, -15, "L1", color=MAROON, fontsize=10)
b.set_xlabel("$y$ [thousand km]"); b.set_ylabel(r"$x-x_{L_1}$ [thousand km]"); b.set_aspect("equal")
b.annotate("toward the Earth", xy=(0, 330), xytext=(0, 250), ha="center", fontsize=9, arrowprops=dict(arrowstyle="-|>", color=INK))
b.set_title(f"Sun--Earth $L_1$ halo (period {2 * thh * TS / DAY:.1f} d): the $xy$ view", fontsize=11.5); despine(b)
c = fig.add_subplot(1, 3, 3)
c.plot(xh, zh, color=PURPLE, lw=1.6); c.scatter([0], [0], color=MAROON, s=25, zorder=5)
c.set_xlabel(r"$x-x_{L_1}$ [thousand km]"); c.set_ylabel("$z$ [thousand km]"); c.set_aspect("equal")
c.set_title(r"the $xz$ view: $A_z=120\,000$ km", fontsize=11.5); despine(c)
fig.tight_layout(); fig.savefig("img/fig-cr3bp-orbits.png", dpi=140, bbox_inches="tight")
print("Jacobi constants:", C1, C2, C3, C4)
print(f"halo: x0 = {x0h:.8f}, vy0 = {vyh:.8f}, period {2 * thh * TS / DAY:.3f} d")
print("wrote img/fig-cr3bp-hill.png, img/fig-cr3bp-points.png, img/fig-cr3bp-orbits.png")