"""Time note: two animations.

  img/anim-time-sidereal.gif   Solar day vs sidereal day, on a planet whose year is only 24 days long
                               (so that the extra turn is a visible 15 deg per day instead of 0.9856 deg).
  img/anim-time-earth-mars.gif Earth-Mars geometry over one synodic period and the light time between them.

Run from the deck directory:   python codes/anim-time.py [sidereal|mars|all]
"""
import sys
import warnings
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Wedge
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _gif import save_gif
from _timeutil import C_KMS, AU, jd_gregorian, calendar_from_jd, heliocentric

warnings.filterwarnings("ignore")
which = sys.argv[1] if len(sys.argv) > 1 else "all"


# ================================================================== solar vs sidereal day
def sidereal_animation(path="img/anim-time-sidereal.gif", N=24, fps=16, hold=24):
    per_day = 16
    total = N*per_day
    ts = np.concatenate([np.zeros(8), np.linspace(0, N, total + 1), np.full(hold, N)])   # solar days
    th_sun0 = np.radians(200.0)                     # Earth's polar angle at t = 0 (Sun is at the origin)
    orb = lambda t: th_sun0 + 2*np.pi*t/N
    spin = lambda t: (th_sun0 - np.pi) + 2*np.pi*(1 + 1/N)*t         # flag angle, inertial
    star_dir = th_sun0 - np.pi                        # direction of a star that lines up with the Sun at t = 0

    fig = plt.figure(figsize=(11.2, 5.6))
    gs = fig.add_gridspec(2, 2, width_ratios=[1, 1.08], height_ratios=[1, 0.5], wspace=0.08, hspace=0.32)
    ax = fig.add_subplot(gs[:, 0]); ax.set_aspect("equal"); ax.axis("off")
    ax.set_xlim(-1.5, 1.5); ax.set_ylim(-1.5, 1.5)
    ax.add_patch(Circle((0, 0), 0.16, color="#f5b301", ec=AMBER, lw=1.2, zorder=3))
    ax.text(0, 0, "Sun", ha="center", va="center", fontsize=9, zorder=4)
    tt = np.linspace(0, 2*np.pi, 300)
    ax.plot(np.cos(tt), np.sin(tt), color="#b8b0a2", lw=0.8, ls="--")
    R = 0.19
    disc = Circle((0, 0), R, color="#cfe3f5", ec=INK, lw=1.2, zorder=5); ax.add_patch(disc)
    night = Wedge((0, 0), R, 0, 1, color="#31435f", alpha=0.55, zorder=6); ax.add_patch(night)
    flag, = ax.plot([], [], color=MAROON, lw=3, zorder=8)
    dot, = ax.plot([], [], "o", color=MAROON, ms=6, zorder=9)
    sunline, = ax.plot([], [], color=AMBER, lw=1.6, zorder=2)
    starline, = ax.plot([], [], color=PURPLE, lw=1.6, ls="--", zorder=2)
    ax.text(0, 1.47, "Top view (a year of only 24 days, so the effect is visible)", fontsize=9.6, va="top", ha="center")
    ax.text(-1.45, -1.35, "purple dashed: direction to a distant star (fixed)\namber: direction to the Sun", fontsize=9, va="bottom")
    flash = Circle((0, 0), R*1.9, fill=False, ec=AMBER, lw=2.5, zorder=7, visible=False); ax.add_patch(flash)
    flash2 = Circle((0, 0), R*1.9, fill=False, ec=PURPLE, lw=2.5, ls="--", zorder=7, visible=False); ax.add_patch(flash2)
    
    # timeline
    tl = fig.add_subplot(gs[0, 1]); tl.set_xlim(-0.6, N + 0.6); tl.set_ylim(-1.15, 1.15)
    for s_ in ("top", "right", "left"): tl.spines[s_].set_visible(False)
    tl.set_yticks([]); tl.set_xlabel("solar days since the start"); tl.set_xticks(range(0, N + 1, 4))
    tl.axhline(0.35, color=AMBER, lw=0.5, alpha=0.5); tl.axhline(-0.35, color=PURPLE, lw=0.5, alpha=0.5)
    tl.text(-0.4, 0.62, "solar noons (Sun on your meridian)", color=AMBER, fontsize=9.6, va="bottom")
    tl.text(-0.4, -0.62, "star transits (Earth turned exactly 360°)", color=PURPLE, fontsize=9.6, va="top")
    noon_pts, = tl.plot([], [], "|", color=AMBER, ms=20, mew=3)
    star_pts, = tl.plot([], [], "|", color=PURPLE, ms=20, mew=3)
    tl.set_title("The two kinds of day, counted", fontsize=11)
    # text panel
    tp = fig.add_subplot(gs[1, 1]); tp.axis("off")
    msg = tp.text(0.0, 1.0, "", va="top", fontsize=10.5)
    counter = tp.text(0.0, 0.0, "", va="bottom", fontsize=11, family="monospace")

    def frame(i):
        t = ts[i]
        a = orb(t); P = np.array([np.cos(a), np.sin(a)])
        disc.center = P
        sd = np.degrees(a + np.pi)                                 # direction Earth -> Sun
        night.set_center(P); night.set_theta1(sd + 90); night.set_theta2(sd + 270)
        s_ = spin(t)
        flag.set_data([P[0] + R*np.cos(s_), P[0] + 1.7*R*np.cos(s_)], [P[1] + R*np.sin(s_), P[1] + 1.7*R*np.sin(s_)])
        dot.set_data([P[0] + R*np.cos(s_)], [P[1] + R*np.sin(s_)])
        su = np.array([np.cos(a + np.pi), np.sin(a + np.pi)]); st = np.array([np.cos(star_dir), np.sin(star_dir)])
        sunline.set_data([P[0], P[0] + 0.55*su[0]], [P[1], P[1] + 0.55*su[1]])
        starline.set_data([P[0], P[0] + 0.55*st[0]], [P[1], P[1] + 0.55*st[1]])
        n_sol, n_sid = int(np.floor(t + 1e-9)), int(np.floor((1 + 1/N)*t + 1e-9))
        counter.set_text("t = %5.2f d:  %2d solar, %2d sidereal days" % (t, n_sol, n_sid))
        ev_noon = np.arange(0, n_sol + 1); ev_star = np.arange(0, n_sid + 1)/(1 + 1/N)
        noon_pts.set_data(ev_noon, np.full_like(ev_noon, 0.35, float)); star_pts.set_data(ev_star, np.full_like(ev_star, -0.35))
        near_noon = min(abs(t - round(t)), 1) < 0.5/per_day and t > 0.05 and i > 8
        near_star = min(abs((1 + 1/N)*t - round((1 + 1/N)*t)), 1) < 0.5/per_day and t > 0.05 and i > 8
        flash.center = P; flash2.center = P
        flash.set_visible(bool(near_noon)); flash2.set_visible(bool(near_star))
        if t >= N - 1e-9 and i > len(ts) - hold + 2:
            msg.set_text("After one year (24 solar days) the Earth has turned\n25 times relative to the stars.\n"
                         "The real ratio is 366.2422 sidereal days\nin 365.2422 solar days.")
        else:
            msg.set_text("A sidereal day is shorter than a solar day\nby 1/%d of a day here (3 m 56 s on the real Earth)." % N)
        return []

    save_gif(fig, frame, len(ts), fps, path)


# ================================================================== Earth-Mars light time
def mars_animation(path="img/anim-time-earth-mars.gif", fps=15, step=4.0):
    jd0 = jd_gregorian(2026, 1, 1)
    days = np.arange(0, 781, step)
    re_, _ = heliocentric("earth", jd0 + days)
    rm_, _ = heliocentric("mars", jd0 + days)
    eps = np.radians(23.4393)
    ecl = lambda v: np.stack([v[:, 0], np.cos(eps)*v[:, 1] + np.sin(eps)*v[:, 2]], -1)/AU
    E, M = ecl(re_), ecl(rm_)
    dist = np.linalg.norm(rm_ - re_, axis=1)
    lt = dist/C_KMS/60
    ts = np.concatenate([np.zeros(6, int), np.arange(len(days)), np.full(20, len(days) - 1)])
    fig = plt.figure(figsize=(10.6, 5.6))
    gs = fig.add_gridspec(1, 2, width_ratios=[1, 1.2], wspace=0.2)
    ax = fig.add_subplot(gs[0]); ax.set_aspect("equal"); ax.set_xlim(-1.9, 1.9); ax.set_ylim(-2.3, 1.9); ax.axis("off")
    tt = np.linspace(0, 2*np.pi, 300)
    ax.plot(np.cos(tt)*1.0, np.sin(tt)*1.0, color="#9db4d0", lw=0.8); ax.plot(1.52*np.cos(tt), 1.52*np.sin(tt), color="#cfa9a9", lw=0.8)
    ax.plot(0, 0, "o", color="#f5b301", ms=14, mec=AMBER)
    ep, = ax.plot([], [], "o", color=BLUE, ms=10, mec=INK); mp_, = ax.plot([], [], "o", color=MAROON, ms=8, mec=INK)
    et, = ax.plot([], [], color=BLUE, lw=1.4); mt, = ax.plot([], [], color=MAROON, lw=1.4)
    link, = ax.plot([], [], color=GREEN, lw=2.2, ls="--")
    pulse, = ax.plot([], [], "o", color=AMBER, ms=8, mec=INK, zorder=6)
    ax.text(-1.85, 1.85, "Earth", color=BLUE, fontsize=11, va="top"); ax.text(-1.85, 1.6, "Mars", color=MAROON, fontsize=11, va="top"); ax.text(-1.85, 1.35, "radio pulse", color=AMBER, fontsize=10, va="top")
    date = ax.text(1.85, 1.85, "", ha="right", va="top", fontsize=11, family="monospace")
    info = ax.text(0, -1.95, "", ha="center", va="bottom", fontsize=11.5, color=GREEN)
    g = fig.add_subplot(gs[1])
    yr = 2026 + days/365.25
    g.plot(yr, lt, color=INK, lw=2.2); g.set_xlim(yr[0], yr[-1]); g.set_ylim(0, 25)
    g.set_xlabel("year"); g.set_ylabel("one-way light time (min)"); despine(g)
    cur, = g.plot([], [], "o", color=GREEN, ms=9, zorder=5); vline = g.axvline(yr[0], color=GREEN, lw=1.0, ls=":")
    g.set_title("Earth to Mars, one synodic period (about 26 months)", fontsize=11)

    def frame(k):
        i = ts[k]
        s0 = max(0, i - 40)
        et.set_data(E[s0:i + 1, 0], E[s0:i + 1, 1]); mt.set_data(M[s0:i + 1, 0], M[s0:i + 1, 1])
        ep.set_data([E[i, 0]], [E[i, 1]]); mp_.set_data([M[i, 0]], [M[i, 1]])
        link.set_data([E[i, 0], M[i, 0]], [E[i, 1], M[i, 1]])
        f = (k*0.09) % 1.0                                         # a pulse racing along the link (schematic speed)
        pulse.set_data([E[i, 0] + f*(M[i, 0] - E[i, 0])], [E[i, 1] + f*(M[i, 1] - E[i, 1])])
        y, m, d = calendar_from_jd(jd0 + days[i]); date.set_text("%d-%02d-%02d" % (y, m, int(d)))
        info.set_text("distance %.2f AU     light time %.1f min\nround trip %.1f min" % (dist[i]/AU, lt[i], 2*lt[i]))
        cur.set_data([yr[i]], [lt[i]]); vline.set_xdata([yr[i], yr[i]])
        return []

    save_gif(fig, frame, len(ts), fps, path)


if which in ("sidereal", "all"):
    sidereal_animation()
if which in ("mars", "all"):
    mars_animation()