"""Where the vernal equinox direction comes from.

The geocentric equatorial frame needs a direction in the reference plane to
measure Omega from, and that direction is not arbitrary: it is the line along
which Earth's *equatorial* plane cuts the *ecliptic* plane.

Two planes meet in a line, and this one meets the orbit twice -- the March and
September equinoxes.  Xhat is the direction from Earth to the Sun at the March
crossing.  Because Earth's spin axis holds a fixed direction in inertial space
while Earth runs round its orbit, that line is fixed too (to first order: it
precesses by 50.3"/yr, which is why one must say *which* equinox -- J2000).

    obliquity   eps = 23.44 deg
    spin axis   nhat = (0, sin eps, cos eps)     -- no X component, so the
                                                    equatorial plane cuts the
                                                    ecliptic exactly along X

Nothing is to scale: the Sun and Earth are drawn ~100x and ~10^4x too large,
and the orbit is drawn as a circle (its true eccentricity is 0.0167).

Run from the deck directory:   python codes/fig-equinox-frame.py
"""

import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

os.makedirs("img", exist_ok=True)

CREAM, INK, DIM = "#fcfbf7", "#2f2a24", "#8b857c"
MAROON, BLUE, GREEN, AMBER, PURPLE = "#6b1f1f", "#2f5ea8", "#2d6a2d", "#b45309", "#7e22ce"
AXIS = "#6b655c"                      # same neutral grey as fig-orbital-elements-3d

plt.rcParams.update({
    "figure.facecolor": CREAM, "savefig.facecolor": CREAM, "axes.facecolor": CREAM,
    "text.color": INK, "font.size": 11,
})

EPS = np.deg2rad(23.44)                       # obliquity of the ecliptic
R = 1.0                                       # orbit radius (1 AU, drawn)
R_E = 0.115                                   # Earth, wildly out of scale
R_EQ = 0.46                                   # drawn size of the equatorial plane
nhat = np.array([0.0, np.sin(EPS), np.cos(EPS)])       # spin axis, fixed
zhat = np.array([0.0, 0.0, 1.0])                       # ecliptic normal

fig = plt.figure(figsize=(11.6, 6.6))
ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], projection="3d")
ax.view_init(elev=27, azim=-62)


def sphere(centre, radius, n=44):
    u = np.linspace(0.0, 2 * np.pi, 2 * n)
    v = np.linspace(0.0, np.pi, n)
    x = radius * np.outer(np.cos(u), np.sin(v)) + centre[0]
    y = radius * np.outer(np.sin(u), np.sin(v)) + centre[1]
    z = radius * np.outer(np.ones_like(u), np.cos(v)) + centre[2]
    return x, y, z


def disc(centre, normal, radius, n=90):
    """Polygon of a disc through `centre` with the given normal."""
    normal = np.asarray(normal, float)
    normal = normal / np.linalg.norm(normal)
    u = np.array([1.0, 0.0, 0.0])
    if abs(u @ normal) > 0.9:
        u = np.array([0.0, 1.0, 0.0])
    u = u - (u @ normal) * normal
    u /= np.linalg.norm(u)
    w = np.cross(normal, u)
    t = np.linspace(0.0, 2 * np.pi, n)
    return np.array([centre + radius * (np.cos(a) * u + np.sin(a) * w)
                     for a in t])


def circle3(centre, normal, radius, n=200):
    return disc(centre, normal, radius, n)


# ── the ecliptic plane, and the orbit ruled on it ───────────────────────────
ecl = disc(np.zeros(3), zhat, 1.52)
ax.add_collection3d(Poly3DCollection([ecl], facecolor=DIM, alpha=0.10,
                                     edgecolor=DIM, linewidths=0.9))
for rad in (0.5, 1.0, 1.5):                    # faint rings, for depth
    c = circle3(np.zeros(3), zhat, rad)
    ax.plot(*c.T, color=DIM, lw=0.5, alpha=0.35)
for ang in np.arange(0.0, 2 * np.pi, np.pi / 6):
    ax.plot([0, 1.52 * np.cos(ang)], [0, 1.52 * np.sin(ang)], [0, 0],
            color=DIM, lw=0.4, alpha=0.25)
ax.text(0.30, -1.34, 0.0, "ecliptic plane", color=DIM, fontsize=12,
        ha="center")

orbit = circle3(np.zeros(3), zhat, R)
ax.plot(*orbit.T, color=INK, lw=2.0)
ax.text(0.30, 1.02, 0.05, "Earth's orbit", color=INK, fontsize=11,
        ha="center", alpha=0.9)

# direction of motion, in the empty quadrant
a0, a1 = np.deg2rad(118.0), np.deg2rad(142.0)
ax.quiver(R * np.cos(a0), R * np.sin(a0), 0.0,
          R * (np.cos(a1) - np.cos(a0)), R * (np.sin(a1) - np.sin(a0)), 0.0,
          color=INK, lw=1.7, arrow_length_ratio=0.42)

# ── the Sun ─────────────────────────────────────────────────────────────────
ax.plot_surface(*sphere(np.zeros(3), 0.105), color=AMBER, shade=True,
                linewidth=0, antialiased=True, zorder=4)
ax.text(0.0, 0.0, 0.20, "Sun", color=AMBER, fontsize=13, ha="center",
        va="bottom")

# ── the equinox line: where the two planes meet ─────────────────────────────
ax.plot([-1.30, 1.62], [0, 0], [0, 0], color=AMBER, lw=1.6, ls=(0, (7, 4)),
        zorder=6)
ax.quiver(1.30, 0, 0, 0.30, 0, 0, color=AMBER, lw=2.4,
          arrow_length_ratio=0.55, zorder=7)
ax.text(1.60, 0.0, -0.14, r"$\hat{\mathbf{I}}$  (vernal equinox)",
        color=AMBER, fontsize=12.5, ha="right", va="top")

# ── the two Earths ──────────────────────────────────────────────────────────
def draw_earth(centre, label, label_off, show_angle):
    centre = np.asarray(centre, float)

    # equatorial plane: normal is the spin axis, so it cuts the ecliptic along X
    eq = disc(centre, nhat, R_EQ)
    ax.add_collection3d(Poly3DCollection([eq], facecolor=BLUE, alpha=0.22,
                                         edgecolor="none"))
    ax.plot(*eq.T, color=BLUE, lw=1.7)

    ax.plot_surface(*sphere(centre, R_E), color="#4a7fc1", shade=True,
                    linewidth=0, antialiased=True, alpha=1.0)
    eqc = circle3(centre, nhat, R_E * 1.002)          # the equator itself
    ax.plot(*eqc.T, color=CREAM, lw=1.2, alpha=0.9)

    # spin axis
    ax.plot(*np.vstack([centre - 0.20 * nhat, centre + 0.30 * nhat]).T,
            color=MAROON, lw=1.8)
    ax.quiver(*(centre + 0.30 * nhat), *(0.10 * nhat), color=MAROON, lw=1.8,
              arrow_length_ratio=0.9)

    label_off = np.asarray(label_off, float)
    ax.text(*(centre + label_off), label, color=INK, fontsize=12, ha="center",
            va="bottom" if label_off[2] > 0 else "top")

    if show_angle:
        # obliquity, between the spin axis and the ecliptic normal
        ax.plot(*np.vstack([centre, centre + 0.52 * zhat]).T, color=DIM,
                lw=1.1, ls=(0, (4, 3)))
        t = np.linspace(0.0, EPS, 60)
        arcpts = np.array([centre + 0.40 * (np.cos(a) * zhat
                                            + np.sin(a) * np.array([0, 1, 0]))
                           for a in t])
        ax.plot(*arcpts.T, color=MAROON, lw=1.6)
        mid = centre + 0.56 * (np.cos(EPS / 2) * zhat
                               + np.sin(EPS / 2) * np.array([0, 1, 0]))
        ax.text(*mid, r"$\varepsilon = 23.44^\circ$", color=MAROON,
                fontsize=12, ha="left", va="center")


draw_earth([-R, 0, 0], "March equinox", [0.0, 0.0, 0.66], show_angle=True)
draw_earth([R, 0, 0], "September equinox", [-0.28, -0.42, -0.34],
           show_angle=False)

ax.text(-R + 0.06, -0.66, -0.16, "equatorial plane", color=BLUE,
        fontsize=12, ha="center", va="top")
ax.text(0.26, 0.52, 0.78, "the spin axis keeps a fixed direction all year",
        color=MAROON, fontsize=11, ha="center", va="center")

# ── cosmetics ───────────────────────────────────────────────────────────────
L = 1.66
ax.set_xlim(-L, L)
ax.set_ylim(-L, L)
ax.set_zlim(-0.62, 0.80)
ax.set_box_aspect((1.0, 1.0, 1.42 / (2 * L)), zoom=1.52)
ax.set_facecolor(CREAM)
ax.axis("off")

OUT = "img/fig-equinox-frame.png"
fig.savefig(OUT, dpi=160, bbox_inches="tight", facecolor=CREAM)
plt.close(fig)

# bbox_inches="tight" does not trim a 3D axes -- it reports its full box, empty
# margins included -- so crop the flat background off afterwards.
from PIL import Image, ImageChops

im = Image.open(OUT).convert("RGB")
bg = Image.new("RGB", im.size, CREAM)
box = ImageChops.difference(im, bg).getbbox()
if box is not None:
    pad = 14
    box = (max(box[0] - pad, 0), max(box[1] - pad, 0),
           min(box[2] + pad, im.width), min(box[3] + pad, im.height))
    im.crop(box).save(OUT)
print("wrote", OUT, Image.open(OUT).size)
