"""Time note, part 2: the family of time scales, the Earth as a (poor) clock, and relativity.

Uses the IERS Earth-orientation series bundled with astropy (astropy-iers-data, 1962-2024).

Run from the deck directory:   python codes/fig-time-scales.py
Writes img/fig-time-scales.png, img/fig-time-ut1.png, img/fig-time-tdb.png.
"""
import warnings
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _timeutil import (MU_E, RE, MU_SUN, AU, LEAP, TT_MINUS_TAI, tai_minus_utc, jd_gregorian, calendar_from_jd,
                       tdb_minus_tt, kepler_E, OMEGA_E, DAY)

warnings.filterwarnings("ignore")
GREY = "#8b8378"
C = 299792458.0

# ================================================================== 1. the ruler of time scales
now_dat = 37.0
rows = [   # (label, offset from UTC in s, colour, family, definition)
    ("UT1", 0.0, GREEN, "rotation", "Earth's rotation angle;  |UT1 − UTC| < 0.9 s  (kept by leap seconds)"),
    ("UTC", 0.0, BLUE, "atomic", "civil time: SI seconds, with leap seconds"),
    ("GPS time", now_dat - 19.0, BLUE, "atomic", "TAI − 19 s;  never has leap seconds (zero at 6 Jan 1980)"),
    ("TAI", now_dat, BLUE, "atomic", "pure atomic time;  TAI − UTC = 37 s since 1 Jan 2017"),
    ("TT", now_dat + TT_MINUS_TAI, MAROON, "dynamical", "TAI + 32.184 s;  time on the geoid, used for Earth-centred ephemerides"),
    ("TDB", now_dat + TT_MINUS_TAI, PURPLE, "dynamical", "TT + periodic terms (±1.66 ms);  barycentric time, used for planetary ephemerides"),
]
fig, ax = plt.subplots(figsize=(13.4, 4.9))
for j, (lab, off, col, fam, txt) in enumerate(rows[::-1]):
    y = j
    ax.plot([0, 80], [y, y], color=col, lw=0.6, alpha=0.35)
    ax.plot([off], [y], "o", color=col, ms=13, zorder=4, mec=INK, mew=0.8)
    ax.text(-2.2, y, lab, ha="right", va="center", fontsize=12.5, color=col, weight="bold")
    ax.text(82.2, y, txt, ha="left", va="center", fontsize=9.6)
    if lab == "UT1":
        ax.add_patch(FancyBboxPatch((-0.9, y - 0.16), 1.8, 0.32, boxstyle="round,pad=0,rounding_size=0.1", color=GREEN, alpha=0.3, zorder=3))
    if lab == "TDB":
        ax.add_patch(FancyBboxPatch((off - 0.9, y - 0.16), 1.8, 0.32, boxstyle="round,pad=0,rounding_size=0.1", color=PURPLE, alpha=0.3, zorder=3))
ys = {r[0]: len(rows) - 1 - k for k, r in enumerate(rows)}
def brace(a, b, txt, y0, dy=-0.62):
    ax.annotate("", xy=(rows[b][1], y0), xytext=(rows[a][1], y0), arrowprops=dict(arrowstyle="<->", color=INK, lw=1.2))
    ax.text(0.5*(rows[a][1] + rows[b][1]), y0 + dy*0.5, txt, ha="center", va="top", fontsize=9.8)
brace(1, 2, "+18 s", ys["GPS time"] - 0.45)
brace(2, 3, "+19 s", ys["TAI"] - 0.45)
brace(3, 4, "+32.184 s (exact)", ys["TT"] - 0.45)
ax.text(0.5*(0 + 37 + 32.184 + 37)/1, ys["UTC"] + 0.62, "", fontsize=9)
ax.annotate("", xy=(now_dat + TT_MINUS_TAI, -0.95), xytext=(0, -0.95), arrowprops=dict(arrowstyle="<->", color=MAROON, lw=1.5))
ax.text(0.5*(now_dat + TT_MINUS_TAI), -1.05, "TT − UTC = 69.184 s  (an ephemeris looked up with the wrong scale is off by this much)", ha="center", va="top", fontsize=10.5, color=MAROON)
ax.set_xlim(-14, 80); ax.set_ylim(-1.9, len(rows) - 0.4)
ax.set_xlabel("offset from UTC at one instant in 2026  (seconds)")
ax.set_yticks([]); ax.spines[["left", "top", "right"]].set_visible(False)
ax.set_title("The same instant on six clocks", fontsize=12.5)
fig.savefig("img/fig-time-scales.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ================================================================== 2. the Earth as a clock (real IERS data)
from astropy.utils import iers
iers.conf.auto_download = False
iers.conf.iers_degraded_accuracy = "ignore"
b = iers.IERS_B.open()
mjd = np.asarray(b["MJD"].value, float)
ut1utc = np.asarray(b["UT1_UTC"].value, float)
sel = mjd >= 41317                                   # 1972-01-01: UTC in its present form
mjd, ut1utc = mjd[sel], ut1utc[sel]
jd = mjd + 2400000.5
year = 2000.0 + (jd - 2451545.0)/365.25
ut1tai = ut1utc - tai_minus_utc(jd)                  # continuous (no leap-second steps)
lod = -np.gradient(ut1tai, mjd)*1e3                  # ms: LOD excess over 86 400 s
k = 365
lod_s = np.convolve(lod, np.ones(k)/k, mode="same")
print("IERS-B: %s to %s, %d days" % (calendar_from_jd(jd[0]), calendar_from_jd(jd[-1]), len(jd)))
print("UT1-UTC range %.3f .. %.3f s" % (ut1utc.min(), ut1utc.max()))
print("LOD excess: mean %.3f ms; min %.3f ms on %s; max %.3f ms on %s" %
      (lod.mean(), lod.min(), calendar_from_jd(jd[np.argmin(lod)]), lod.max(), calendar_from_jd(jd[np.argmax(lod)])))
print("1-year mean LOD: min %.3f (%.1f) max %.3f (%.1f)" % (lod_s[k:-k].min(), year[k:-k][np.argmin(lod_s[k:-k])], lod_s[k:-k].max(), year[k:-k][np.argmax(lod_s[k:-k])]))
print("equatorial displacement for 0.9 s: %.0f m" % (0.9*OMEGA_E*RE*1e3))
# secular trend of the day length since 1972 (tidal braking + everything else)
p = np.polyfit(year[k:-k], lod_s[k:-k], 1)
print("linear trend of 1-yr mean LOD 1973-2023: %.3f ms/century" % (p[0]*100))
n_leap = len(LEAP) - 1
print("leap seconds inserted 1972-2017: %d;  TAI-UTC = %d s" % (n_leap, LEAP[-1][1]))

fig, (u1, u2, u3) = plt.subplots(1, 3, figsize=(15.4, 4.3))
yl = np.array([y + (m - 1)/12 + (d - 1)/365 for (y, m, d), _ in LEAP])
stairs_x = np.concatenate([[1972.0], np.repeat(yl[1:], 2), [2026.9]])
stairs_y = np.concatenate([np.repeat([v for _, v in LEAP], 2)])
u1.step(np.concatenate([yl, [2026.9]]), [v for _, v in LEAP] + [37], where="post", color=BLUE, lw=2.4)
u1.set_xlabel("year"); u1.set_ylabel("TAI − UTC  (s)"); despine(u1)
u1.set_title("27 leap seconds, the last on 31 Dec 2016", fontsize=11.5)
u1.annotate("10 s at the start\n(1 Jan 1972)", (1972.0, 10), xytext=(1974, 26), fontsize=9.5, arrowprops=dict(arrowstyle="-", lw=0.6))
u1.annotate("37 s", (2017.0, 37), xytext=(2005, 38.5), fontsize=9.5, arrowprops=dict(arrowstyle="-", lw=0.6))
u1.set_ylim(8, 41)

u2.plot(year, ut1utc, color=GREEN, lw=1.3)
for (y_, m_, d_), v in LEAP[1:]:
    u2.axvline(y_ + (m_ - 1)/12, color=BLUE, lw=0.4, alpha=0.6)
u2.axhline(0, color=INK, lw=0.5); u2.axhline(0.9, color=MAROON, ls="--", lw=0.9); u2.axhline(-0.9, color=MAROON, ls="--", lw=0.9)
u2.text(1973, 0.93, "±0.9 s: the limit that triggers a leap second", fontsize=8.8, color=MAROON, va="bottom")
u2.set_ylim(-1.05, 1.15); u2.set_xlabel("year"); u2.set_ylabel("UT1 − UTC  (s)"); despine(u2)
u2.set_title("The Earth drifts from atomic time; leap seconds (blue) reset it", fontsize=11.5)
sec = u2.secondary_yaxis("right", functions=(lambda s: s*OMEGA_E*RE*1e3, lambda m: m/(OMEGA_E*RE*1e3)))
sec.set_ylabel("displacement of an equatorial site (m)")

u3.plot(year, lod, color=GREY, lw=0.4, alpha=0.7, label="daily")
u3.plot(year[k:-k], lod_s[k:-k], color=INK, lw=2.4, label="1-year mean")
u3.axhline(0, color=MAROON, lw=0.9, ls="--"); u3.text(1973, 0.05, "86 400 s exactly", color=MAROON, fontsize=9, va="bottom")
u3.set_xlabel("year"); u3.set_ylabel("length of day − 86 400 s  (ms)"); despine(u3)
u3.legend(frameon=False, fontsize=9, loc="upper right")
u3.set_title("The day length wanders by about 1–2 ms", fontsize=11.5)
fig.tight_layout(); fig.savefig("img/fig-time-ut1.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ================================================================== 3. relativity: TDB-TT, and GPS clocks
jd0 = jd_gregorian(2026, 1, 1)
days = np.arange(0, 3*365.25, 1.0)
series = tdb_minus_tt(jd0 + days + 69.184/86400)*1e3
# closed form: (2 sqrt(mu a)/c^2) e sin E  for a Keplerian Earth
e_E, a_E = 0.016709, AU*1e3
amp = 2*np.sqrt(MU_SUN*1e9*a_E)/C**2*e_E
nn = jd0 + days - 2451545.0
Mn = np.radians(357.528 + 0.9856003*nn)
closed = amp*np.sin(kepler_E(Mn, e_E))*1e3
print("closed-form amplitude 2 sqrt(mu a) e / c^2 = %.4f ms; series peak %.4f ms; max |series - closed| = %.1f us" %
      (amp*1e3, np.abs(series).max(), 1e3*np.abs(series - closed).max()))
print("Earth's speed x 1.657 ms = %.1f m ; Mars' speed x 1.657 ms = %.1f m" % (29.78e3*1.657e-3, 24.13e3*1.657e-3))

# GPS: rate of a satellite clock relative to a clock on the geoid
a_gps = 26560e3
GM = MU_E*1e9
Re_ = RE*1e3
grav = GM/C**2*(1/Re_ - 1/a_gps)                    # satellite runs faster (higher potential)
vel = -0.5*(GM/a_gps)/C**2                          # ... and slower (moving)
# the geoid clock also moves with the Earth: its own -0.5 (omega R)^2/c^2 is absorbed in the geoid potential (U0)
net = grav + vel
print("GPS: gravitational +%.2f us/day, velocity %.2f us/day, net +%.2f us/day -> %.1f km/day of ranging error" %
      (grav*DAY*1e6, vel*DAY*1e6, net*DAY*1e6, net*DAY*C/1e3))
print("     fractional rate offset %.3e (compare a 1e-13 Cs clock)" % net)

fig, (r1, r2) = plt.subplots(1, 2, figsize=(13.2, 4.3), gridspec_kw=dict(width_ratios=[1.6, 1]))
r1.plot(days/365.25 + 2026, series, color=PURPLE, lw=2.4, label="series (Astronomical Almanac)")
r1.plot(days[::14]/365.25 + 2026, closed[::14], "o", color=AMBER, ms=5.5, label=r"$\dfrac{2\sqrt{\mu a}}{c^2}\,e\sin E$  (Kepler orbit)")
r1.axhline(0, color=INK, lw=0.5)
r1.set_xlabel("year"); r1.set_ylabel("TDB − TT  (ms)"); despine(r1)
r1.set_title("Barycentric time runs ±1.66 ms around the Earth's clock", fontsize=11.5)
r1.legend(frameon=False, fontsize=9.5, loc="upper right", ncol=1)
r1.set_ylim(-2.4, 2.9)
labels = ["gravity\n(higher up)", "velocity\n(moving)", "net"]
vals = np.array([grav, vel, net])*DAY*1e6
r2.bar(range(3), vals, color=[GREEN, MAROON, INK], width=0.6)
for j, v in enumerate(vals):
    r2.text(j, v + (2 if v > 0 else -2), "%+.1f µs" % v, ha="center", va="bottom" if v > 0 else "top", fontsize=10.5)
r2.set_xticks(range(3)); r2.set_xticklabels(labels); r2.axhline(0, color=INK, lw=0.6)
r2.set_ylim(-16, 58); r2.set_ylabel("clock rate vs. the ground (µs per day)")
r2.set_title("GPS satellite clocks: +38 µs/day,\nabout 11 km/day if ignored", fontsize=11.5); despine(r2)
fig.tight_layout(); fig.savefig("img/fig-time-tdb.png", dpi=160, bbox_inches="tight"); plt.close(fig)