"""Time note, part 3: what a clock error does to an interplanetary mission (numerical experiments).

  1. Light time Earth-Mars, 2026-2031 (real geometry from the built-in ephemeris).
  2. Oscillator classes: time error growing as eps*t, and the range error eps*rho.
  3. A representative 2026 Earth-Mars transfer (Lambert), and how a velocity error at departure
     turns into a miss distance at Mars -- with the clock stability that would cause it.
  4. The "ledger": consequences of using the wrong time scale.

Run from the deck directory:   python codes/fig-time-nav.py
Writes img/fig-time-lighttime.png, img/fig-time-clocks.png, img/fig-time-nav.png and prints the numbers.
"""
import warnings
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _timeutil import (C_KMS, AU, MU_SUN, DAY, OMEGA_E, RE, jd_gregorian, calendar_from_jd, heliocentric, lambert,
                       kepler_propagate, tai_minus_utc, TT_MINUS_TAI)

warnings.filterwarnings("ignore")
GREY = "#8b8378"
EPS = np.radians(23.4393)


def ecl(v):
    """ICRS/equatorial -> ecliptic axes (rotation about x)."""
    v = np.asarray(v)
    c, s = np.cos(EPS), np.sin(EPS)
    return np.stack([v[..., 0], c*v[..., 1] + s*v[..., 2], -s*v[..., 1] + c*v[..., 2]], -1)


def date_str(jd):
    y, m, d = calendar_from_jd(jd)
    return "%d-%02d-%02d" % (y, m, int(d))


# ================================================================== 1. light time
jd0 = jd_gregorian(2026, 1, 1)
days = np.arange(0, 365.25*5.0, 1.0)
re_, _ = heliocentric("earth", jd0 + days)
rm_, _ = heliocentric("mars", jd0 + days)
dist = np.linalg.norm(rm_ - re_, axis=1)
lt = dist/C_KMS/60
# extrema
def local_ext(y, kind):
    i = np.arange(1, len(y) - 1)
    m = (y[i] < y[i-1]) & (y[i] <= y[i+1]) if kind == "min" else (y[i] > y[i-1]) & (y[i] >= y[i+1])
    return i[m]
imin, imax = local_ext(dist, "min"), local_ext(dist, "max")
print("Earth-Mars: closest approaches:", [(date_str(jd0 + days[i]), round(dist[i]/AU, 3), round(lt[i], 2)) for i in imin])
print("Earth-Mars: greatest distances :", [(date_str(jd0 + days[i]), round(dist[i]/AU, 3), round(lt[i], 2)) for i in imax])
print("light time range 2026-2031: %.2f .. %.2f min" % (lt.min(), lt.max()))
# 2014 Mars Orbiter Mission insertion, 24 Sep 2014 ~ 02:00 UTC
jdm = jd_gregorian(2014, 9, 24) + 2/24 + (tai_minus_utc(jd_gregorian(2014, 9, 24)) + TT_MINUS_TAI)/DAY
rE, _ = heliocentric("earth", jdm); rM, _ = heliocentric("mars", jdm)
dmom = np.linalg.norm(rM - rE)
print("MOM orbit insertion 2014-09-24: distance %.3f AU = %.0f million km, one-way light time %.2f min, round trip %.2f min" %
      (dmom/AU, dmom/1e6, dmom/C_KMS/60, 2*dmom/C_KMS/60))
# extremes over many synodic periods (2000-2050), the "3 to 22 minutes" statement
dd = np.arange(0, 365.25*50, 3.0)
r1, _ = heliocentric("earth", jd_gregorian(2000, 1, 1) + dd); r2, _ = heliocentric("mars", jd_gregorian(2000, 1, 1) + dd)
dl = np.linalg.norm(r2 - r1, axis=1)/C_KMS/60
print("2000-2050 light time range: %.2f .. %.2f min  (%.3f .. %.3f AU)" % (dl.min(), dl.max(), dl.min()*60*C_KMS/AU, dl.max()*60*C_KMS/AU))

fig, (l1, l2) = plt.subplots(1, 2, figsize=(14.4, 4.7), gridspec_kw=dict(width_ratios=[1.7, 1]))
x = 2026 + days/365.25
l1.plot(x, lt, color=INK, lw=2.4)
l1.set_ylabel("one-way light time, Earth to Mars (min)"); l1.set_xlabel("year")
ax2 = l1.secondary_yaxis("right", functions=(lambda m: m*60*C_KMS/AU, lambda a: a/(60*C_KMS/AU)))
ax2.set_ylabel("distance (AU)")
for i in imin:
    l1.plot(x[i], lt[i], "o", color=GREEN, ms=8); l1.text(x[i], lt[i] - 1.7, "opposition\n%s\n%.1f min" % (date_str(jd0 + days[i]), lt[i]), ha="center", va="top", fontsize=8.6, color=GREEN)
for i in imax:
    l1.plot(x[i], lt[i], "o", color=MAROON, ms=8); l1.text(x[i], lt[i] + 0.7, "conjunction\n%s\n%.1f min" % (date_str(jd0 + days[i]), lt[i]), ha="center", va="bottom", fontsize=8.6, color=MAROON)
l1.set_ylim(-1, 26); despine(l1)
l1.set_title("A signal to Mars takes 5 to 21 minutes (3 to 22 over the decades)", fontsize=11.5)

# top view at opposition and conjunction
def orbit_xy(body, jd_start, n_days):
    dj = np.arange(0, n_days, 3.0)
    r, _ = heliocentric(body, jd_start + dj)
    return ecl(r)[:, :2]/AU
eo, mo = orbit_xy("earth", jd0, 366), orbit_xy("mars", jd0, 688)
l2.plot(eo[:, 0], eo[:, 1], color=BLUE, lw=1.0); l2.plot(mo[:, 0], mo[:, 1], color=MAROON, lw=1.0)
l2.plot(0, 0, "o", color="#f5b301", ms=13, mec=AMBER)
ii = [imin[1] if len(imin) > 1 else imin[0], imax[0] if len(imax) else 0]
for i, tag in ((imin[0], "opposition"), (imax[0], "conjunction")):
    pe, pm = ecl(re_[i])[:2]/AU, ecl(rm_[i])[:2]/AU
    l2.plot(*pe, "o", color=BLUE, ms=8); l2.plot(*pm, "o", color=MAROON, ms=8)
    l2.plot([pe[0], pm[0]], [pe[1], pm[1]], color=GREEN if tag == "opposition" else MAROON, lw=1.8, ls="--")
    l2.text(*(0.5*(pe + pm) + (np.array([0.05, 0.12]) if tag == "opposition" else np.array([0.12, -0.3]))), "%.1f min" % lt[i], fontsize=9.5, color=GREEN if tag == "opposition" else MAROON)
l2.set_aspect("equal"); l2.set_xlim(-2, 2); l2.set_ylim(-2, 2)
l2.set_xlabel("x (AU)"); l2.set_ylabel("y (AU)"); l2.set_title("Top view: the two extremes", fontsize=11.5); despine(l2)
fig.tight_layout(); fig.savefig("img/fig-time-lighttime.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ================================================================== 2. clocks
clocks = [   # name, fractional frequency error eps (order of magnitude), colour
    ("quartz crystal", 1e-6, GREY), ("oven-controlled quartz", 1e-9, BLUE), ("rubidium", 1e-11, GREEN),
    ("caesium beam", 1e-13, AMBER), ("hydrogen maser", 1e-15, MAROON), ("optical lattice (lab)", 1e-17, PURPLE)]
print("\nCLOCKS   eps   time error after 1 s / 1 h / 1 day / 1 year  |  c*dt after 1 day  |  range error at 1.5 AU  |  at 2.5 AU")
for name, e, _ in clocks:
    print("%-24s %.0e  %.2e s / %.2e s / %.2e s / %.2e s | %s | %.3g m | %.3g m" % (
        name, e, e, e*3600, e*DAY, e*365.25*DAY,
        "%.3g m" % (C_KMS*1e3*e*DAY), e*1.5*AU*1e3, e*2.5*AU*1e3))

fig, (c1, c2) = plt.subplots(1, 2, figsize=(14.0, 5.0))
t = np.logspace(0, np.log10(365.25*DAY), 200)
for name, e, col in clocks:
    c1.loglog(t, e*t, color=col, lw=2.4)
    c1.text(t[-1]*1.15, e*t[-1], "%s (%.0e)" % (name, e), color=col, fontsize=9.2, va="center", clip_on=False)
for yy, s_ in ((1/C_KMS/1e3, "1 m"), (1e3/(C_KMS*1e3), "1 km"), (1e6/(C_KMS*1e3), "1000 km")):
    c1.axhline(yy, color=INK, lw=0.7, ls=":"); c1.text(1.15, yy*1.3, "light travels " + s_ if s_ != "1 m" else "light travels 1 m (3.3 ns)", fontsize=8.8, bbox=dict(fc=CREAM, ec="none", pad=0.5))
for xx, s_ in ((60, "1 min"), (3600, "1 h"), (DAY, "1 day"), (365.25*DAY, "1 year")):
    c1.axvline(xx, color=INK, lw=0.5, ls=":", alpha=0.6); c1.text(xx/1.12, 60, s_, rotation=90, fontsize=8.8, va="top", ha="right")
c1.set_xlim(1, 365.25*DAY); c1.set_ylim(1e-17, 1e2)
c1.set_xlabel("time since the clock was set right (s)"); c1.set_ylabel("accumulated clock error (s)")
c1.set_title("A clock that is slightly fast or slow: error = $\\varepsilon\\,t$", fontsize=11.5); despine(c1)

rho = np.linspace(0.5, 2.7, 100)*AU*1e3   # m
for name, e, col in clocks:
    c2.loglog(rho/(AU*1e3), e*rho, color=col, lw=2.4)
c2.axhline(1, color=INK, lw=0.7, ls=":"); c2.text(0.52, 1.3, "1 m", fontsize=9)
c2.axhline(3.39e6, color=INK, lw=0.7, ls=":"); c2.text(0.52, 3.39e6*1.3, "Mars radius", fontsize=9)
for name, e, col in clocks:
    c2.text(2.75, e*2.7*AU*1e3, name, color=col, fontsize=9.2, va="center", clip_on=False)
c2.set_xlim(0.5, 2.7); c2.set_ylim(1e-6, 1e10)
c2.minorticks_off(); c2.set_xticks([0.5, 1, 1.5, 2, 2.5]); c2.set_xticklabels(["0.5", "1", "1.5", "2", "2.5"]); c2.set_yticks(10.0**np.arange(-6, 11, 2))
c2.set_xlabel("Earth–Mars distance $\\rho$ (AU)"); c2.set_ylabel("two-way range error  $\\delta\\rho=\\varepsilon\\,\\rho$  (m)")
c2.set_title("Ranging with a clock of fractional error $\\varepsilon$", fontsize=11.5); despine(c2)
fig.tight_layout(); fig.savefig("img/fig-time-clocks.png", dpi=160, bbox_inches="tight"); plt.close(fig)
print("range error needed 1 m at 1.5 AU -> eps = %.2e ; 10 m at 2.5 AU -> eps = %.2e" % (1/(1.5*AU*1e3), 10/(2.5*AU*1e3)))
print("timing needed for 1 m two-way ranging: %.2f ns" % (2*1/(C_KMS*1e3)*1e9))

# ================================================================== 3. Lambert transfer and miss distance
jd_dep, tof = jd_gregorian(2026, 11, 2), 290.0
jd_arr = jd_dep + tof
rE0, vE0 = heliocentric("earth", jd_dep)
rM1, vM1 = heliocentric("mars", jd_arr)
v1, v2 = lambert(rE0, rM1, tof*DAY)
print("\nTRANSFER: depart %s, arrive %s (%.0f days)" % (date_str(jd_dep), date_str(jd_arr), tof))
print("  v_inf departure %.3f km/s (C3 = %.2f km^2/s^2), v_inf arrival %.3f km/s" % (np.linalg.norm(v1 - vE0), np.linalg.norm(v1 - vE0)**2, np.linalg.norm(v2 - vM1)))
print("  distance Earth-Mars at arrival %.3f AU; sun distance of transfer: r_dep %.3f AU, r_arr %.3f AU" % (
    np.linalg.norm(rM1 - heliocentric('earth', jd_arr)[0])/AU, np.linalg.norm(rE0)/AU, np.linalg.norm(rM1)/AU))
Rh = rE0/np.linalg.norm(rE0); Nh = np.cross(rE0, v1); Nh /= np.linalg.norm(Nh); Th = np.cross(Nh, Rh)
dirs = {"radial": Rh, "along-track": Th, "normal to the orbit": Nh}
dvs = np.logspace(-12, 0, 49)                     # km/s   (1 mm/s = 1e-6 km/s)
miss = {}
for name, u in dirs.items():
    miss[name] = np.array([np.linalg.norm(kepler_propagate(rE0, v1 + d*u, tof*DAY)[0] - rM1) for d in dvs])
print("  miss distance at Mars for a 1 mm/s error at departure: " +
      ", ".join("%s %.1f km" % (n, np.interp(1e-6, dvs, m)) for n, m in miss.items()))
print("  ... for 1 cm/s: " + ", ".join("%s %.0f km" % (n, np.interp(1e-5, dvs, m)) for n, m in miss.items()))
print("  ... for 1 m/s: " + ", ".join("%s %.0f km" % (n, np.interp(1e-3, dvs, m)) for n, m in miss.items()))
Phi = np.zeros((3, 3)); h = 1e-7
for k in range(3):
    dv = np.zeros(3); dv[k] = h
    Phi[:, k] = (kepler_propagate(rE0, v1 + dv, tof*DAY)[0] - kepler_propagate(rE0, v1 - dv, tof*DAY)[0])/(2*h)
sv = np.linalg.svd(Phi, compute_uv=False)
print("  singular values of dr(arrival)/dv(departure): %s km per (mm/s)" % np.round(sv*1e-6, 2))
print("  velocity accuracy needed for a 10 km miss (worst direction): %.3f mm/s" % (10/(sv[0]*1e-6)))
print("  ... equivalent fractional frequency error c*eps = dv:  eps = %.2e" % ((10/(sv[0]*1e-6))*1e-6/C_KMS))
# equivalent for clocks: one-way delta v = c eps
for name, e, _ in clocks:
    dvk = C_KMS*e
    print("   clock %-24s eps %.0e -> dv = %.3g mm/s -> worst miss ~ %.3g km" % (name, e, dvk*1e6, sv[0]*dvk))

fig, (n1, n2) = plt.subplots(1, 2, figsize=(14.2, 5.6), gridspec_kw=dict(width_ratios=[1, 1.25]))
# (a) trajectory
tt = np.linspace(0, tof*DAY, 300)
path = np.array([kepler_propagate(rE0, v1, s)[0] for s in tt])
eo = orbit_xy("earth", jd0, 366); mo = orbit_xy("mars", jd0, 688)
n1.plot(eo[:, 0], eo[:, 1], color=BLUE, lw=0.9, alpha=0.7); n1.plot(mo[:, 0], mo[:, 1], color=MAROON, lw=0.9, alpha=0.7)
P = ecl(path)[:, :2]/AU
n1.plot(P[:, 0], P[:, 1], color=INK, lw=2.6)
n1.plot(0, 0, "o", color="#f5b301", ms=13, mec=AMBER)
pe = ecl(rE0)[:2]/AU; pm = ecl(rM1)[:2]/AU
n1.plot(*pe, "o", color=BLUE, ms=10, mec=INK); n1.text(pe[0] + 0.05, pe[1] - 0.16, "Earth\n%s" % date_str(jd_dep), color=BLUE, fontsize=9.5)
n1.plot(*pm, "o", color=MAROON, ms=10, mec=INK); n1.text(pm[0] + 0.1, pm[1] - 0.05, "Mars\n%s" % date_str(jd_arr), color=MAROON, fontsize=9.5, ha="left", va="top")
n1.set_aspect("equal"); n1.set_xlim(-1.9, 1.9); n1.set_ylim(-1.9, 1.9)
n1.set_xlabel("x (AU)"); n1.set_ylabel("y (AU)"); despine(n1)
n1.set_title("A representative 2026 transfer (Lambert), %.0f days" % tof, fontsize=11.5)

# (b) miss vs velocity error
cs = {"radial": GREEN, "along-track": MAROON, "normal to the orbit": BLUE}
for name in dirs:
    n2.loglog(dvs*1e6, miss[name], color=cs[name], lw=2.6, label=name)
n2.axhline(10, color=INK, lw=0.7, ls=":"); n2.text(1.3e-6, 12, "10 km: typical Mars approach target", fontsize=9)
n2.axhline(3389.5, color=INK, lw=0.7, ls=":"); n2.text(1.3e-6, 3389.5*1.25, "Mars radius: the spacecraft misses the planet", fontsize=9)
n2.axhline(1, color=INK, lw=0.7, ls=":"); n2.text(1.3e-6, 1.25, "1 km", fontsize=9)
n2.set_xlim(1e-6, 1e6); n2.set_ylim(1e-4, 3e8)
n2.set_xlabel("velocity error at departure  $\\delta v$  (mm/s)"); n2.set_ylabel("miss distance at Mars, 290 days later (km)")
n2.legend(frameon=False, fontsize=10, loc="lower right", bbox_to_anchor=(1.0, 0.14), title="error direction", title_fontsize=9.5)
despine(n2); n2.set_title("A tiny velocity error, left uncorrected, becomes a large miss", fontsize=11.5)
top = n2.secondary_xaxis("top", functions=(lambda v: v*1e-6/C_KMS, lambda e: e*C_KMS*1e6))
top.set_xlabel("equivalent fractional frequency error  $\\varepsilon=\\delta v/c$  (one-way Doppler)", fontsize=9.8)
top.set_xscale("log")
top.minorticks_off(); top.set_xticks([1e-17, 1e-15, 1e-13, 1e-11, 1e-9, 1e-7]); top.set_xticklabels(["$10^{-17}$", "$10^{-15}$", "$10^{-13}$", "$10^{-11}$", "$10^{-9}$", "$10^{-7}$"])
for name, e, col in clocks[1:]:
    n2.axvline(C_KMS*e*1e6, color=col, lw=1.0, ls="--", alpha=0.8)
    n2.text(C_KMS*e*1e6/1.15, 1.5e8, name, color=col, rotation=90, fontsize=8.6, va="top", ha="right")
fig.tight_layout(); fig.savefig("img/fig-time-nav.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ================================================================== 4. ledger of time-scale mistakes
print("\nLEDGER")
jd_t = jd_arr
for body in ("earth", "mars"):
    r0, v0 = heliocentric(body, jd_t); r1, _ = heliocentric(body, jd_t + 69.184/DAY)
    print("  %s: |v| = %.2f km/s; ephemeris shifted by 69.184 s moves it %.0f km; by 1 s: %.1f km; by 37 s: %.0f km" %
          (body, np.linalg.norm(v0), np.linalg.norm(r1 - r0), np.linalg.norm(v0), 37*np.linalg.norm(v0)))
vLEO = np.sqrt(398600.4418/(RE + 500)); vGEO = np.sqrt(398600.4418/42164.17)
print("  500 km LEO |v| = %.3f km/s: 18 s (GPS vs UTC) -> %.0f km along track; 1 s -> %.2f km" % (vLEO, 18*vLEO, vLEO))
print("  equator surface speed %.3f km/s: UT1-UTC = 0.9 s -> %.0f m; 0.3 s -> %.0f m" % (OMEGA_E*RE, 0.9*OMEGA_E*RE*1e3, 0.3*OMEGA_E*RE*1e3))
print("  TDB-TT = 1.657 ms: Earth %.0f m, Mars %.0f m" % (29.78e3*1.657e-3, 24.13e3*1.657e-3))
print("  delta-DOR: baseline 8000 km; timing 30 ps -> angle %.2f nrad -> %.0f m at 1.5 AU" % (
    30e-12*C_KMS/8000*1e9, 30e-12*C_KMS/8000*1.5*AU*1e3))