"""Figures for objective-questions.qmd (fig-oq-*.png).

  fig-oq-nbody-configs   three equal masses on a triangle, four on a square
  fig-oq-conics          four trajectories through one periapsis, labelled 1-4
  fig-oq-ellipse-points  an ellipse with P, A, B, L, C, F, F* marked
  fig-oq-elements-3d     the orbit plane, the equator and three numbered angles
  fig-oq-coaxial         two aligned coaxial ellipses and the two apsidal routes X, Y

Run from the sfm-slides directory:   python3 codes/fig-oq-objective.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE


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


# ------------------------------------------------------------------ n-body configurations
fig, (a1, a2) = plt.subplots(1, 2, figsize=(10, 4.4))
tri = np.array([[0, 0], [1, 0], [0.5, np.sqrt(3) / 2]])
a1.fill(tri[:, 0], tri[:, 1], color=BLUE, alpha=0.06)
a1.plot(*np.vstack([tri, tri[:1]]).T, color=INK, lw=1.2)
a1.scatter(tri[:, 0], tri[:, 1], s=260, color=MAROON, zorder=3)
for p in tri:
    a1.text(p[0], p[1] - 0.12 if p[1] < 0.1 else p[1] + 0.09, "$m$", ha="center", fontsize=14)
cm = tri.mean(axis=0)
a1.scatter(*cm, marker="+", s=160, color=INK)
a1.text(cm[0] + 0.04, cm[1] - 0.07, "C", fontsize=12)
a1.text(0.5, -0.09, "$d$", ha="center", fontsize=13)
a1.set_title("(a) equilateral triangle, side $d$", fontsize=12)
sq = np.array([[0, 0], [1, 0], [1, 1], [0, 1]])
a2.fill(sq[:, 0], sq[:, 1], color=BLUE, alpha=0.06)
a2.plot(*np.vstack([sq, sq[:1]]).T, color=INK, lw=1.2)
a2.plot([0, 1], [0, 1], color=INK, lw=0.8, ls=":")
a2.plot([0, 1], [1, 0], color=INK, lw=0.8, ls=":")
a2.scatter(sq[:, 0], sq[:, 1], s=260, color=MAROON, zorder=3)
for p in sq:
    a2.text(p[0] + (-0.1 if p[0] < 0.5 else 0.1), p[1] + (-0.1 if p[1] < 0.5 else 0.1), "$m$",
            ha="center", va="center", fontsize=14)
a2.text(0.5, -0.11, "$d$", ha="center", fontsize=13)
a2.text(0.6, 0.47, r"$\sqrt{2}\,d$", fontsize=12)
a2.set_title("(b) square, side $d$", fontsize=12)
for ax in (a1, a2):
    ax.set_aspect("equal"); ax.axis("off")
    ax.set_xlim(-0.3, 1.3); ax.set_ylim(-0.3, 1.25)
save(fig, "fig-oq-nbody-configs")

# ------------------------------------------------------------------ four conics, one periapsis
rp = 1.0
fig, ax = plt.subplots(figsize=(7.4, 6.4))
curves = [(0.0, 3, BLUE), (0.6, 1, GREEN), (1.0, 4, AMBER), (1.6, 2, MAROON)]   # (e, label, colour)
for e, lab, col in curves:
    p = rp * (1 + e)
    fmax = np.pi if e < 1 else (np.arccos(-1 / e) - 0.02 if e > 1 else np.radians(150))
    f = np.linspace(-fmax, fmax, 1200)
    r = p / (1 + e * np.cos(f))
    ok = r < 6.5
    ax.plot((r * np.cos(f))[ok], (r * np.sin(f))[ok], color=col, lw=2.2)
    fl = {0.0: np.radians(120), 0.6: np.radians(150), 1.0: np.radians(118), 1.6: np.radians(90)}[e]
    rl = p / (1 + e * np.cos(fl))
    ax.text(rl * np.cos(fl) + 0.12, rl * np.sin(fl) + 0.12, str(lab), color=col, fontsize=16, fontweight="bold")
ax.scatter([0], [0], s=90, color=INK, zorder=3)
ax.text(0.08, -0.32, "F", fontsize=12)
ax.scatter([rp], [0], s=30, color=INK, zorder=3)
ax.text(rp + 0.08, -0.3, "periapsis", fontsize=10)
ax.set_xlim(-4.6, 1.8); ax.set_ylim(-4.2, 4.2)
ax.set_aspect("equal"); ax.axis("off")
save(fig, "fig-oq-conics")

# ------------------------------------------------------------------ ellipse with marked points
a, e = 1.0, 0.5
b, c, p = a * np.sqrt(1 - e * e), a * e, a * (1 - e * e)
fig, ax = plt.subplots(figsize=(8.4, 5.2))
t = np.linspace(0, 2 * np.pi, 600)
ax.plot(-c + a * np.cos(t), b * np.sin(t), color=BLUE, lw=2.2)
ax.plot([-c - a, -c + a], [0, 0], color=INK, lw=0.8, ls=":")
ax.plot([-c, -c], [-b, b], color=INK, lw=0.8, ls=":")
ax.plot([0, 0], [0, p], color=INK, lw=0.8, ls=":")
pts = {"P": (a - c, 0), "A": (-a - c, 0), "B": (-c, b), "L": (0, p)}
for k, (x, y) in pts.items():
    ax.scatter([x], [y], s=55, color=MAROON, zorder=3)
off = {"P": (0.05, 0.05), "A": (-0.13, 0.05), "B": (-0.03, 0.07), "L": (0.04, 0.06)}
for k, (x, y) in pts.items():
    ax.text(x + off[k][0], y + off[k][1], k, color=MAROON, fontsize=15, fontweight="bold")
ax.scatter([0], [0], s=90, color=INK, zorder=3); ax.text(0.03, -0.11, "F", fontsize=13)
ax.scatter([-2 * c], [0], s=30, color=INK, zorder=3); ax.text(-2 * c - 0.03, -0.11, r"F$^*$", fontsize=13)
ax.scatter([-c], [0], s=20, color=INK, marker="x"); ax.text(-c + 0.02, -0.11, "C", fontsize=12)
ax.annotate("", xy=(np.cos(np.radians(40)) * 0.22, np.sin(np.radians(40)) * 0.22), xytext=(0.22, 0),
            arrowprops=dict(arrowstyle="-|>", color=GREEN, connectionstyle="arc3,rad=0.4"))
ax.text(0.25, 0.08, "motion", color=GREEN, fontsize=10)
ax.set_title(r"Ellipse, $e=0.5$: F is the attracting focus", fontsize=12, pad=22)
ax.set_aspect("equal"); ax.axis("off")
save(fig, "fig-oq-ellipse-points")

# ------------------------------------------------------------------ three angles of the orbit
Om, inc, w = np.radians(40), np.radians(42), np.radians(115)


def R3(x):
    return np.array([[np.cos(x), -np.sin(x), 0], [np.sin(x), np.cos(x), 0], [0, 0, 1]])


def R1(x):
    return np.array([[1, 0, 0], [0, np.cos(x), -np.sin(x)], [0, np.sin(x), np.cos(x)]])


Q = R3(Om) @ R1(inc)                         # node frame -> ECI (orbit plane drawn as a circle)
fig = plt.figure(figsize=(7.6, 7.0))
ax = fig.add_subplot(projection="3d")
ax.set_facecolor(CREAM)
ax.view_init(elev=24, azim=-62)
tt = np.linspace(0, 2 * np.pi, 300)
ax.plot(np.cos(tt), np.sin(tt), 0 * tt, color=INK, lw=0.9, ls=(0, (4, 3)))
xx, yy = np.meshgrid(np.linspace(-1.2, 1.2, 2), np.linspace(-1.2, 1.2, 2))
ax.plot_surface(xx, yy, 0 * xx, color=BLUE, alpha=0.06, edgecolor="none")
orb = Q @ np.vstack([np.cos(tt), np.sin(tt), 0 * tt])
ax.plot(*orb, color=MAROON, lw=2.2)
for v, lab in ((np.array([1.35, 0, 0]), r"$\mathbf{\hat{I}}$"), (np.array([0, 1.35, 0]), r"$\mathbf{\hat{J}}$"),
               (np.array([0, 0, 1.25]), r"$\mathbf{\hat{K}}$")):
    ax.quiver(0, 0, 0, *v, color=INK, lw=1.2, arrow_length_ratio=0.08)
    ax.text(*(1.07 * v), lab, fontsize=13)
hv = Q @ np.array([0, 0, 0.95])
ax.quiver(0, 0, 0, *hv, color=MAROON, lw=1.6, arrow_length_ratio=0.08)
ax.text(*(1.08 * hv), r"$\mathbf{h}$", color=MAROON, fontsize=13)
nd = np.array([np.cos(Om), np.sin(Om), 0])
ax.plot(*np.vstack([-nd, nd]).T, color=INK, lw=0.8, ls=":")
ax.scatter(*nd, s=40, color=GREEN); ax.text(*(1.12 * nd), "N", color=GREEN, fontsize=13)
per = Q @ np.array([np.cos(w), np.sin(w), 0])
ax.plot(*np.vstack([[0, 0, 0], per]).T, color=PURPLE, lw=1.4)
ax.scatter(*per, s=40, color=PURPLE); ax.text(*(1.08 * per), "periapsis", color=PURPLE, fontsize=11)
s = np.linspace(0, Om, 40)
ax.plot(0.55 * np.cos(s), 0.55 * np.sin(s), 0 * s, color=GREEN, lw=2)
ax.text(0.64 * np.cos(Om / 2), 0.52 * np.sin(Om / 2), 0, r"$\theta_1$", color=GREEN, fontsize=14)
s = np.linspace(0, w, 40)
arc = Q @ np.vstack([0.5 * np.cos(s), 0.5 * np.sin(s), 0 * s])
ax.plot(*arc, color=PURPLE, lw=2)
mid = Q @ np.array([0.6 * np.cos(w / 2), 0.6 * np.sin(w / 2), 0])
ax.text(*mid, r"$\theta_3$", color=PURPLE, fontsize=14)
# angle between the planes, drawn at the node in the plane perpendicular to the node line
side = np.array([-np.sin(Om), np.cos(Om), 0])
s = np.linspace(0, inc, 30)
arc = np.array([nd * 1.0 + 0.5 * (np.cos(x) * side + np.sin(x) * np.array([0, 0, 1])) for x in s]).T
ax.plot(*arc, color=AMBER, lw=2)
pt = nd + 0.62 * (np.cos(inc / 2) * side + np.sin(inc / 2) * np.array([0, 0, 1]))
ax.text(*pt, r"$\theta_2$", color=AMBER, fontsize=14)
ax.text(0.15, -1.2, 0, "equatorial plane", fontsize=10, color=BLUE)
ax.set_xlim(-1.1, 1.1); ax.set_ylim(-1.1, 1.1); ax.set_zlim(-0.8, 1.1)
ax.set_box_aspect((1, 1, 0.85))
ax.set_axis_off()
save(fig, "fig-oq-elements-3d")

# ------------------------------------------------------------------ aligned coaxial ellipses
rp1, ra1, rp2, ra2 = 8400.0, 15600.0, 21000.0, 63000.0


def ell(rp, ra, sign=1, f0=-np.pi, f1=np.pi):
    a, e = (rp + ra) / 2, (ra - rp) / (ra + rp)
    f = np.linspace(f0, f1, 500)
    r = a * (1 - e * e) / (1 + e * np.cos(f))
    return sign * r * np.cos(f), r * np.sin(f)


fig, ax = plt.subplots(figsize=(9.4, 5.6))
ax.plot(*ell(rp1, ra1), color=BLUE, lw=2.0)
ax.plot(*ell(rp2, ra2), color=INK, lw=2.0)
ax.plot(*ell(rp1, ra2, f0=0, f1=np.pi), color=MAROON, lw=2.0, ls=(0, (6, 3)))           # Y: rp1 -> ra2, upper
x, y = ell(ra1, rp2, f0=0, f1=np.pi)                                                    # X: ra1 -> rp2, lower
ax.plot(-x, -y, color=PURPLE, lw=2.0, ls=(0, (6, 3)))
ax.scatter([0], [0], s=60, color=INK, zorder=3)
ax.text(rp1 * 0.15, -5000, "F", fontsize=12)
ax.text(rp1 + 600, -3000, "1", color=BLUE, fontsize=14, fontweight="bold")
ax.text(rp2 + 900, 4000, "2", color=INK, fontsize=14, fontweight="bold")
ax.text(-30000, 26000, "Y", color=MAROON, fontsize=15, fontweight="bold")
ax.text(-3000, -21000, "X", color=PURPLE, fontsize=15, fontweight="bold")
ax.text(rp2 - 4000, 33000, "both periapses on this side $\\rightarrow$", fontsize=10, ha="right")
ax.set_aspect("equal"); ax.axis("off")
save(fig, "fig-oq-coaxial")
print("wrote five figures")