"""A sun-synchronous orbit: why it needs J2, and what it buys.

A sun-synchronous orbit (SSO) is *not* an orbit that follows the Sun.  It is an
orbit whose **plane** turns at exactly the rate the Sun's right ascension turns,

    Omega_dot = 360 deg / 365.2422 d = +0.9856 deg/day,

so that the angle between the orbit plane and the Earth-Sun line never changes.
The satellite therefore crosses the equator at the same *local solar time* on
every pass, for the life of the mission -- same sun elevation, same shadow
lengths, same illumination in every image, and a solar-array geometry that
never drifts.

Where the turn comes from.  A spherical Earth cannot do it: two-body motion
holds Omega fixed forever.  The equatorial bulge can.  To first order in J2 the
node regresses at (the formula quoted in the orbital-elements deck)

    Omega_dot = -3/2 * sqrt(mu) J2 R_E^2 / [ (1-e^2)^2 a^{7/2} ] * cos i .

The bracket is positive, so the sign is carried entirely by cos i:

    i <  90 deg  ->  Omega_dot < 0,  the node walks west  (prograde orbit)
    i =  90 deg  ->  Omega_dot = 0,  a true polar orbit does not precess
    i >  90 deg  ->  Omega_dot > 0,  the node walks east

Sun-synchrony needs the node to chase the Sun *eastward*, so it needs cos i < 0:
**every sun-synchronous orbit is retrograde**, i ~ 98 deg at low altitude.  That
is the whole content of the "i ~ 98 deg" that gets quoted without explanation.

Left panel   INERTIAL, heliocentric, seen from the ecliptic north pole -- the
             cause.  Earth runs round the Sun, so the Earth-Sun line sweeps a
             full turn in a year.  The blue (SSO) orbit plane turns with it;
             the grey plane, held fixed in inertial space as a spherical Earth
             would hold it, does not.
Right panel  SUN-FIXED, geocentric, seen from the north celestial pole -- the
             effect.  The Sun is pinned to the right, so the local-solar-time
             dial is stationary and can simply be read.  The blue node stays
             nailed to its design local time all year; the grey node walks
             backwards through all 24 hours, about 4 minutes a day.

Angles use the *mean* Sun, uniform in right ascension, which is what local mean
solar time and hence the sun-synchronous condition are defined against, so the
blue local time is constant to the last digit.  The true Sun adds the equation
of time (-14 to +16 min) to the *apparent* time of the crossing; it does not
move the mean-time crossing.

Nothing in the left panel is to scale: the drawn orbit is some 30 times the
Earth-Sun distance divided by its true ratio, and the satellite's motion along
the orbit is a display rate, not 14.4 revolutions a day.

Run from the deck directory:

    python codes/sso-animation.py --verify
    python codes/sso-animation.py --png img/fig-sso-geometry.png
    python codes/sso-animation.py --save img/sso-animation.gif
"""

import argparse
import os

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.patches import Circle, Wedge

CREAM, INK, DIM = "#fcfbf7", "#2f2a24", "#8b857c"
MAROON, BLUE, GREEN, AMBER, PURPLE = "#6b1f1f", "#2f5ea8", "#2d6a2d", "#b45309", "#7e22ce"
SUNY = "#e0a11b"
DAY, NIGHT = "#5b8fd0", "#26333f"      # the two halves of the Earth disc

FROZEN = "#8b857c"                      # the no-J2 comparison orbit reads as grey

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,
})

MU = 398600.4418        # km^3/s^2
J2 = 1.08262668e-3
R_E = 6378.137          # km
YEAR = 365.242190       # tropical year, days
EPS = np.radians(23.439)
OMEGA_DOT_SS = 360.0 / YEAR     # deg/day -- the rate the Sun's RA advances


# ── the J2 node rate, and the inclination that makes it sun-synchronous ──────

def node_rate(a, e, inc):
    """Secular nodal drift from J2, deg/day.  `inc` in radians, `a` in km."""
    return (-1.5 * np.sqrt(MU) * J2 * R_E**2
            / ((1.0 - e**2)**2 * a**3.5) * np.cos(inc)) * 86400.0 * 180.0 / np.pi


def sso_inclination(a, e=0.0):
    """The inclination that makes the node keep pace with the Sun, radians.

    Omega_dot is linear in cos i, so no iteration is needed: divide out the
    bracket and take the arccos.  The result exceeds 90 deg whenever the
    required drift is eastward, which for sun-synchrony it always is.
    """
    bracket = (1.5 * np.sqrt(MU) * J2 * R_E**2
               / ((1.0 - e**2)**2 * a**3.5)) * 86400.0 * 180.0 / np.pi
    c = -OMEGA_DOT_SS / bracket
    if abs(c) > 1.0:
        raise ValueError(f"no sun-synchronous inclination exists at a = {a:.1f} km")
    return np.arccos(c)


def period_minutes(a):
    return 2.0 * np.pi * np.sqrt(a**3 / MU) / 60.0


# ── frames, Sun, and the local time of a node ────────────────────────────────

def mean_sun_ra(t_days, ra0=0.0):
    """Right ascension of the *mean* Sun, radians -- uniform by definition."""
    return (ra0 + 2.0 * np.pi * t_days / YEAR) % (2.0 * np.pi)


def ltan_hours(Om, ra_sun):
    """Local mean solar time of the ascending node, hours in [0, 24).

    A point whose right ascension equals the Sun's is at local noon, and local
    time runs with increasing RA, so the node's clock reading is 12 h plus the
    RA gap converted at 15 deg per hour.
    """
    return (12.0 + np.degrees((Om - ra_sun)) / 15.0) % 24.0


def hhmm(hours):
    h = int(hours) % 24
    m = int(round((hours - int(hours)) * 60.0))
    if m == 60:
        h, m = (h + 1) % 24, 0
    return f"{h:02d}:{m:02d}"


def eq_to_ecl(v):
    """Equatorial components -> ecliptic components (rotation by eps about I)."""
    c, s = np.cos(EPS), np.sin(EPS)
    x, y, z = v[..., 0], v[..., 1], v[..., 2]
    return np.stack([x, y * c + z * s, -y * s + z * c], axis=-1)


def orbit_points(Om, inc, radius, n=240):
    """A circular orbit of the given node and inclination, equatorial frame."""
    u = np.linspace(0.0, 2.0 * np.pi, n)                # argument of latitude
    # r = R3(-Om) R1(-i) [cos u, sin u, 0]: the standard 3-1-3 with omega = 0
    x = np.cos(Om) * np.cos(u) - np.sin(Om) * np.sin(u) * np.cos(inc)
    y = np.sin(Om) * np.cos(u) + np.cos(Om) * np.sin(u) * np.cos(inc)
    z = np.sin(u) * np.sin(inc)
    return radius * np.stack([x, y, z], axis=-1)


def sat_point(Om, inc, radius, u):
    """One point of that orbit, at argument of latitude u."""
    return radius * np.array([
        np.cos(Om) * np.cos(u) - np.sin(Om) * np.sin(u) * np.cos(inc),
        np.sin(Om) * np.cos(u) + np.cos(Om) * np.sin(u) * np.cos(inc),
        np.sin(u) * np.sin(inc)])


# ── panel A: inertial, heliocentric, from the ecliptic north pole ────────────

A_LIM = 1.46            # astronomical units drawn
R_ORB_A = 0.24          # the satellite orbit, wildly enlarged


def draw_helio(ax, t, Om_sso, Om_frz, inc, u_sat):
    ax.clear()
    ax.set_xlim(-A_LIM, A_LIM)
    ax.set_ylim(-A_LIM, A_LIM)
    ax.set_aspect("equal")
    ax.axis("off")

    ra_s = mean_sun_ra(t)
    # Earth sits opposite the Sun as seen from Earth: its heliocentric
    # longitude is the Sun's geocentric longitude plus 180 deg.
    lam_E = ra_s + np.pi
    E = np.array([np.cos(lam_E), np.sin(lam_E)])

    # Earth's orbit, and the fixed inertial direction the whole frame hangs on
    ax.add_patch(Circle((0, 0), 1.0, fill=False, ec=DIM, lw=0.8,
                        ls=(0, (2, 4)), alpha=0.7))
    for m in range(12):                       # a tick a month, to pace the year
        th = 2.0 * np.pi * m / 12.0
        c, s_ = np.cos(th), np.sin(th)
        ax.plot([0.975 * c, 1.025 * c], [0.975 * s_, 1.025 * s_], color=DIM,
                lw=0.8, alpha=0.5)
    ax.annotate("", xy=(0.58, 0), xytext=(0.20, 0),
                arrowprops=dict(arrowstyle="-|>", color=GREEN, lw=1.3,
                                shrinkA=0, shrinkB=0, alpha=0.8))
    ax.text(0.61, 0.04, r"$\hat{\mathbf{I}}$", color=GREEN, fontsize=14,
            ha="left", va="bottom")

    # the Sun
    ax.add_patch(Circle((0, 0), 0.105, color=SUNY, zorder=5))
    ax.add_patch(Circle((0, 0), 0.165, color=SUNY, alpha=0.22, zorder=4))
    ax.text(0, -0.22, "Sun", color=AMBER, fontsize=12, ha="center", va="top")

    # the Earth-Sun line, carried past the Earth into the night side
    ax.plot([0, E[0] * 1.30], [0, E[1] * 1.30], color=DIM, lw=0.9,
            ls=(0, (3, 3)), zorder=2)

    # The view is down the *ecliptic* pole, so every equatorial vector is
    # tilted by eps before it is flattened onto the page.
    def node_dir(Om):
        return eq_to_ecl(np.array([np.cos(Om), np.sin(Om), 0.0]))[:2]

    # The orbit ring is only scene-setting here -- what has to be read at this
    # scale is the *direction* of the node line, so that carries the weight.
    for Om, colour, lw, ls, z in ((Om_frz, FROZEN, 2.0, (0, (5, 3)), 6),
                                  (Om_sso, BLUE, 2.8, "-", 7)):
        p = eq_to_ecl(orbit_points(Om, inc, R_ORB_A))[:, :2] + E
        ax.plot(p[:, 0], p[:, 1], color=colour, lw=0.9, ls="-", alpha=0.45,
                zorder=z)
        nd = node_dir(Om)
        ax.plot([E[0] - 1.42 * R_ORB_A * nd[0], E[0] + 1.42 * R_ORB_A * nd[0]],
                [E[1] - 1.42 * R_ORB_A * nd[1], E[1] + 1.42 * R_ORB_A * nd[1]],
                color=colour, lw=lw, ls=ls, zorder=z)
        ax.plot(E[0] + 1.42 * R_ORB_A * nd[0], E[1] + 1.42 * R_ORB_A * nd[1],
                "o", color=colour, ms=6.5, zorder=z + 2)

    s = eq_to_ecl(sat_point(Om_sso, inc, R_ORB_A, u_sat))[:2] + E
    ax.plot(*s, "o", color=MAROON, ms=5.0, zorder=12)

    ax.add_patch(Circle(E, 0.040, color=DAY, zorder=8, ec=INK, lw=0.6))
    lab = E * (1.0 + 1.42 * R_ORB_A + 0.075)
    ax.text(lab[0], lab[1], "Earth", color=INK, fontsize=10.5, ha="center",
            va="center", zorder=8)


# ── panel B: sun-fixed, geocentric, from the north celestial pole ────────────

B_LIM = 2.28
R_ORB_B = 1.30          # orbit radius in Earth radii, a little enlarged
R_TICK0, R_TICK1, R_RING, R_LAB = 1.56, 1.70, 1.63, 1.88


def draw_clock(ax, Om_sso, Om_frz, inc, ra_s, u_sat):
    ax.clear()
    ax.set_xlim(-B_LIM, B_LIM)
    ax.set_ylim(-B_LIM, B_LIM)
    ax.set_aspect("equal")
    ax.axis("off")

    # The frame follows the Sun: pin the sub-solar direction to +x.  Local time
    # then runs counterclockwise, because that is the way the Earth turns seen
    # from the north, carrying a place from dawn through noon to dusk.
    def ang(Om):
        return (Om - ra_s) % (2.0 * np.pi)

    ax.add_patch(Wedge((0, 0), 1.0, -90, 90, color=DAY, zorder=3))
    ax.add_patch(Wedge((0, 0), 1.0, 90, 270, color=NIGHT, zorder=3))
    ax.add_patch(Circle((0, 0), 1.0, fill=False, ec=INK, lw=0.8, zorder=4))
    ax.plot([0, 0], [-1, 1], color=CREAM, lw=1.0, alpha=0.5, zorder=4)

    ax.annotate("", xy=(2.24, 0), xytext=(2.00, 0),
                arrowprops=dict(arrowstyle="-|>", color=SUNY, lw=2.2,
                                shrinkA=0, shrinkB=0))
    ax.text(2.10, 0.09, "to Sun", color=AMBER, fontsize=11.5, ha="center",
            va="bottom")

    # the local-time dial
    ax.add_patch(Circle((0, 0), R_RING, fill=False, ec=DIM, lw=0.6, alpha=0.4))
    for h in range(24):
        th = np.radians((h - 12) * 15.0)
        c, s = np.cos(th), np.sin(th)
        major = (h % 3 == 0)
        ax.plot([R_TICK0 * c, R_TICK1 * c], [R_TICK0 * s, R_TICK1 * s],
                color=DIM, lw=1.4 if major else 0.7, alpha=0.9 if major else 0.5)
        if major:
            ax.text(R_LAB * c, R_LAB * s, f"{h:02d}", color=DIM, fontsize=9.5,
                    ha="center", va="center")

    # the two orbits.  Seen down the pole a near-polar orbit projects to a thin
    # ellipse whose long axis *is* the node line, which is what we want read.
    for Om, colour, lw, ls, z in ((Om_frz, FROZEN, 1.6, (0, (5, 3)), 5),
                                  (Om_sso, BLUE, 2.4, "-", 6)):
        a = ang(Om)
        p = orbit_points(Om, inc, R_ORB_B)
        ca, sa = np.cos(a - Om), np.sin(a - Om)
        ax.plot(p[:, 0] * ca - p[:, 1] * sa, p[:, 0] * sa + p[:, 1] * ca,
                color=colour, lw=lw, ls=ls, zorder=z)
        for th, fill in ((a, True), (a + np.pi, False)):
            x, y = R_ORB_B * np.cos(th), R_ORB_B * np.sin(th)
            if fill:
                ax.plot(x, y, "o", color=colour, ms=9, zorder=z + 3)
            else:
                ax.plot(x, y, "o", mfc=CREAM, mec=colour, mew=1.6, ms=7,
                        zorder=z + 3)
        # the pointer: this node's reading on the dial
        c, s = np.cos(a), np.sin(a)
        ax.plot([R_TICK0 * c, R_TICK1 * c], [R_TICK0 * s, R_TICK1 * s],
                color=colour, lw=3.4, zorder=z + 1, solid_capstyle="butt")
        ax.plot([R_ORB_B * c, R_TICK0 * c], [R_ORB_B * s, R_TICK0 * s],
                color=colour, lw=0.8, ls=(0, (1, 2)), alpha=0.7, zorder=z)

    a_sso = ang(Om_sso)
    sp = sat_point(Om_sso, inc, R_ORB_B, u_sat)
    ca, sa = np.cos(a_sso - Om_sso), np.sin(a_sso - Om_sso)
    ax.plot(sp[0] * ca - sp[1] * sa, sp[0] * sa + sp[1] * ca, "o",
            color=MAROON, ms=7, zorder=14)




# ── assembling one frame ─────────────────────────────────────────────────────

def make_figure():
    fig = plt.figure(figsize=(12.0, 6.1))
    xA, xB = 0.228, 0.733
    axA = fig.add_axes([xA - 0.164, 0.140, 0.328, 0.640])
    axB = fig.add_axes([xB - 0.164, 0.140, 0.328, 0.640])
    head = fig.text(0.5, 0.955, "", ha="center", fontsize=14.5, color=INK)
    fig.text(xA, 0.888, "the cause: the Earth$-$Sun line turns "
             r"$360^\circ$ in a year", ha="center", fontsize=12.5, color=INK)
    fig.text(xA, 0.845, "inertial frame, from the ecliptic north pole",
             ha="center", fontsize=10, color=DIM)
    fig.text(xB, 0.888, "the effect: the equator crossing keeps its clock time",
             ha="center", fontsize=12.5, color=INK)
    fig.text(xB, 0.845, "frame pinned to the Sun, from the north celestial pole",
             ha="center", fontsize=10, color=DIM)
    sub = fig.text(xA, 0.800, "", ha="center", fontsize=11, color=DIM)
    rd1 = fig.text(xB - 0.055, 0.800, "", ha="right", fontsize=11.5, color=BLUE)
    rd2 = fig.text(xB - 0.040, 0.800, "", ha="left", fontsize=11.5, color=FROZEN)
    fig.text(xA, 0.105, "solid blue  $-$  $J_2$ turns the plane with the Sun\n"
             "dashed grey  $-$  no $J_2$: the plane is frozen in space",
             ha="center", va="top", fontsize=10, color=DIM, linespacing=1.5)
    fig.text(xB, 0.105, "the dial is local solar time; the Earth turns\n"
             r"counterclockwise, dawn $\to$ noon $\to$ dusk",
             ha="center", va="top", fontsize=10, color=DIM, linespacing=1.5)
    return fig, axA, axB, head, sub, rd1, rd2


def render_frame(axA, axB, texts, t, Om0, inc, alt, u_sat, ltan0):
    head, sub, rd1, rd2 = texts
    ra_s = mean_sun_ra(t)
    Om_sso = (Om0 + np.radians(OMEGA_DOT_SS * t)) % (2.0 * np.pi)
    Om_frz = Om0                                   # a round Earth freezes it
    draw_helio(axA, t, Om_sso, Om_frz, inc, u_sat)
    draw_clock(axB, Om_sso, Om_frz, inc, ra_s, u_sat)
    lt_s = ltan_hours(Om_sso, ra_s)
    lt_f = ltan_hours(Om_frz, ra_s)
    head.set_text(f"A sun-synchronous orbit   $-$   day "
                  f"{int(round(t)):3d} of the year")
    sub.set_text(rf"$h = {alt:.0f}$ km,   $i = {np.degrees(inc):.2f}^\circ$,   "
                 rf"$T = {period_minutes(R_E + alt):.1f}$ min,   "
                 rf"$\dot\Omega = +{OMEGA_DOT_SS:.4f}^\circ$/day")
    rd1.set_text(f"node at {hhmm(lt_s)}")
    rd2.set_text(f"   ·   {hhmm(lt_f)} without $J_2$")
    return lt_s, lt_f


# ── outputs ──────────────────────────────────────────────────────────────────

def save_png(path, alt=700.0, day=137.0, ltan=13.5):
    a = R_E + alt
    inc = sso_inclination(a)
    Om0 = (mean_sun_ra(0.0) + np.radians((ltan - 12.0) * 15.0)) % (2.0 * np.pi)
    fig, axA, axB, *texts = make_figure()
    render_frame(axA, axB, texts, day, Om0, inc, alt, 2.1, ltan)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    fig.savefig(path, dpi=150, facecolor=CREAM)
    plt.close(fig)
    print("wrote", path)


def render(path, alt=700.0, step=4.0, ltan=13.5, fps=14, dpi=92):
    a = R_E + alt
    inc = sso_inclination(a)
    Om0 = (mean_sun_ra(0.0) + np.radians((ltan - 12.0) * 15.0)) % (2.0 * np.pi)
    days = np.arange(0.0, YEAR + step, step)
    fig, axA, axB, *texts = make_figure()

    def frame(k):
        # the satellite really goes round ~14.4 times a day; at four days a
        # frame no honest phase is drawable, so this is a display rate only.
        render_frame(axA, axB, texts, days[k], Om0, inc, alt, 0.9 * k, ltan)
        return ()

    ani = animation.FuncAnimation(fig, frame, frames=len(days), 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(days)} 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")


def verify():
    print("sun-synchronous inclination, circular orbits\n")
    print(f"{'alt km':>8}{'a km':>11}{'i deg':>9}{'T min':>9}"
          f"{'Omega_dot deg/day':>20}")
    for alt in (300.0, 400.0, 600.0, 700.0, 800.0, 1000.0, 1200.0):
        a = R_E + alt
        inc = sso_inclination(a)
        print(f"{alt:8.0f}{a:11.2f}{np.degrees(inc):9.3f}"
              f"{period_minutes(a):9.2f}{node_rate(a, 0.0, inc):20.6f}")
    print(f"\n  target rate 360/{YEAR} = {OMEGA_DOT_SS:.6f} deg/day"
          "   (every inclination above exceeds 90 deg: retrograde)")

    print("\nthe sign of the drift is carried by cos i  (a = 7078.137 km)\n")
    for ideg in (0.0, 45.0, 89.0, 90.0, 91.0, 98.19, 135.0, 180.0):
        r = node_rate(7078.137, 0.0, np.radians(ideg))
        sense = "east" if r > 1e-9 else ("west" if r < -1e-9 else "none")
        print(f"  i = {ideg:6.2f} deg   Omega_dot = {r:+9.5f} deg/day   {sense}")

    print("\nlocal time of the ascending node over one year\n")
    alt, ltan = 700.0, 13.5
    a = R_E + alt
    inc = sso_inclination(a)
    Om0 = (mean_sun_ra(0.0) + np.radians((ltan - 12.0) * 15.0)) % (2 * np.pi)
    t = np.linspace(0.0, YEAR, 2000)
    ra = mean_sun_ra(t)
    lt_s = ltan_hours(Om0 + np.radians(OMEGA_DOT_SS * t), ra)
    lt_f = ltan_hours(Om0 * np.ones_like(t), ra)
    dev = np.max(np.abs((lt_s - ltan + 12.0) % 24.0 - 12.0)) * 60.0
    print(f"  sun-synchronous : {hhmm(lt_s[0])} -> {hhmm(lt_s[-1])}, "
          f"worst departure from {hhmm(ltan)} over the year = {dev:.2e} min")
    span = np.unwrap(np.radians(lt_f * 15.0)) / np.radians(15.0)
    print(f"  no J2           : {hhmm(lt_f[0])} -> {hhmm(lt_f[-1])}, "
          f"total walk = {span[-1] - span[0]:+.3f} h  "
          f"({(span[-1] - span[0]) * 60.0 / YEAR:+.2f} min/day)")

    print("\nthe J2 rate against the deck's own numbers\n")
    a = R_E + 700.0
    inc = sso_inclination(a)
    print(f"  i = {np.degrees(inc):.3f} deg gives Omega_dot = "
          f"{node_rate(a, 0.0, inc):.6f} deg/day, wanted {OMEGA_DOT_SS:.6f}")
    print(f"  residual = {abs(node_rate(a, 0.0, inc) - OMEGA_DOT_SS):.3e} deg/day")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--alt", type=float, default=700.0, help="altitude, km")
    ap.add_argument("--ltan", type=float, default=13.5,
                    help="design local time of the ascending node, hours")
    ap.add_argument("--step", type=float, default=4.0, help="days per frame")
    ap.add_argument("--fps", type=int, default=14)
    ap.add_argument("--dpi", type=int, default=92)
    ap.add_argument("--verify", action="store_true")
    ap.add_argument("--png", default=None)
    ap.add_argument("--save", default=None)
    args = ap.parse_args()

    if args.verify:
        verify()
    if args.png:
        save_png(args.png, alt=args.alt, ltan=args.ltan)
    if args.save:
        render(args.save, alt=args.alt, step=args.step, ltan=args.ltan,
               fps=args.fps, dpi=args.dpi)
    if not (args.verify or args.png or args.save):
        verify()
