"""Ground-track note: the method.

    img/fig-gt-pipeline.png   the six steps, elements -> latitude/longitude
    img/fig-gt-geometry.png   the spherical right triangle behind the circular-orbit formula
    img/fig-gt-j2check.png    how good is the mean-element (secular J2) ground track?

Run from the deck directory:   python codes/fig-gt-method.py
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
from scipy.integrate import solve_ivp
from scipy.optimize import least_squares
from _gt import *
from _timeutil import MU_E, RE, J2, OMEGA_E, D2R, DAY, j2_rates
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine

# ================================================================== 1. pipeline
fig, ax = plt.subplots(figsize=(15.0, 5.0)); ax.axis("off"); ax.set_xlim(0, 15); ax.set_ylim(0, 5.2)
steps = [
    ("1", "Elements\nat epoch", "$a,\\ e,\\ i,\\ \\Omega,\\ \\omega,\\ M_0$\nand the Greenwich\nangle $\\theta_0$", GREY := "#8b8378"),
    ("2", "Advance\nthe clock", "$M=M_0+n\\,(t-t_0)$\nsolve\n$E-e\\sin E=M$", BLUE),
    ("3", "Position in the\norbit plane", "$x_P=a(\\cos E-e)$\n$y_Q=a\\sqrt{1-e^2}\\,\\sin E$", BLUE),
    ("4", "Rotate to\ninertial axes (ECI)", "$\\mathbf{r}_{\\rm ECI}=R_3(-\\Omega)\\,R_1(-i)$\n$\\times R_3(-\\omega)\\,\\mathbf{r}_{PQW}$", BLUE),
    ("5", "Rotate to\nEarth-fixed (ECEF)", "$\\mathbf{r}_{\\rm ECEF}=R_3(\\theta)\\,\\mathbf{r}_{\\rm ECI}$\n$\\theta=\\theta_0+\\omega_\\oplus t$", GREEN),
    ("6", "Latitude and\nlongitude", "$\\lambda=\\mathrm{atan2}(y,x)$\n$\\varphi=\\arcsin(z/r)$\nor geodetic", GREEN),
]
xs = np.linspace(1.25, 13.75, 6)
for (n, head, body, col), x in zip(steps, xs):
    ax.add_patch(FancyBboxPatch((x - 1.08, 1.5), 2.16, 2.5, boxstyle="round,pad=0.03,rounding_size=0.12", fc="white", ec=col, lw=2.4))
    ax.text(x, 3.72, n, ha="center", va="center", fontsize=15, color="white", bbox=dict(boxstyle="circle,pad=0.25", fc=col, ec="none"))
    ax.text(x, 3.12, head, ha="center", va="center", fontsize=9.6, weight="bold", color=col, linespacing=1.15)
    ax.text(x, 2.15, body, ha="center", va="center", fontsize=9.4, linespacing=1.5)
for x0, x1 in zip(xs[:-1], xs[1:]):
    ax.add_patch(FancyArrowPatch((x0 + 1.12, 2.75), (x1 - 1.12, 2.75), arrowstyle="-|>", mutation_scale=18, color=INK, lw=1.6))
ax.annotate("", xy=(xs[3] + 1.08, 1.05), xytext=(xs[1] - 1.08, 1.05), arrowprops=dict(arrowstyle="<->", color=BLUE, lw=1.4))
ax.text(0.5*(xs[1] + xs[3]), 0.72, "the orbit is a fixed conic in inertial space (steps 2-4)", ha="center", color=BLUE, fontsize=10.5)
ax.annotate("", xy=(xs[5] + 1.08, 1.05), xytext=(xs[4] - 1.08, 1.05), arrowprops=dict(arrowstyle="<->", color=GREEN, lw=1.4))
ax.text(0.5*(xs[4] + xs[5]), 0.72, "the Earth turns underneath (steps 5-6)", ha="center", color=GREEN, fontsize=10.5)
ax.text(7.5, 0.15, "Only step 5 involves the rotation of the Earth, and only step 5 involves time as measured by the Earth's clock ($\\theta$).", ha="center", fontsize=10.5, style="italic")
fig.savefig("img/fig-gt-pipeline.png", dpi=160, bbox_inches="tight"); plt.close(fig)

# ================================================================== 2. spherical triangle
inc, u_s = np.radians(40.0), np.radians(62.0)
fig = plt.figure(figsize=(8.6, 7.4)); ax = fig.add_subplot(111, projection="3d")
uu, vv = np.meshgrid(np.linspace(0, 2*np.pi, 60), np.linspace(-np.pi/2, np.pi/2, 30))
ax.plot_wireframe(np.cos(vv)*np.cos(uu), np.cos(vv)*np.sin(uu), np.sin(vv), color="#c9d3dc", lw=0.35, alpha=0.7)
tt = np.linspace(0, 2*np.pi, 300)
ax.plot(np.cos(tt), np.sin(tt), 0*tt, color=GREEN, lw=1.6)
orb = lambda u: np.array([np.cos(u), np.sin(u)*np.cos(inc), np.sin(u)*np.sin(inc)])
ax.plot(*orb(tt), color=BLUE, lw=1.6)
uarc = np.linspace(0, u_s, 60); O = orb(uarc); ax.plot(*O, color=BLUE, lw=4.4)
alpha_s = np.arctan2(np.cos(inc)*np.sin(u_s), np.cos(u_s))
aarc = np.linspace(0, alpha_s, 60); ax.plot(np.cos(aarc), np.sin(aarc), 0*aarc, color=GREEN, lw=4.4)
S = orb(u_s); F = np.array([np.cos(alpha_s), np.sin(alpha_s), 0])
phi_s = np.arcsin(S[2]); parc = np.linspace(0, phi_s, 40)
ax.plot(np.cos(parc)*np.cos(alpha_s), np.cos(parc)*np.sin(alpha_s), np.sin(parc), color=MAROON, lw=4.4)
for P, lab, off, ha in ((np.array([1, 0, 0]), "N: ascending node", (0.0, 0.0, -0.10), "right"), (S, "S: satellite", (0.03, 0.0, 0.06), "left"), (F, "F", (0.05, 0.0, -0.08), "left")):
    ax.scatter(*P, color=INK, s=48, zorder=10); ax.text(*(P + np.array(off)), lab, fontsize=11, weight="bold", ha=ha)
ax.plot([0, 0], [0, 0], [-1.25, 1.25], color=INK, lw=1); ax.text(0, 0, 1.3, "north", ha="center", fontsize=10)
ax.text(*(orb(0.5*u_s)*1.06 + np.array([-0.05, 0.06, 0.13])), "$u$", color=BLUE, fontsize=17, ha="center")
ax.text(np.cos(0.5*alpha_s)*1.05, np.sin(0.5*alpha_s)*1.05, -0.16, "$\\alpha-\\Omega$", color=GREEN, fontsize=15, ha="center")
ax.text(np.cos(0.5*phi_s)*np.cos(alpha_s)*1.13, np.cos(0.5*phi_s)*np.sin(alpha_s)*1.13 + 0.06, 0.5*S[2] + 0.02, "$\\varphi$", color=MAROON, fontsize=17)
ax.text(0.90, 0.20, 0.0, "$i$", fontsize=16, color=INK, ha="center")
ax.set_box_aspect((1, 1, 0.9), zoom=1.25); ax.set_axis_off(); tri = np.array([1, 0, 0]) + F + S; tri /= np.linalg.norm(tri)
nrm = np.array([0, -np.sin(inc), np.cos(inc)])
cam = tri + 0.75*nrm; cam /= np.linalg.norm(cam)
ax.view_init(np.degrees(np.arcsin(cam[2])), np.degrees(np.arctan2(cam[1], cam[0])))
fig.savefig("img/fig-gt-geometry.png", dpi=160, bbox_inches="tight"); plt.close(fig)
print("triangle check: sin(phi)=%.6f vs sin i sin u=%.6f ; tan(alpha)=%.6f vs cos i tan u=%.6f" % (
    np.sin(phi_s), np.sin(inc)*np.sin(u_s), np.tan(alpha_s), np.cos(inc)*np.tan(u_s)))

# ================================================================== 3. mean elements vs numerical integration
def a_j2(r):
    rn = np.linalg.norm(r); x, y, z = r
    k = -1.5*J2*MU_E*RE**2/rn**5
    return -MU_E*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)])

def elements_to_state(a, e, i, Om, w, M):
    E = float(kepler_E(M, e)); xp, yp = a*(np.cos(E) - e), a*np.sqrt(1 - e*e)*np.sin(E)
    vx, vy = -np.sqrt(MU_E*a)/np.hypot(xp, yp)*np.sin(E), np.sqrt(MU_E*a*(1 - e*e))/np.hypot(xp, yp)*np.cos(E)
    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 @ np.array([xp, yp, 0.0]), Q @ np.array([vx, vy, 0.0])

a0, e0, i0 = RE + 800, 0.0011, np.radians(98.6)
r0, v0 = elements_to_state(a0, e0, i0, np.radians(30), np.radians(90), np.radians(0))
tt_ = np.arange(0, 3*DAY + 1, 30.0)
sol = solve_ivp(lambda t, y: np.concatenate([y[3:], a_j2(y[:3])]), (0, tt_[-1]), np.concatenate([r0, v0]), t_eval=tt_, rtol=1e-11, atol=1e-9, method="DOP853")
X = sol.y[:3]
th = OMEGA_E*tt_
ecef = np.array([np.cos(th)*X[0] + np.sin(th)*X[1], -np.sin(th)*X[0] + np.cos(th)*X[1], X[2]])
lat_n = np.degrees(np.arcsin(ecef[2]/np.linalg.norm(ecef, axis=0))); lon_n = np.degrees(np.arctan2(ecef[1], ecef[0]))

def gc_dist(la1, lo1, la2, lo2):
    la1, lo1, la2, lo2 = [np.radians(v) for v in (la1, lo1, la2, lo2)]
    c = np.sin(la1)*np.sin(la2) + np.cos(la1)*np.cos(la2)*np.cos(lo1 - lo2)
    return np.degrees(np.arccos(np.clip(c, -1, 1)))*np.pi/180*RE

def model(p, tarr):
    a, e, inc, om, w, M0 = p
    tr = track(a, e, inc, om, w, M0, tarr, j2=True)
    return tr["lat_gc"], tr["lon"]

p0 = [a0, e0, np.degrees(i0), 30.0, 90.0, 0.0]
naive = model(p0, tt_)   # osculating elements used as if they were mean elements
sep_naive = gc_dist(*naive, lat_n, lon_n)
res = least_squares(lambda p: gc_dist(*model(p, tt_), lat_n, lon_n), p0, x_scale=[10, 1e-3, 0.1, 0.1, 5.0, 5.0],
                    bounds=([6800, 0, 97, 20, 0, -180], [7600, 0.01, 100, 40, 360, 180]))
best = model(res.x, tt_)
sep_fit = gc_dist(*best, lat_n, lon_n)
print("J2 check: osculating a = %.2f km -> fitted MEAN a = %.2f km (difference %.2f km); e: %.5f -> %.5f; i: %.4f -> %.4f deg" % (a0, res.x[0], res.x[0] - a0, e0, res.x[1], np.degrees(i0), res.x[2]))
print("  naive (osculating used as mean): separation after 1 orbit %.0f km, 1 day %.0f km, 3 days %.0f km" % (
    sep_naive[int(6060/30)], sep_naive[int(DAY/30)], sep_naive[-1]))
print("  fitted mean elements: rms %.2f km, max %.2f km over 3 days" % (np.sqrt((sep_fit**2).mean()), sep_fit.max()))
fig, (c1, c2) = plt.subplots(1, 2, figsize=(14.0, 4.6), gridspec_kw=dict(width_ratios=[1, 1.15]))
m = tt_ <= 3*6060
c1.plot(tt_[m]/3600, sep_naive[m], color=MAROON, lw=2.2, label="osculating elements used as mean elements")
c1.plot(tt_[m]/3600, sep_fit[m], color=GREEN, lw=2.2, label="mean elements fitted to the integration")
c1.set_xlabel("time (h)"); c1.set_ylabel("distance between the two ground tracks (km)"); despine(c1)
c1.set_title("Secular-$J_2$ model vs numerical integration (first 3 orbits)", fontsize=11); c1.legend(frameon=False, fontsize=9)
c2.plot(tt_/86400, sep_fit, color=GREEN, lw=1.3); c2.set_xlabel("days"); c2.set_ylabel("separation (km)"); despine(c2)
c2.set_title("Fitted mean elements over 3 days: the error is periodic and stays below %.0f km" % np.ceil(sep_fit.max()), fontsize=11)
fig.tight_layout(); fig.savefig("img/fig-gt-j2check.png", dpi=160, bbox_inches="tight"); plt.close(fig)