"""Earth-gravity note: coefficient spectrum, J2 acceleration field, GEO longitude drift (C22, S22)."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

MU, RE, J2 = 398600.4418, 6378.137, 1.08262668e-3       # km^3/s^2, km
C22, S22 = 1.5745e-6, -9.0387e-7                        # unnormalised EGM-type values

# ------------------------------------------------------------------ spectrum
fig, ax = plt.subplots(figsize=(5.4, 4.3))
n = np.arange(2, 21)
ax.semilogy(n, 1e-5/n**2, color=INK, ls="--", lw=1.6, label=r"Kaula's rule $\sim 10^{-5}/n^2$")
for nn, val, c in ((2, 1.08263e-3, MAROON), (3, 2.5327e-6, BLUE), (4, 1.6196e-6, GREEN)):
    ax.plot(nn, val, "o", color=c, ms=9)
    ax.text(nn+.4, val*1.15, rf"$|J_{nn}|$", color=c, fontsize=11)
ax.axhline(1, color=AMBER, lw=1.6); ax.text(9, 1.6, r"point mass ($n=0$): $1$", color=AMBER, fontsize=10)
ax.set_ylim(1e-8, 10); ax.set_xticks([2, 4, 6, 8, 10, 12, 14, 16, 18, 20]); ax.set_xlabel("degree $n$"); ax.set_ylabel(r"size of term relative to $\mu/r$ at $r=R$")
ax.legend(frameon=False, fontsize=9.5, loc="lower left"); despine(ax)
fig.tight_layout(); fig.savefig("img/fig-grav-spectrum.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ------------------------------------------------------------------ J2 perturbing acceleration
def a_cart(x, z):
    r = np.hypot(x, z); k = -1.5*J2*MU*RE**2/r**5
    return k*x*(1 - 5*z**2/r**2), k*z*(3 - 5*z**2/r**2)         # (a_x, a_z), y = 0 plane
def a_sph(r, ph):
    k = 1.5*J2*MU*RE**2/r**4
    return k*(3*np.sin(ph)**2 - 1), -k*np.sin(2*ph)              # (a_r, a_phi)
# self-check: Cartesian and spherical forms agree
r0, p0 = 7078., np.radians(37.); x0, z0 = r0*np.cos(p0), r0*np.sin(p0)
ax_, az_ = a_cart(x0, z0); ar, ap = a_sph(r0, p0)
er = np.array([np.cos(p0), np.sin(p0)]); ep = np.array([-np.sin(p0), np.cos(p0)])
assert np.allclose([ax_, az_], ar*er + ap*ep, rtol=1e-12), "J2 forms disagree"

fig, (a1, a2) = plt.subplots(1, 2, figsize=(11.5, 5.2), gridspec_kw=dict(width_ratios=[1.15, 1]))
gx, gz = np.meshgrid(np.linspace(0.15, 3.0, 15), np.linspace(-3.0, 3.0, 21))
X, Z = gx*RE, gz*RE; Rr = np.hypot(X, Z); mask = Rr > 1.12*RE
Ax, Az = a_cart(X, Z); mag = np.hypot(Ax, Az)/(MU/Rr**2)
q = a1.quiver(X[mask]/RE, Z[mask]/RE, (Ax/np.hypot(Ax, Az))[mask], (Az/np.hypot(Ax, Az))[mask],
              mag[mask]*1e3, cmap="viridis", pivot="mid", angles="xy", scale_units="xy", scale=3.6, width=.005)
cb = fig.colorbar(q, ax=a1, fraction=.046, pad=.02); cb.set_label(r"$|\mathbf{a}_{J_2}|\,/\,(\mu/r^2)$  ($\times10^{-3}$)")
tt = np.linspace(0, 2*np.pi, 300)
from matplotlib.patches import Ellipse
a1.add_patch(Ellipse((0, 0), 2, 2*(1-.12), color="#cfdcec", zorder=3)); a1.plot(np.cos(tt), (1-.12)*np.sin(tt), color=BLUE, lw=1.5, zorder=4)
a1.text(0.45, 0.0, "Earth", ha="center", va="center", fontsize=10, zorder=5)
a1.set_xlim(0, 3.2); a1.set_ylim(-3.2, 3.2); a1.set_aspect("equal")
a1.set_xlabel(r"distance from spin axis  $\rho/R$"); a1.set_ylabel(r"$z/R$")
a1.set_title("Direction of the $J_2$ perturbing acceleration (meridian plane)", fontsize=11.5); despine(a1)

ph = np.radians(np.linspace(-90, 90, 361)); rr = RE + 700
ar, ap = a_sph(rr, ph); s = MU/rr**2
a2.plot(np.degrees(ph), ar/s*1e3, color=BLUE, lw=2.2, label=r"radial  $a_r$  (+ outward)")
a2.plot(np.degrees(ph), ap/s*1e3, color=MAROON, lw=2.2, label=r"northward  $a_\varphi$")
a2.axhline(0, color=INK, lw=.6); a2.set_xlim(-90, 90); a2.set_xticks(range(-90, 91, 30))
a2.set_xlabel(r"latitude $\varphi$ (deg)"); a2.set_ylabel(r"acceleration / $(\mu/r^2)$   ($\times10^{-3}$)")
a2.set_title(r"Components at $r=R+700$ km", fontsize=11.5); a2.legend(frameon=False, fontsize=10, loc="upper center"); despine(a2)
fig.tight_layout(); fig.savefig("img/fig-grav-j2-field.png", dpi=160, bbox_inches="tight"); plt.close(fig)
print("J2 accel ratio at 700 km, equator/pole (units of mu/r^2):", a_sph(rr, 0)[0]/(MU/rr**2), a_sph(rr, np.pi/2)[0]/(MU/rr**2))

# ------------------------------------------------------------------ GEO longitude drift from C22, S22
J22 = np.hypot(C22, S22); lam22 = 0.5*np.arctan2(S22, C22)
ag = 42164.17
K = 6*MU*RE**2*J22/ag**4                                        # km/s^2
lamddot = lambda lam: (-3/ag)*(-K*np.sin(2*(lam-lam22)))         # rad/s^2
conv = (86400**2)*180/np.pi
print("J22 = %.4e, lam22 = %.2f deg, peak lamddot = %.5f deg/day^2" % (J22, np.degrees(lam22), 3*K/ag*conv))
print("stable:", np.degrees(lam22)+90, np.degrees(lam22)-90, " unstable:", np.degrees(lam22), np.degrees(lam22)+180)

fig, (a1, a2) = plt.subplots(1, 2, figsize=(11.5, 4.9), gridspec_kw=dict(width_ratios=[1, 1.25]))
t = np.linspace(0, 2*np.pi, 500)
rad = 1 + .09*np.cos(2*(t-lam22))
a1.fill(rad*np.cos(t), rad*np.sin(t), color="#cfdcec"); a1.plot(rad*np.cos(t), rad*np.sin(t), color=BLUE, lw=1.6)
a1.text(0, 0, "Earth\n(equator, ellipticity\nexaggerated)", ha="center", va="center", fontsize=9.5)
Rg = 1.75; a1.plot(Rg*np.cos(t), Rg*np.sin(t), color=INK, lw=.8, ls=":")
for ld in np.radians(np.arange(-180, 180, 15)):
    at = -K*np.sin(2*(ld-lam22))/K                                  # in [-1, 1]
    p = Rg*np.array([np.cos(ld), np.sin(ld)]); tang = np.array([-np.sin(ld), np.cos(ld)])
    if abs(at) > .05:
        a1.annotate("", xy=p + .32*at*tang, xytext=p, arrowprops=dict(arrowstyle="-|>", color=PURPLE, lw=2, mutation_scale=13))
    a1.plot(*p, "o", color=INK, ms=2.5)
for ld, c, mk, lab in ((lam22+np.pi/2, GREEN, "o", "stable 75°E"), (lam22-np.pi/2, GREEN, "o", "stable 105°W"),
                       (lam22, MAROON, "s", "unstable"), (lam22+np.pi, MAROON, "s", "unstable")):
    p = 2.2*np.array([np.cos(ld), np.sin(ld)]); a1.plot(*(Rg*np.array([np.cos(ld), np.sin(ld)])), mk, color=c, ms=10, zorder=6)
    a1.text(*p, lab, color=c, fontsize=9.5, ha="center", va="center")
a1.text(0, -2.75, "arrows: tangential acceleration on a GEO satellite (east = counter-clockwise)", fontsize=8.8, ha="center")
a1.set_xlim(-2.9, 2.9); a1.set_ylim(-2.9, 2.6); a1.set_aspect("equal"); a1.axis("off")
a1.set_title("Looking down on the equator", fontsize=11.5)

lam = np.radians(np.linspace(-180, 180, 721))
a2.plot(np.degrees(lam), lamddot(lam)*conv, color=INK, lw=2.4)
a2.axhline(0, color=INK, lw=.6)
for ld, c, mk in ((lam22+np.pi/2, GREEN, "o"), (lam22-np.pi/2, GREEN, "o"), (lam22, MAROON, "s"), (lam22+np.pi, MAROON, "s")):
    a2.plot(np.degrees(ld), 0, mk, color=c, ms=10, zorder=5)
a2.set_xlabel("longitude of the satellite (deg east)"); a2.set_ylabel(r"$\ddot\lambda$  (deg/day$^2$)")
a2.set_title(r"Longitude acceleration $\ddot\lambda=-3a_t/a$: a pendulum with two wells", fontsize=11.5)
a2.set_xticks(range(-180, 181, 45)); despine(a2)
fig.tight_layout(); fig.savefig("img/fig-grav-geo-drift.png", dpi=160, bbox_inches="tight"); plt.close(fig)