"""The Sun's year, watched from a point on the equator at noon.

The geocentric equatorial frame is defined by two planes and one direction in
their intersection.  The cheapest way to *see* that direction is to stand on
the equator and look at the Sun at local noon, every day for a year.

    * Declination is the angle from the celestial equator, measured along the
      object's hour circle.  At *apparent* noon the Sun is on the observer's
      meridian, which is an hour circle, and on the equator the zenith itself
      lies on the celestial equator (delta_zenith = phi = 0).  The two arcs
      coincide, so the noon zenith distance is exactly |delta|.  (At latitude
      phi it would be |phi - delta|.)  Every frame here is apparent noon, so
      the Sun sits on the meridian and the equation of time never enters.
    * Over a year delta swings between +23.44 deg and -23.44 deg: the Sun is
      not fixed in the equatorial plane, it runs round the *ecliptic*.
    * Twice a year delta = 0 and the Sun stands exactly overhead -- a vertical
      rod casts no shadow at noon.  Those are the equinoxes: the two days on
      which the Sun lies *in* the equatorial plane.
    * The crossing on which the Sun is moving north (delta going from - to +)
      is the **vernal equinox**.  The Earth-to-Sun direction at that instant is
      Ihat.

The apparent Sun is taken from the low-precision series of the Astronomical
Almanac (about 0.01 deg over 1950--2050), so the equinox dates that come out
of the animation are the real ones for the year requested.

Left panel   geocentric equatorial frame: equator (blue) and ecliptic (amber)
             are two great circles inclined by eps, meeting at the two nodes;
             the Sun walks the ecliptic and its (alpha, delta) are read off.
Top right    what the observer sees on the meridian at noon, with the shadow
             of a vertical rod.
Bottom right delta against the date, with the two zero crossings.

Run from the deck directory:

    python codes/sun-year-animation.py
    python codes/sun-year-animation.py --year 2026 --step 3 --fps 15
"""

import argparse
import datetime as dt
import os

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

CREAM, INK, DIM = "#fcfbf7", "#2f2a24", "#8b857c"
MAROON, BLUE, GREEN, AMBER, PURPLE = "#6b1f1f", "#2f5ea8", "#2d6a2d", "#b45309", "#7e22ce"
SUNY = "#e0a11b"

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


# ── the apparent Sun ─────────────────────────────────────────────────────────

def julian_day(year, month, day_frac):
    """Julian date of a Gregorian calendar date; day_frac counts from midnight."""
    if month <= 2:
        year, month = year - 1, month + 12
    a = year // 100
    b = 2 - a + a // 4
    return (int(365.25 * (year + 4716)) + int(30.6001 * (month + 1))
            + day_frac + b - 1524.5)


def sun_apparent(jd):
    """Apparent right ascension, declination and obliquity (radians).

    Astronomical Almanac low-precision series.  Also returns the equation of
    time in degrees, EoT = apparent solar time - mean solar time.
    """
    n = np.asarray(jd, float) - 2451545.0
    L = np.radians(280.460 + 0.9856474 * n)              # mean longitude
    g = np.radians(357.528 + 0.9856003 * n)              # mean anomaly
    lam = L + np.radians(1.915) * np.sin(g) + np.radians(0.020) * np.sin(2 * g)
    eps = np.radians(23.439 - 4.0e-7 * n)                # obliquity
    alpha = np.arctan2(np.cos(eps) * np.sin(lam), np.cos(lam))
    delta = np.arcsin(np.sin(eps) * np.sin(lam))
    eot = (np.degrees(L - alpha) + 180.0) % 360.0 - 180.0
    return alpha % (2 * np.pi), delta, eps, eot


def unit_ijk(alpha, delta):
    """Geocentric unit vector to the Sun in the equatorial frame."""
    return np.array([np.cos(delta) * np.cos(alpha),
                     np.cos(delta) * np.sin(alpha),
                     np.sin(delta)])


def crossings(jd0, ndays):
    """Days (fractional, from jd0) at which delta changes sign, and the sense."""
    t = np.arange(0.0, ndays, 0.01)
    d = sun_apparent(jd0 + t)[1]
    out = []
    for k in np.nonzero(np.diff(np.sign(d)))[0]:
        # linear root of delta between the two samples
        f = d[k] / (d[k] - d[k + 1])
        out.append((t[k] + f * (t[k + 1] - t[k]), np.sign(d[k + 1])))
    return out


# ── geometry helpers ─────────────────────────────────────────────────────────

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


def ring(normal, radius, n=240):
    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 @ normal) * normal
    u /= np.linalg.norm(u)
    w = np.cross(normal, u)
    t = np.linspace(0.0, 2 * np.pi, n)
    return (np.cos(t)[:, None] * u + np.sin(t)[:, None] * w) * radius


def ra_arc(alpha, radius, n=120):
    """The angle alpha itself: from Ihat, counterclockwise, in the equator."""
    t = np.linspace(0.0, alpha, n)
    return radius * np.column_stack([np.cos(t), np.sin(t), np.zeros_like(t)])


def dec_arc(alpha, delta, radius, n=60):
    """The angle delta: from the equator up to the Sun, along its meridian."""
    t = np.linspace(0.0, delta, n)[:, None]
    p = np.array([np.cos(alpha), np.sin(alpha), 0.0])
    k = np.array([0.0, 0.0, 1.0])
    return radius * (np.cos(t) * p + np.sin(t) * k)


# ── the three panels ─────────────────────────────────────────────────────────

R_E = 0.13          # Earth, some 10^4 times too large
R_SKY = 1.0         # radius of the drawn celestial sphere


def draw_frame3d(ax, alpha, delta, ecl, trail, is_equinox):
    ax.clear()
    ax.set_axis_off()
    ax.view_init(elev=20, azim=-56)

    # equatorial plane, with the equator ruled on it
    eq_disc = ring([0, 0, 1], 1.30)
    ax.add_collection3d(Poly3DCollection([eq_disc], facecolor=BLUE, alpha=0.10,
                                         edgecolor="none"))
    ax.plot(*eq_disc.T, color=BLUE, lw=0.5, alpha=0.35)
    eq = ring([0, 0, 1], R_SKY)
    ax.plot(*eq.T, color=BLUE, lw=2.0)
    ax.text(0.10, -1.44, -0.06, "celestial equator", color=BLUE, fontsize=11,
            ha="center")

    # the ecliptic: the Sun's apparent path, tilted by the obliquity
    ax.plot(*ecl.T, color=AMBER, lw=1.6, alpha=0.55)
    ax.text(-1.14, 0.80, 0.48, "ecliptic", color=AMBER, fontsize=11.5,
            ha="center")

    # the line of nodes -- the intersection of the two planes
    ax.plot([-1.24, 1.42], [0, 0], [0, 0], color=GREEN, lw=1.4, ls=(0, (7, 4)))
    ax.quiver(1.02, 0, 0, 0.40, 0, 0, color=GREEN, lw=2.4, arrow_length_ratio=0.5)
    ax.text(1.50, 0.0, -0.12, r"$\hat{\mathbf{I}}$", color=GREEN, fontsize=15,
            ha="left", va="top")
    ax.scatter([1.0], [0.0], [0.0], s=34, color=GREEN, depthshade=False)
    ax.text(0.99, 0.10, -0.30, "vernal\nequinox", color=GREEN, fontsize=10,
            ha="center", va="top")
    ax.scatter([-1.0], [0.0], [0.0], s=26, color=GREEN, alpha=0.55,
               depthshade=False)

    # spin axis
    ax.plot([0, 0], [0, 0], [0, 1.28], color=MAROON, lw=1.6)
    ax.quiver(0, 0, 1.22, 0, 0, 0.22, color=MAROON, lw=2.0, arrow_length_ratio=0.6)
    ax.text(0.0, 0.06, 1.50, r"$\hat{\mathbf{K}}$", color=MAROON, fontsize=15,
            ha="center")

    # the Earth
    ax.plot_surface(*sphere(np.zeros(3), R_E), color="#4a7fc1", shade=True,
                    linewidth=0, antialiased=True)
    ax.plot(*(ring([0, 0, 1], R_E * 1.003)).T, color=CREAM, lw=1.0, alpha=0.9)

    # the Sun: direction, right ascension and declination
    s = unit_ijk(alpha, delta) * R_SKY
    p = np.array([s[0], s[1], 0.0])                    # projection on the equator
    ax.plot(*trail.T, color=AMBER, lw=2.6)
    ax.plot([0, s[0]], [0, s[1]], [0, s[2]], color=DIM, lw=0.9, ls=(0, (3, 3)))
    ax.plot([s[0], p[0]], [s[1], p[1]], [s[2], p[2]], color=SUNY, lw=1.2,
            ls=(0, (2, 2)))
    ax.plot(*ra_arc(alpha, 0.52).T, color=BLUE, lw=2.0)
    ax.plot(*dec_arc(alpha, delta, 0.74).T, color=SUNY, lw=1.8)
    ax.text(*ra_arc(alpha, 0.62)[60], r"$\alpha$", color=BLUE,
            fontsize=13, ha="center", va="center")
    ax.text(*dec_arc(alpha, delta, 0.86)[-1], r"$\delta$", color=SUNY,
            fontsize=13, ha="center", va="center")
    ax.plot_surface(*sphere(s, 0.075), color=SUNY, shade=True, linewidth=0,
                    antialiased=True)

    if is_equinox:
        ax.text(0.0, 0.0, -0.86, r"$\delta = 0$: the Sun is in the equatorial plane",
                color=GREEN, fontsize=12.5, ha="center")

    ax.set_xlim(-1.35, 1.35)
    ax.set_ylim(-1.35, 1.35)
    ax.set_zlim(-0.95, 1.55)
    ax.set_box_aspect((1.0, 1.0, 2.5 / 2.7), zoom=1.42)


def draw_sky(ax, delta, past, is_equinox):
    ax.clear()
    ax.set_aspect("equal")
    ax.axis("off")

    th = np.linspace(0.0, np.pi, 200)
    ax.plot(np.cos(th), np.sin(th), color=DIM, lw=1.0)                # the dome
    ax.plot([-1.16, 1.16], [0, 0], color=INK, lw=1.6)                 # horizon
    ax.plot([0, 0], [0, 1.12], color=DIM, lw=0.9, ls=(0, (3, 3)))     # zenith line
    ax.text(0.0, 1.16, "zenith", color=DIM, fontsize=10.5, ha="center")
    ax.text(-1.18, -0.055, "N", color=INK, fontsize=12, ha="right", va="top")
    ax.text(1.18, -0.055, "S", color=INK, fontsize=12, ha="left", va="top")

    # the band the noon Sun is confined to: |delta| <= eps
    e = np.radians(23.44)
    for d0, lab, dy in ((e, "June solstice", 0.10), (-e, "December solstice", 0.10)):
        ax.plot([-np.sin(d0)], [np.cos(d0)], marker="|", color=AMBER, ms=9)
        ax.text(-np.sin(d0) * 1.19, np.cos(d0) * 1.19 - 0.02, lab, color=AMBER,
                fontsize=9.5, ha="center")

    # every noon position so far -- the dots crowd where the Sun turns round
    if len(past):
        ax.plot(-np.sin(past), np.cos(past), ls="none", marker="o", ms=2.6,
                color=AMBER, alpha=0.30)

    # a vertical rod and its noon shadow
    h = 0.26
    ax.plot([0, 0], [0, h], color=INK, lw=3.0, solid_capstyle="butt")
    tip = h * np.tan(delta)
    ax.plot([0, tip], [0, 0], color=MAROON, lw=4.0, solid_capstyle="butt",
            alpha=0.85)
    ax.plot([-np.sin(delta), tip], [np.cos(delta), 0.0], color=SUNY, lw=0.8,
            ls=(0, (4, 3)))

    ax.plot([-np.sin(delta)], [np.cos(delta)], marker="o", ms=13, color=SUNY,
            mec=AMBER, mew=1.0)

    ax.text(0.0, -0.30, "the Sun is at the zenith: no shadow at noon" if is_equinox
            else f"noon Sun {abs(np.degrees(delta)):.1f}$^\\circ$ from the zenith,"
                 f" {'north' if delta > 0 else 'south'}",
            color=GREEN if is_equinox else INK, fontsize=11.5, ha="center")
    ax.set_title("apparent noon, on the equator", fontsize=12, color=INK, pad=2)
    ax.set_xlim(-1.30, 1.30)
    ax.set_ylim(-0.42, 1.34)


def draw_dec(ax, days, decs, k, eqx, year):
    ax.clear()
    ax.plot(days, np.degrees(decs), color=AMBER, lw=1.1, alpha=0.45)
    ax.plot(days[:k + 1], np.degrees(decs[:k + 1]), color=AMBER, lw=2.0)
    ax.axhline(0.0, color=BLUE, lw=1.4)
    ax.text(112.0, 2.4, "celestial equator", color=BLUE, fontsize=9.5)

    for t0, sense in eqx:
        ax.axvline(t0, color=GREEN, lw=1.0, ls=(0, (4, 3)))
        d0 = dt.datetime(year, 1, 1, 12) + dt.timedelta(days=float(t0))
        ax.text(t0 + 4.0, -20.5 if sense > 0 else 12.0,
                ("vernal equinox\n" if sense > 0 else "autumnal equinox\n")
                + d0.strftime("%d %b"), color=GREEN, fontsize=9.5,
                ha="left", va="bottom" if sense > 0 else "top")

    ax.plot([days[k]], [np.degrees(decs[k])], marker="o", ms=8, color=SUNY,
            mec=AMBER)
    ax.set_xlim(0, days[-1])
    ax.set_ylim(-29, 29)
    ax.set_yticks([-23.44, 0, 23.44])
    ax.set_yticklabels([r"$-23.44^\circ$", r"$0$", r"$+23.44^\circ$"], fontsize=9.5)
    starts = [(dt.date(year, m, 1) - dt.date(year, 1, 1)).days for m in range(1, 13)]
    ax.set_xticks(starts)
    ax.set_xticklabels(list("JFMAMJJASOND"), fontsize=9.5)
    ax.set_ylabel(r"declination $\delta$", fontsize=10.5)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)


# ── the animation ────────────────────────────────────────────────────────────

def build(year, step, path, fps, dpi, hold):
    jd0 = julian_day(year, 1, 1.5)                    # 1 January, 12h UT
    ndays = 366 if (year % 4 == 0 and year % 100 != 0) or year % 400 == 0 else 365

    days = np.arange(0.0, ndays, 1.0)
    alphas, decs, _, _ = sun_apparent(jd0 + days)
    ecl = np.array([unit_ijk(a, d) for a, d in zip(alphas, decs)]) * R_SKY
    eqx = crossings(jd0, ndays)

    order = list(range(0, ndays, step))
    if order[-1] != ndays - 1:
        order.append(ndays - 1)

    # linger on the four days that name the seasons
    keys = [t for t, _ in eqx] + [float(np.argmax(decs)), float(np.argmin(decs))]
    stops = {min(order, key=lambda j: abs(j - key)) for key in keys}
    equinox_frames = {min(order, key=lambda j: abs(j - t)) for t, _ in eqx}
    seq = []
    for k in order:
        seq += [k] * (1 + hold * (k in stops))

    fig = plt.figure(figsize=(12.0, 6.1))
    gs = fig.add_gridspec(2, 2, width_ratios=[1.30, 1.0], height_ratios=[1.0, 0.66],
                          left=0.005, right=0.975, bottom=0.085, top=0.895,
                          wspace=0.02, hspace=0.34)
    ax3 = fig.add_subplot(gs[:, 0], projection="3d")
    axs = fig.add_subplot(gs[0, 1])
    axd = fig.add_subplot(gs[1, 1])
    ax3.set_facecolor(CREAM)

    head = fig.text(0.5, 0.955, "", ha="center", fontsize=14, color=INK)
    sub = fig.text(0.5, 0.915, "", ha="center", fontsize=11.5, color=DIM)

    def frame(n):
        k = seq[n]
        d = decs[k]
        is_eqx = k in equinox_frames
        draw_frame3d(ax3, alphas[k], d, ecl, ecl[:k + 1], is_eqx)
        draw_sky(axs, d, decs[:k + 1], is_eqx)
        draw_dec(axd, days, decs, k, eqx, year)
        date = dt.date(year, 1, 1) + dt.timedelta(days=int(k))
        head.set_text(date.strftime("%d %B %Y") + "   $-$   the Sun at apparent"
                      " noon, seen from the equator")
        sub.set_text(rf"$\alpha = {np.degrees(alphas[k]):.1f}^\circ$,   "
                     rf"$\delta = {np.degrees(d):+.2f}^\circ$")
        return ()

    ani = animation.FuncAnimation(fig, frame, frames=len(seq), blit=False)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    writer = "ffmpeg" if path.endswith(".mp4") else "pillow"
    print(f"writing {path} ({len(seq)} frames) ...")
    ani.save(path, writer=writer, fps=fps, dpi=dpi,
             savefig_kwargs=dict(facecolor=CREAM))
    plt.close(fig)
    print("done:", path, f"{os.path.getsize(path)/1e6:.2f} MB")
    for t, sense in eqx:
        # t counts from 1 January 12h UT
        when = dt.datetime(year, 1, 1, 12) + dt.timedelta(days=float(t))
        print(f"  delta = 0 going {'north' if sense > 0 else 'south'}: "
              f"{when.strftime('%d %b %Y  %H:%M')} UT")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--year", type=int, default=2026)
    ap.add_argument("--step", type=int, default=3, help="days per frame")
    ap.add_argument("--hold", type=int, default=8,
                    help="extra frames held at the equinoxes and solstices")
    ap.add_argument("--fps", type=int, default=14)
    ap.add_argument("--dpi", type=int, default=92)
    ap.add_argument("--save", default="img/sun-year-animation.gif")
    a = ap.parse_args()
    build(a.year, a.step, a.save, a.fps, a.dpi, a.hold)
