"""Ground-track note: the standard families of ground tracks.

    img/fig-gt-leo.png        three low orbits, one day each (6, 28.5, 51.6 deg)
    img/fig-gt-polar.png      a polar orbit: one day, coloured by orbit number
    img/fig-gt-sso.png        a sun-synchronous orbit: track, and local time of the node over a year
    img/fig-gt-gso.png        geostationary and inclined geosynchronous orbits (NavIC-like), and the small-i figure 8
    img/fig-gt-molniya.png    Molniya and Tundra: highly elliptic tracks with hourly ticks

Run from the deck directory:   python codes/fig-gt-tracks.py
"""
import numpy as np
import matplotlib.pyplot as plt
from _gt import *
from _timeutil import T_STELLAR, DAY, OMEGA_E, MU_E, RE, D2R, j2_rates
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

W = 11.5


def new_map(figsize=(W, 6.3), **kw):
    fig, ax = plt.subplots(figsize=figsize)
    draw_map(ax, **kw)
    return fig, ax


def mark_ends(ax, lon, lat, color=INK, label=True):
    ax.plot(((lon[0] + 180) % 360) - 180, lat[0], "o", color=color, mec="white", ms=8, zorder=8)
    ax.plot(((lon[-1] + 180) % 360) - 180, lat[-1], "s", color=color, mec="white", ms=7, zorder=8)


# ------------------------------------------------------------------ LEO
t = np.arange(0, DAY + 1, 15.0)
fig, ax = new_map()
leos = [("Astrosat-like: $i=6^\\circ$, 650 km", RE + 650, 6.0, AMBER, 20.0),
        ("Cape Canaveral launch inclination: $i=28.5^\\circ$, 500 km", RE + 500, 28.5, GREEN, 60.0),
        ("International Space Station: $i=51.6^\\circ$, 420 km", RE + 420, 51.64, BLUE, 100.0)]
for lab, a, inc, col, om in leos:
    tr = track(a, 0.0, inc, om, 0.0, 0.0, t)
    plot_track(ax, tr["lon"], tr["lat"], color=col, lw=1.1, alpha=0.95)
    ax.axhline(inc, color=col, lw=0.8, ls="--", alpha=0.6); ax.axhline(-inc, color=col, lw=0.8, ls="--", alpha=0.6)
    ax.plot([], [], color=col, lw=2, label=lab)
    mark_ends(ax, tr["lon"], tr["lat"], color=col)
    T = period(a)
    print("LEO %-52s T = %.2f min, %.3f rev/day, node shifts %.2f deg west per orbit, %.0f km at the equator" %
          (lab.replace("$", "").replace("\\circ", "deg"), T/60, DAY/T, np.degrees(OMEGA_E*T), np.degrees(OMEGA_E*T)*RE*np.pi/180))
ax.legend(loc="lower left", fontsize=9.5, frameon=True, facecolor="white", framealpha=0.9)
ax.set_title("Ground tracks of three low orbits, 24 hours each (circle: start, square: end)", fontsize=12)
fig.tight_layout(); fig.savefig("img/fig-gt-leo.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ polar
a = RE + 800
tr = track(a, 0.0, 90.0, 0.0, 0.0, 0.0, t)
norb = tr["t"]/period(a)
fig, ax = new_map()
lc = plot_track(ax, tr["lon"], tr["lat"], c=norb, cmap="viridis", lw=1.3, vmin=0, vmax=norb.max())
cb = fig.colorbar(lc, ax=ax, fraction=0.022, pad=0.01); cb.set_label("orbit number", fontsize=10)
mark_ends(ax, tr["lon"], tr["lat"])
rho = earth_central_angle(800, 10)
ax.set_title("Polar orbit, $i=90^\\circ$, 800 km: every latitude is reached, and the tracks cross at the poles", fontsize=12)
fig.tight_layout(); fig.savefig("img/fig-gt-polar.png", dpi=160, bbox_inches="tight"); plt.close(fig)
T = period(a)
print("POLAR 800 km: T = %.2f min, %.3f rev/day, node shift %.2f deg/orbit; Omega_dot(J2) = %.4f deg/day; coverage half-angle at 10 deg elevation %.1f deg (swath %.0f km)" %
      (T/60, DAY/T, np.degrees(OMEGA_E*T), np.degrees(j2_rates(a, 0, np.pi/2)[0])*DAY, rho, 2*rho*RE*np.pi/180))

# ------------------------------------------------------------------ sun-synchronous
a = RE + 800
i_sso = sso_inclination(a)
print("SSO at 800 km: i = %.3f deg" % i_sso)
Omd, ud = nodal_rates(a, 0.0, i_sso)
tr = track(a, 0.0, i_sso, 0.0, 0.0, 0.0, t, j2=True)
norb = tr["t"]/period(a)
fig = plt.figure(figsize=(W + 4.3, 6.1))
gs = fig.add_gridspec(1, 2, width_ratios=[2.05, 1], wspace=0.22)
ax = fig.add_subplot(gs[0]); draw_map(ax)
lc = plot_track(ax, tr["lon"], tr["lat"], c=norb, cmap="plasma", lw=1.3, vmin=0, vmax=norb.max())
mark_ends(ax, tr["lon"], tr["lat"])
ax.axhline(180 - i_sso, color=INK, lw=0.7, ls="--"); ax.axhline(-(180 - i_sso), color=INK, lw=0.7, ls="--")
ax.text(-178, 180 - i_sso + 3, "max latitude $180^\\circ-i=%.1f^\\circ$" % (180 - i_sso), fontsize=9)
ax.set_title("Sun-synchronous orbit, $i=%.2f^\\circ$, 800 km, 24 h" % i_sso, fontsize=12)
# local time of the ascending node over a year
b = fig.add_subplot(gs[1])
days = np.arange(0, 366, 1.0)
lst_hours = lambda Om_dot, Om0: (12 + (Om0 + Om_dot*days*DAY - np.radians(0.9856473)*days)/np.radians(15)) % 24
for lab, inc, col in (("sun-synchronous, $i=%.1f^\\circ$" % i_sso, i_sso, MAROON), ("polar, $i=90^\\circ$", 90.0, BLUE), ("ISS-like, $i=51.6^\\circ$", 51.64, GREEN)):
    Od, _ = nodal_rates(a if inc != 51.64 else RE + 420, 0.0, inc)
    y = lst_hours(Od, np.radians(10.5*15 - 180))
    y_ = y.copy(); y_[np.r_[False, np.abs(np.diff(y)) > 12]] = np.nan
    b.plot(days, y_, color=col, lw=2.4 if col == MAROON else 1.8, label=lab)
    print("  local time of ascending node drifts %.2f h per 30 days for %s" % ((np.diff(np.unwrap(y*np.pi/12))*12/np.pi).mean()*30, lab.replace("$", "").replace("\\circ", "deg")))
b.set_xlabel("days since launch"); b.set_ylabel("local mean solar time of the ascending node (h)")
b.set_ylim(0, 24); b.set_yticks(range(0, 25, 3)); b.legend(fontsize=8.6, frameon=False, loc="upper center", bbox_to_anchor=(0.5, -0.14)); despine(b)
b.set_title("Same local time every pass, all year", fontsize=11.5)
fig.savefig("img/fig-gt-sso.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ geosynchronous: NavIC-like
aG = (MU_E*(T_STELLAR/(2*np.pi))**2)**(1/3)
tg = np.arange(0, T_STELLAR + 1, 60.0)
fig = plt.figure(figsize=(W + 4.6, 6.0))
gs = fig.add_gridspec(1, 2, width_ratios=[1.75, 1], wspace=0.25)
ax = fig.add_subplot(gs[0]); draw_map(ax, lon0=10, lon1=160, lat0=-50, lat1=50, grid=10)
ax.set_xticks(range(10, 161, 20)); ax.set_yticks(range(-50, 51, 10))
for lon_node, col, lab in ((55.0, BLUE, "IGSO, node at 55°E"), (111.75, GREEN, "IGSO, node at 111.75°E")):
    Om = np.radians(lon_node)*0 + lon_node        # theta0 = 0 : node longitude at t=0 equals Omega
    tr = track(aG, 0.0, 29.0, Om, 0.0, 0.0, tg)
    plot_track(ax, tr["lon"], tr["lat"], color=col, lw=2.6)
    ax.plot([], [], color=col, lw=2.6, label=lab + ", $i=29^\\circ$")
for lon_g in (32.5, 83.0, 131.5):
    ax.plot(lon_g, 0, "o", color=MAROON, ms=9, mec="white", zorder=8)
    ax.text(lon_g, -4.2, "GEO\n%.1f°E" % lon_g, ha="center", va="top", fontsize=8.6, color=MAROON)
ax.plot([], [], "o", color=MAROON, label="geostationary ($i=0$: a point)")
ax.legend(loc="upper right", fontsize=9, frameon=True, facecolor="white", framealpha=0.92)
ax.set_title("A regional navigation constellation (NavIC-like): 3 GEO + inclined GSO", fontsize=11.5)
# small-inclination figure 8
b = fig.add_subplot(gs[1])
u = np.linspace(0, 2*np.pi, 721)
for inc, col in ((5, GREEN), (10, AMBER), (20, MAROON), (29, BLUE)):
    i_ = inc*D2R
    lat_ = np.degrees(np.arcsin(np.sin(i_)*np.sin(u)))
    dlon = np.degrees(np.arctan2(np.cos(i_)*np.sin(u), np.cos(u)) - u)
    dlon = (dlon + 180) % 360 - 180
    b.plot(dlon, lat_, color=col, lw=2.0, label="$i=%d^\\circ$" % inc)
    b.plot(np.degrees(-(i_**2/4)*np.sin(2*u)), lat_, color=col, lw=1.0, ls=":")
    print("GSO i=%d deg: exact lon excursion +-%.3f deg, approx i^2/4 = %.3f deg; lat excursion +-%d" % (inc, dlon.max(), np.degrees(i_**2/4), inc))
b.set_xlabel("longitude $-$ longitude of the node (deg)"); b.set_ylabel("latitude (deg)")
b.axhline(0, color=INK, lw=0.5); b.axvline(0, color=INK, lw=0.5); b.legend(frameon=False, fontsize=9, loc="lower right"); despine(b)
b.set_title("The figure 8: height $2i$, width $\\approx i^2/2$\n(dotted: small-$i$ formula)", fontsize=11)
fig.savefig("img/fig-gt-gso.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ Molniya and Tundra
def apogee_lon_setup(a, e, i_deg, lon_apo_deg):
    """Choose Omega so that the first apogee (M0 = 0 at perigee, omega = 270 deg) is at lon_apo_deg."""
    T = period(a)
    theta = OMEGA_E*T/2
    return lon_apo_deg - 90.0 + np.degrees(theta)          # apogee at u = 90 deg -> alpha = Omega + 90 deg

aM = (MU_E*(T_STELLAR/2/(2*np.pi))**2)**(1/3)
eM = 0.74
OmM = apogee_lon_setup(aM, eM, 63.4, 90.0)
TM = period(aM)
tm = np.arange(0, 2*TM + 1, 30.0)
aT, eT = aG, 0.27
OmT = apogee_lon_setup(aT, eT, 63.4, 100.0)
tt_ = np.arange(0, T_STELLAR + 1, 60.0)
fig = plt.figure(figsize=(W + 4, 5.6))
gs = fig.add_gridspec(1, 2, wspace=0.12)
for k, (lab, a_, e_, Om_, tarr, col, hours) in enumerate((
        ("Molniya: $T=11$ h 58 m, $e=0.74$, $i=63.4^\\circ$ (two orbits)", aM, eM, OmM, tm, MAROON, 1.0),
        ("Tundra: $T=23$ h 56 m, $e=0.27$, $i=63.4^\\circ$ (one orbit)", aT, eT, OmT, tt_, BLUE, 2.0))):
    ax = fig.add_subplot(gs[k]); draw_map(ax, lon0=-180, lon1=180, lat0=-90, lat1=90, grid=30, labels=(k == 0))
    if k == 1:
        ax.set_xticks(range(-180, 181, 60)); ax.set_xlabel("longitude (deg E)")
    tr = track(a_, e_, 63.4, Om_, 270.0, 0.0, tarr, j2=True)
    plot_track(ax, tr["lon"], tr["lat"], color=col, lw=2.0)
    step = int(hours*3600/(tarr[1] - tarr[0]))
    idx = np.arange(0, len(tarr), step)
    ax.plot(((tr["lon"][idx] + 180) % 360) - 180, tr["lat"][idx], "o", color=AMBER, mec=INK, ms=5.5, zorder=7)
    ax.set_title(lab + "\nmarkers every %g h: the satellite lingers near apogee" % hours, fontsize=10.5)
    print("%s: a = %.1f km, perigee alt %.0f km, apogee alt %.0f km, T = %.2f min" % (lab.split(":")[0], a_, a_*(1 - e_) - RE, a_*(1 + e_) - RE, period(a_)/60))
    print("   min/max latitude %.1f / %.1f ; max alt speed ratio v_peri/v_apo = %.2f" % (tr["lat"].min(), tr["lat"].max(), (1 + e_)/(1 - e_)))
fig.savefig("img/fig-gt-molniya.png", dpi=160, bbox_inches="tight"); plt.close(fig)
# dwell time above 55 deg north for Molniya
trm = track(aM, eM, 63.4, OmM, 270.0, 0.0, np.arange(0, TM, 10.0), j2=True)
print("Molniya: time above 55N per orbit = %.2f h of %.2f h" % ((trm["lat"] > 55).sum()*10/3600, TM/3600))