"""Earth-gravity note: torque picture, secular J2 rates, numerical validation."""
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

MU, RE, J2 = 398600.4418, 6378.137, 1.08262668e-3
DAY = 86400.; D2R = np.pi/180

def omega_dot(a, e, i):   # rad/s
    n = np.sqrt(MU/a**3); return -1.5*n*J2*RE**2*np.cos(i)/(a**2*(1-e**2)**2)
def peri_dot(a, e, i):
    n = np.sqrt(MU/a**3); return 0.75*n*J2*RE**2*(5*np.cos(i)**2-1)/(a**2*(1-e**2)**2)

# ------------------------------------------------------------------ torque picture (3-D)
fig = plt.figure(figsize=(7.6, 6.6)); ax = fig.add_subplot(111, projection="3d")
i = np.radians(55); Ro = 1.9
t = np.linspace(0, 2*np.pi, 300)
orb = Ro*np.array([np.cos(t), np.sin(t)*np.cos(i), np.sin(t)*np.sin(i)])
ax.plot(*orb, color=BLUE, lw=2.4)
eqr = 1.0*np.array([np.cos(t), np.sin(t), 0*t]); ax.plot(*eqr, color=AMBER, lw=5, alpha=.9)
big = 2.5*np.array([np.cos(t), np.sin(t), 0*t]); ax.plot(*big, color=INK, lw=.6, ls=":")
u, v = np.meshgrid(np.linspace(0, 2*np.pi, 40), np.linspace(-np.pi/2, np.pi/2, 20))
ax.plot_wireframe(.95*np.cos(v)*np.cos(u), .95*np.cos(v)*np.sin(u), .95*np.sin(v), color="#9db4d0", lw=.3, alpha=.6)
top = Ro*np.array([0, np.cos(i), np.sin(i)]); bot = -top
for P, dz in ((top, -1), (bot, +1)):
    ax.quiver(*P, 0, 0, dz*.9, color=PURPLE, lw=2.6, arrow_length_ratio=.25)
    ax.plot(*P, "o", color=BLUE, ms=7)
ax.text(*(top+[.35, .05, .1]), "pulled down toward\nthe equatorial bulge", color=PURPLE, fontsize=10)
ax.text(*(bot+[0, -.1, -.85]), "pulled up", color=PURPLE, fontsize=10)
h = np.array([0, -np.sin(i), np.cos(i)])*1.6
ax.quiver(0, 0, 0, *h, color=GREEN, lw=3, arrow_length_ratio=.12); ax.text(*(h*1.02+[.05, 0, .05]), r"$\mathbf{h}$", color=GREEN, fontsize=15)
ax.quiver(0, 0, 0, -1.7, 0, 0, color=MAROON, lw=3, arrow_length_ratio=.12); ax.text(-2.4, .05, .2, r"torque $\tau$", color=MAROON, fontsize=12)
ax.plot([0, 2.5], [0, 0], [0, 0], color=INK, lw=1.1); ax.text(1.55, 0, .25, "ascending node", fontsize=9.5)
ax.quiver(2.35, 0, 0, 0, -.9, 0, color=GREEN, lw=2.4, arrow_length_ratio=.3); ax.text(.4, -1.4, -.75, r"node moves west ($\dot\Omega<0$)", fontsize=9.5, color=GREEN)
ax.plot([0, 0], [0, 0], [-1.6, 2.2], color=INK, lw=1); ax.text(-.9, 0, 2.3, "spin axis", fontsize=9.5)
ax.set_box_aspect((1, 1, .8), zoom=1.35); ax.set_xlim(-2.6, 3.2); ax.set_ylim(-2.2, 2.2); ax.set_zlim(-2, 2.4)
ax.set_axis_off(); ax.view_init(20, -78)
fig.savefig("img/fig-grav-torque.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ rates vs inclination, SSO, critical inclination
fig, (a1, a2, a3) = plt.subplots(1, 3, figsize=(12.6, 4.2))
inc = np.radians(np.linspace(0, 180, 721))
for h, c in ((400, BLUE), (800, GREEN), (1500, AMBER), (35786, MAROON)):
    a = RE + h; a1.plot(np.degrees(inc), omega_dot(a, 0, inc)*DAY/D2R, color=c, lw=2, label=f"$h={h}$ km")
a1.axhline(0.98565, color=INK, ls="--", lw=1); a1.text(3, 1.4, "+0.9856°/day (Sun)", fontsize=9)
a1.set_ylim(-8, 8); a1.set_xlabel("inclination $i$ (deg)"); a1.set_ylabel(r"$\dot\Omega$  (deg/day)")
a1.set_title(r"Node drift, circular orbits", fontsize=11.5); a1.legend(frameon=False, fontsize=9, loc="lower left"); despine(a1)
hh = np.linspace(300, 1500, 300); a = RE + hh; n = np.sqrt(MU/a**3)
target = 2*np.pi/(365.2422*DAY)
isso = np.degrees(np.arccos(-target*2*a**2/(3*n*J2*RE**2)))
a2.plot(hh, isso, color=BLUE, lw=2.4); a2.set_xlabel("altitude (km)"); a2.set_ylabel("inclination for SSO (deg)")
a7 = RE+700; i7 = np.degrees(np.arccos(-target*2*a7**2/(3*np.sqrt(MU/a7**3)*J2*RE**2)))
a2.plot(700, i7, "o", color=MAROON, ms=9); a2.text(760, i7-.9, f"700 km: {i7:.2f}°", color=MAROON, fontsize=10)
a2.set_title("Sun-synchronous inclination", fontsize=11.5); despine(a2)
print("SSO inclination at 700 km: %.3f deg" % i7)
for h in (300, 400, 500, 600, 800, 1000): 
    a_ = RE+h; n_ = np.sqrt(MU/a_**3); print("SSO", h, np.degrees(np.arccos(-target*2*a_**2/(3*n_*J2*RE**2))))
for a_, e_, c, lab in ((RE+700, 0, BLUE, "LEO 700 km, $e=0$"), (26562, .74, MAROON, "Molniya, $e=0.74$")):
    a3.plot(np.degrees(inc), peri_dot(a_, e_, inc)*DAY/D2R, color=c, lw=2, label=lab)
a3.axhline(0, color=INK, lw=.6)
for x in (np.degrees(np.arccos(1/np.sqrt(5))), 180-np.degrees(np.arccos(1/np.sqrt(5)))): a3.axvline(x, color=INK, ls="--", lw=1)
a3.text(66, 7.2, "63.4°", fontsize=10); a3.text(119, 7.2, "116.6°", fontsize=10)
a3.set_ylim(-7, 8); a3.set_xlabel("inclination $i$ (deg)"); a3.set_ylabel(r"$\dot\omega$  (deg/day)")
a3.set_title("Perigee drift stops at the critical inclination", fontsize=11.5); a3.legend(frameon=False, fontsize=9, loc="lower right"); despine(a3)
fig.tight_layout(); fig.savefig("img/fig-grav-j2-rates.png", dpi=160, bbox_inches="tight"); plt.close(fig)
for lab, a_, e_, i_ in (("ISS", RE+420, 0.0005, 51.6), ("LEO 700 / 55", RE+700, 0, 55.), ("GPS", 26560, 0., 55.), ("Molniya-like i=63.4", 26562, .74, 63.435)):
    print(lab, "Omega_dot=%.4f deg/day  omega_dot=%.4f deg/day" % (omega_dot(a_, e_, i_*D2R)*DAY/D2R, peri_dot(a_, e_, i_*D2R)*DAY/D2R))

# ------------------------------------------------------------------ numerical validation
def acc(t, s):
    r = s[:3]; v = s[3:]; rn = np.linalg.norm(r); x, y, z = r
    k = -1.5*J2*MU*RE**2/rn**5
    a = -MU*r/rn**3 + k*np.array([x*(1-5*z**2/rn**2), y*(1-5*z**2/rn**2), z*(3-5*z**2/rn**2)])
    return np.concatenate([v, a])
def coe2rv(a, e, i, Om, w, nu):
    p = a*(1-e**2); rp = p/(1+e*np.cos(nu))*np.array([np.cos(nu), np.sin(nu), 0]); vp = np.sqrt(MU/p)*np.array([-np.sin(nu), e+np.cos(nu), 0])
    def R3(x): c, s = np.cos(x), np.sin(x); return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
    def R1(x): c, s = np.cos(x), np.sin(x); return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])
    Q = R3(Om) @ R1(i) @ R3(w); return Q @ rp, Q @ vp
def rv2coe(r, v):
    h = np.cross(r, v); hn = np.linalg.norm(h); ev = np.cross(v, h)/MU - r/np.linalg.norm(r)
    i = np.arccos(h[2]/hn); N = np.array([-h[1], h[0], 0]); Om = np.arctan2(h[0], -h[1])
    w = np.arctan2(np.dot(np.cross(N, ev), h)/hn, np.dot(N, ev)); return i, Om, w, np.linalg.norm(ev)
a0, e0, i0, Om0, w0 = 8000., 0.1, 50*D2R, 40*D2R, 30*D2R
r0, v0 = coe2rv(a0, e0, i0, Om0, w0, 0.)
T = 5*DAY; tt = np.linspace(0, T, 6000)
sol = solve_ivp(acc, (0, T), np.concatenate([r0, v0]), t_eval=tt, rtol=1e-12, atol=1e-9, method="DOP853")
el = np.array([rv2coe(sol.y[:3, k], sol.y[3:, k]) for k in range(len(tt))])
Om = np.unwrap(el[:, 1]); w = np.unwrap(el[:, 2])
Om_pred = Om0 + omega_dot(a0, e0, i0)*tt; w_pred = w0 + peri_dot(a0, e0, i0)*tt
slope = lambda y: np.polyfit(tt, y, 1)[0]*DAY/D2R
print("Omega slope: numeric %.4f analytic %.4f deg/day" % (slope(Om), omega_dot(a0, e0, i0)*DAY/D2R))
print("omega slope: numeric %.4f analytic %.4f deg/day" % (slope(w), peri_dot(a0, e0, i0)*DAY/D2R))
print("i range (deg): %.4f .. %.4f" % (el[:, 0].min()/D2R, el[:, 0].max()/D2R))
fig, (a1, a2, a3) = plt.subplots(1, 3, figsize=(12.6, 4.2))
td = tt/DAY
a1.plot(td, Om/D2R, color=BLUE, lw=2.4, label="numerical (with $J_2$)"); a1.plot(td, Om_pred/D2R, color=AMBER, lw=1.6, ls="--", label=r"$\Omega_0+\dot\Omega\,t$ (formula)")
a1.set_xlabel("time (days)"); a1.set_ylabel(r"node $\Omega$ (deg)"); a1.legend(frameon=False, fontsize=9.5); a1.set_title("Node regresses linearly", fontsize=11.5); despine(a1)
a2.plot(td, w/D2R, color=BLUE, lw=2.4, label="numerical"); a2.plot(td, w_pred/D2R, color=AMBER, lw=1.6, ls="--", label="formula")
a2.set_xlabel("time (days)"); a2.set_ylabel(r"perigee $\omega$ (deg)"); a2.legend(frameon=False, fontsize=9.5); a2.set_title("Perigee advances ($i<63.4°$)", fontsize=11.5); despine(a2)
a3.plot(td[tt < 0.5*DAY*4], (el[:, 0]/D2R)[tt < 0.5*DAY*4], color=MAROON, lw=2)
a3.set_xlabel("time (days)"); a3.set_ylabel("inclination $i$ (deg)"); a3.set_title("Inclination only wobbles (no secular drift)", fontsize=11.5); despine(a3)
fig.tight_layout(); fig.savefig("img/fig-grav-validation.png", dpi=160, bbox_inches="tight"); plt.close(fig)