"""Figures for orbit-determination.qmd (img/fig-od-*.png).

  fig-od-scenario  the pass used in every example: ground track and sky plot, with the observation times marked
  fig-od-gibbs     the three positions, the eccentricity vector and the vector S in the orbital plane
  fig-od-hg        Gibbs against Herrick-Gibbs: error against the spacing of the positions, with and without noise
  fig-od-laplace   the two relations rho(r) of Laplace's method (and Gauss's), and where they cross
  fig-od-gauss     convergence of the iterated Gauss method
  fig-od-noise     error of Laplace, Gauss and iterated Gauss against the arc, for 2 arcsec of noise
  fig-od-lsq       least squares: convergence, the error cloud (a needle along the line of sight), and the gain over the IOD

Run from the sfm-slides directory:   python3 codes/fig-od.py
The Monte Carlo runs are in _odmc.py (fixed seeds, cached in codes/_cache).
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
import _od as od
from _od import truth_state, observe, gibbs, laplace, gauss, theta_g, SITE, D2R, MU_E
from _odmc import crossover, noise_vs_arc, lsq_experiment, ARCSEC
from _topo import R3
from _uv import coe

r0, v0 = truth_state()
GREY = "#8a8478"
T5 = [-300, -60, 0, 60, 300]


def lab(t, math):
    s = "0" if t == 0 else (f"+{t}" if t > 0 else f"\u2212{-t}")
    return f"t = {s} s" if math else f"{s} s"


def save(fig, name):
    fig.savefig(f"img/{name}.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    print("wrote", name)


# ------------------------------------------------------------------------------------------ scenario
def fig_scenario():
    ts = np.arange(-700, 801, 5.0)
    ob = observe(ts)
    lon, lat = [], []
    for o in ob:
        re = R3(theta_g(o["t"])).T @ o["r"]
        lon.append(np.degrees(np.arctan2(re[1], re[0]))); lat.append(np.degrees(np.arcsin(re[2]/np.linalg.norm(re))))
    lon, lat = np.array(lon), np.array(lat)
    el = np.degrees([o["el"] for o in ob])
    az = np.array([o["az"] for o in ob])
    mark = {int(t): i for i, t in enumerate(ts) if t in T5}
    cols = {-300: GREEN, -60: AMBER, 0: BLUE, 60: AMBER, 300: GREEN}
    fig = plt.figure(figsize=(13.2, 5.4))
    a = fig.add_subplot(1, 2, 1)
    a.plot(lon, lat, color=GREY, lw=1.4)
    up = el > 0
    a.plot(lon[up], lat[up], color=MAROON, lw=2.8, label="above the horizon")
    for t, i in mark.items():
        a.scatter([lon[i]], [lat[i]], color=cols[t], s=55, zorder=5)
        a.annotate(lab(t, True), (lon[i], lat[i]), textcoords="offset points", xytext=(8, -4), fontsize=10, color=cols[t])
    a.scatter([np.degrees(SITE["lam"])], [np.degrees(SITE["phi"])], marker="*", s=230, color=INK, zorder=6)
    a.annotate("station", (np.degrees(SITE["lam"]), np.degrees(SITE["phi"])), textcoords="offset points", xytext=(-48, -14), fontsize=10)
    a.set_xlabel("east longitude [deg]"); a.set_ylabel("sub-satellite latitude [deg]")
    a.set_title("Ground track", fontsize=12); a.legend(frameon=False, loc="upper right", fontsize=9); despine(a)
    b = fig.add_subplot(1, 2, 2, projection="polar")
    b.set_theta_zero_location("N"); b.set_theta_direction(-1)
    b.plot(az[up], 90 - el[up], color=MAROON, lw=2.8)
    for t, i in mark.items():
        b.scatter([az[i]], [90 - el[i]], color=cols[t], s=55, zorder=5)
        b.annotate(lab(t, False), (az[i], 90 - el[i]), textcoords="offset points", xytext=(9, 4), fontsize=10, color=cols[t])
    b.set_ylim(0, 90); b.set_yticks([0, 30, 60, 90]); b.set_yticklabels(["", "30°", "60°", "horizon"], fontsize=8)
    b.set_title("Sky plot (zenith at the centre)", fontsize=12, pad=14)
    fig.tight_layout()
    save(fig, "fig-od-scenario")


# ------------------------------------------------------------------------------------------ Gibbs geometry
def fig_gibbs():
    rr = [observe([t])[0]["r"] for t in (-300.0, 0.0, 300.0)]
    h = np.cross(r0, v0); hh = h/np.linalg.norm(h)
    ev = np.cross(v0, h)/MU_E - r0/np.linalg.norm(r0)
    P = ev/np.linalg.norm(ev); Q = np.cross(hh, P)
    xy = lambda x: np.array([x @ P, x @ Q])
    a_, e_ = coe(r0, v0)[:2]
    p_ = a_*(1 - e_**2)
    f = np.linspace(0, 2*np.pi, 721)
    orb = np.array([p_/(1 + e_*np.cos(x))*np.array([np.cos(x), np.sin(x)]) for x in f])
    n = [np.linalg.norm(x) for x in rr]
    D = np.cross(rr[0], rr[1]) + np.cross(rr[1], rr[2]) + np.cross(rr[2], rr[0])
    S = rr[0]*(n[1] - n[2]) + rr[1]*(n[2] - n[0]) + rr[2]*(n[0] - n[1])
    fig, ax = plt.subplots(figsize=(7.6, 6.6))
    ax.add_patch(plt.Circle((0, 0), 6378.137, color="#d9e6f2", zorder=0))
    ax.plot(orb[:, 0], orb[:, 1], color=GREY, lw=1.3)
    for k, (x, lab, c) in enumerate(zip(rr, (r"$\mathbf{r}_1$", r"$\mathbf{r}_2$", r"$\mathbf{r}_3$"), (GREEN, BLUE, GREEN))):
        q = xy(x)
        ax.annotate("", xy=q, xytext=(0, 0), arrowprops=dict(arrowstyle="-|>", color=c, lw=2))
        ax.annotate(lab, q, textcoords="offset points", xytext=((-24, 6) if k == 2 else (6, 4)), fontsize=13, color=c)
    ax.annotate("", xy=(e_*a_*3, 0), xytext=(0, 0), arrowprops=dict(arrowstyle="-|>", color=MAROON, lw=2.4))
    ax.annotate(r"$\mathbf{e}$ (to perigee, drawn $3\times$)", (e_*a_*3, 0), textcoords="offset points", xytext=(6, 6), fontsize=10.5, color=MAROON)
    sd = xy(S)/np.linalg.norm(S)*3200
    ax.annotate("", xy=sd, xytext=(0, 0), arrowprops=dict(arrowstyle="-|>", color=PURPLE, lw=2.4))
    ax.annotate(r"$\mathbf{S}\perp\mathbf{e}$ (direction)", sd, textcoords="offset points", xytext=(-110, 6), fontsize=10.5, color=PURPLE)
    v2 = xy(gibbs(*rr))
    q2 = xy(rr[1])
    ax.annotate("", xy=q2 + v2*300, xytext=q2, arrowprops=dict(arrowstyle="-|>", color=AMBER, lw=2.4))
    ax.annotate(r"$\mathbf{v}_2$ (Gibbs)", q2 + v2*300, xytext=(-7200, 8800), fontsize=12, color=AMBER, arrowprops=dict(arrowstyle="-", color=AMBER, lw=0.8))
    ax.scatter([0], [0], color=INK, s=30, zorder=5)
    ax.text(300, -700, "focus", fontsize=10)
    ax.set_aspect("equal"); ax.set_xlim(-9800, 9800); ax.set_ylim(-9800, 9800)
    ax.set_xlabel("along the eccentricity vector [km]"); ax.set_ylabel("along $\\hat{h}\\times\\hat{e}$ [km]")
    ax.set_title("The orbital plane seen from $\\mathbf{D}$ (which points out of the page)", fontsize=12); despine(ax)
    save(fig, "fig-od-gibbs")


# ------------------------------------------------------------------------------------------ Herrick-Gibbs
def fig_hg():
    X = crossover()
    tau, sep, g0, h0, gm, hm = X.T
    fig, (a, b) = plt.subplots(1, 2, figsize=(12.6, 4.9))
    a.loglog(sep, np.maximum(g0, 1e-16)*1e3, "--o", color=MAROON, ms=4, lw=1.6, label="Gibbs, exact positions")
    a.loglog(sep, np.maximum(h0, 1e-16)*1e3, "--o", color=BLUE, ms=4, lw=1.6, label="Herrick–Gibbs, exact positions")
    a.loglog(sep, gm*1e3, "-o", color=MAROON, ms=5, lw=2.6, label="Gibbs, 10 m noise")
    a.loglog(sep, hm*1e3, "-o", color=BLUE, ms=5, lw=2.6, label="Herrick–Gibbs, 10 m noise")
    a.axvline(0.742, color=GREY, lw=1, ls=":"); a.axvline(21.2, color=GREY, lw=1, ls=":")
    a.text(0.78, 3e-2, "exact data:\nequal at $0.7^\\circ$", fontsize=9, color=GREY, va="center")
    a.text(22.5, 3e3, "noisy data:\nequal at $21^\\circ$", fontsize=9, color=GREY, va="top")
    a.set_ylim(1e-10, 1e4)
    a.set_xlabel(r"angle between $\mathbf{r}_1$ and $\mathbf{r}_3$ [deg]"); a.set_ylabel(r"error in $\mathbf{v}_2$ [m/s]")
    a.set_title("Error against the spacing of the three positions", fontsize=12); a.legend(frameon=False, fontsize=9, loc="lower right", bbox_to_anchor=(1.0, 0.13)); despine(a)
    # cancellation in D: |D|/|r1 x r2| against the half-spacing angle psi (circle: 2 sin psi - sin 2 psi over sin psi ~ psi^2)
    psis, ratio = [], []
    for t in tau:
        rp = [observe([s])[0]["r"] for s in (-t, 0.0, t)]
        D = np.cross(rp[0], rp[1]) + np.cross(rp[1], rp[2]) + np.cross(rp[2], rp[0])
        psis.append(np.arccos(rp[0] @ rp[1]/np.linalg.norm(rp[0])/np.linalg.norm(rp[1])))
        ratio.append(np.linalg.norm(D)/np.linalg.norm(np.cross(rp[0], rp[1])))
    psis, ratio = np.array(psis), np.array(ratio)
    b.loglog(np.degrees(psis), ratio, "o-", color=PURPLE, lw=2.4, ms=5, label=r"$|\mathbf{D}|\,/\,|\mathbf{r}_1\times\mathbf{r}_2|$")
    b.loglog(np.degrees(psis), psis**2, "--", color=INK, lw=1.2, label=r"$\psi^{2}$")
    b.set_xlabel(r"angle $\psi$ between neighbouring positions [deg]"); b.set_ylabel("size of $\\mathbf{D}$ relative to its terms")
    b.set_title("What is lost in $\\mathbf{D}$ when the points are close", fontsize=12); b.legend(frameon=False, fontsize=10); despine(b)
    fig.tight_layout()
    save(fig, "fig-od-hg")


# ------------------------------------------------------------------------------------------ Laplace curves
def fig_laplace():
    TW = [-60.0, 0.0, 60.0]
    ob = observe(TW)
    ra = np.round(np.degrees([o["ra"] for o in ob]), 6)*D2R
    de = np.round(np.degrees([o["dec"] for o in ob]), 6)*D2R
    L = [np.array([np.cos(d)*np.cos(a), np.cos(d)*np.sin(a), np.sin(d)]) for a, d in zip(ra, de)]
    R = [o["R"] for o in ob]
    _, _, dl = laplace(L, R, ob[1]["Rd"], ob[1]["Rdd"], TW, details=True)
    _, _, dg = gauss(L, R, TW, details=True)
    C, Rn2 = dl["C"], R[1] @ R[1]
    geo = lambda r: -C + np.sqrt(r**2 - Rn2 + C**2)
    lap = lambda r: dl["A"] + dl["B"]/r**3
    gau = lambda r: dg["A"] + MU_E*dg["B"]/r**3
    fig, (a, b) = plt.subplots(1, 2, figsize=(12.6, 4.9))
    r = np.linspace(6500, 9600, 600)
    a.plot(r, geo(r), color=MAROON, lw=2.6, label=r"geometry: $\rho=-C+\sqrt{r^{2}-R^{2}+C^{2}}$")
    a.plot(r, lap(r), color=BLUE, lw=2.6, label=r"dynamics (Laplace): $\rho=A+B/r^{3}$")
    a.scatter([dl["r"]], [dl["rho"]], color=INK, zorder=5, s=40)
    a.annotate(f"root  $r={dl['r']:.1f}$ km", (dl["r"], dl["rho"]), textcoords="offset points", xytext=(10, -22), fontsize=10.5)
    a.set_ylim(-200, 3500); a.set_xlabel("distance $r$ from the centre of the Earth [km]"); a.set_ylabel(r"range $\rho$ [km]")
    a.set_title("Two relations between $\\rho$ and $r$ cross once", fontsize=12); a.legend(frameon=False, fontsize=9.5, loc="upper left"); despine(a)
    r = np.linspace(7540, 7680, 300)
    b.plot(r, geo(r), color=MAROON, lw=2.6, label="geometry")
    b.plot(r, lap(r), color=BLUE, lw=2.6, label=f"Laplace  (root {dl['r']:.1f} km)")
    b.plot(r, gau(r), color=GREEN, lw=2.6, label=f"Gauss, first pass  (root {dg['roots'][-1]:.1f} km)")
    b.axvline(7600.0, color=INK, lw=1.2, ls="--"); b.text(7601, 1452, "truth $7600$ km", fontsize=9.5, va="bottom")
    b.set_xlabel("$r$ [km]"); b.set_ylabel(r"$\rho$ [km]"); b.set_title("Zoom: the dynamics decide where they cross", fontsize=12)
    b.set_ylim(1240, 1475); b.legend(frameon=False, fontsize=9.5, loc="lower right"); despine(b)
    fig.tight_layout()
    save(fig, "fig-od-laplace")
    return dl, dg


# ------------------------------------------------------------------------------------------ Gauss iteration
def fig_gauss():
    fig, (a, b) = plt.subplots(1, 2, figsize=(12.6, 4.9))
    cols = {10: PURPLE, 60: BLUE, 120: GREEN, 240: AMBER}
    for h, c in cols.items():
        ob = observe([-h, 0, h])
        L = [o["u"] for o in ob]; R = [o["R"] for o in ob]
        _, _, d = gauss(L, R, [-h, 0, h], details=True)
        e = np.array([np.linalg.norm(x[1] - r0) for x in d["hist"]])
        a.semilogy(np.arange(len(e)), np.maximum(e, 1e-12), "o-", color=c, ms=3.5, lw=1.8, label=f"arc $\\pm{h}$ s")
    k = np.arange(0, 40)
    a.semilogy(k, 1.24*0.5**k, ":", color=INK, lw=1.3)
    a.text(1, 2e-10, "dotted: halving at every iteration", fontsize=9.5)
    a.set_ylim(1e-11, 1e2); a.set_xlabel("iteration (0 = first pass with the series for $F$, $G$)"); a.set_ylabel(r"position error $|\mathbf{r}_2-\mathbf{r}_{\mathrm{true}}|$ [km]")
    a.set_title("Exact data: the error halves each iteration", fontsize=12); a.legend(frameon=False, fontsize=9.5); despine(a)
    ob = observe([-60, 0, 60])
    _, _, d = gauss([o["u"] for o in ob], [o["R"] for o in ob], [-60, 0, 60], details=True)
    rho_err = np.array([x[0][1] - ob[1]["range"] for x in d["hist"]])[:30]
    it = np.arange(len(rho_err))
    pos = rho_err > 0
    b.semilogy(it, np.abs(rho_err), "-", color=GREY, lw=1.2)
    b.semilogy(it[pos], np.abs(rho_err[pos]), "^", color=MAROON, ms=7, label=r"$\rho_2$ too large")
    b.semilogy(it[~pos], np.abs(rho_err[~pos]), "v", color=BLUE, ms=7, label=r"$\rho_2$ too small")
    b.set_xlabel("iteration"); b.set_ylabel(r"range error $|\rho_2-\rho_{2,\mathrm{true}}|$ [km]")
    b.set_title("The range overshoots and undershoots alternately", fontsize=12); b.legend(frameon=False, fontsize=10); despine(b)
    fig.tight_layout()
    save(fig, "fig-od-gauss")


# ------------------------------------------------------------------------------------------ noise against arc
def fig_noise():
    NA = noise_vs_arc()
    h = np.array([r["h"] for r in NA], float)
    fig, (a, b) = plt.subplots(1, 2, figsize=(12.6, 4.9))
    for key, c, lab in (("laplace", MAROON, "Laplace"), ("gauss1", AMBER, "Gauss, first pass"), ("gauss", BLUE, "Gauss, iterated")):
        q = np.array([r[key] for r in NA])
        a.loglog(h, q[:, 1], "o-", color=c, lw=2.6, ms=5, label=lab)
        a.fill_between(h, q[:, 0], q[:, 2], color=c, alpha=0.18, lw=0)
    # the same without noise
    hs = np.array([10, 20, 40, 60, 120, 240, 480])
    e_l, e_g1 = [], []
    for hh in hs:
        ob = observe([-hh, 0, hh])
        L = [o["u"] for o in ob]; R = [o["R"] for o in ob]
        e_l.append(np.linalg.norm(laplace(L, R, ob[1]["Rd"], ob[1]["Rdd"], [-hh, 0, hh])[0] - r0))
        e_g1.append(np.linalg.norm(gauss(L, R, [-hh, 0, hh], iterate=False)[0] - r0))
    a.loglog(hs, e_l, "--", color=MAROON, lw=1.3); a.loglog(hs, e_g1, "--", color=AMBER, lw=1.3)
    a.text(28, 0.02, "dashed: the same methods\nwithout noise", fontsize=9.5, color=GREY, ha="left")
    a.set_xlabel(r"half-width $h$ of the arc ($t=-h,0,+h$) [s]"); a.set_ylabel(r"position error $|\mathbf{r}_2-\mathbf{r}_{\mathrm{true}}|$ [km]")
    a.set_title("2″ noise: median (line) and 16–84 % (band), 300 trials", fontsize=12); a.legend(frameon=False, fontsize=9.5, loc="upper left"); despine(a)
    dev = np.array([r["dev"] for r in NA])
    b.loglog(h, dev, "o-", color=PURPLE, lw=2.6, ms=5, label="out-of-plane deviation of the middle line of sight")
    b.axhline(2.0, color=INK, lw=1.2, ls="--"); b.text(11, 2.4, "measurement noise, 2″", fontsize=10)
    b.set_xlabel(r"half-width $h$ of the arc [s]"); b.set_ylabel("arcseconds")
    b.set_title("The signal that carries the range", fontsize=12); b.legend(frameon=False, fontsize=9.5, loc="upper left"); despine(b)
    fig.tight_layout()
    save(fig, "fig-od-noise")


# ------------------------------------------------------------------------------------------ least squares
def fig_lsq():
    LS = lsq_experiment()
    demo = LS["demo"]
    hist = demo["hist"]
    fig = plt.figure(figsize=(14.0, 4.6))
    a = fig.add_subplot(1, 3, 1)
    it = np.arange(len(hist))
    rms = np.array([x[1] for x in hist])/ARCSEC
    err = np.array([np.linalg.norm(x[0][:3] - r0) for x in hist])
    a.semilogy(it, rms, "o-", color=BLUE, lw=2.4, ms=5)
    a.axhline(2.0, color=INK, lw=1.1, ls="--"); a.text(len(hist) - 1.1, 2.3, "noise 2″", ha="right", fontsize=9.5)
    a.set_xlabel("iteration"); a.set_ylabel("RMS residual [arcsec]", color=BLUE)
    a2 = a.twinx(); a2.semilogy(it, err, "s-", color=MAROON, lw=2.2, ms=5); a2.set_ylabel("position error [km]", color=MAROON)
    a.set_title("Differential correction", fontsize=12); despine(a)
    b = fig.add_subplot(1, 3, 2)
    los = LS["los"]*1e3
    b.scatter(los[:, 0], los[:, 1], s=9, color=BLUE, alpha=0.55)
    P = demo["B"] @ demo["P"][:3, :3] @ demo["B"].T*1e6
    for k, ls_ in ((1, "-"), (3, "--")):
        b.add_patch(Ellipse((0, 0), 2*k*np.sqrt(P[0, 0]), 2*k*np.sqrt(P[1, 1]), fill=False, color=MAROON, lw=1.6, ls=ls_))
    b.text(0.97, 0.93, "ellipses: formal 1σ, 3σ", transform=b.transAxes, ha="right", va="top", color=MAROON, fontsize=9.5)
    b.set_xlim(-900, 900); b.set_ylim(-45, 45)
    b.set_xlabel("error along the line of sight [m]"); b.set_ylabel("error across [m]")
    b.set_title("200 repeats: a needle (note the scales)", fontsize=12); despine(b)
    c = fig.add_subplot(1, 3, 3)
    bins = np.logspace(-2, 1, 28)
    c.hist(LS["iod"], bins=bins, color=AMBER, alpha=0.75, label=f"3 observations, Gauss (median {np.median(LS['iod']):.2f} km)")
    c.hist(LS["lsq"], bins=bins, color=BLUE, alpha=0.75, label=f"all 49, least squares (median {np.median(LS['lsq']):.2f} km)")
    c.set_xscale("log"); c.set_xlabel(r"position error [km]"); c.set_ylabel("trials")
    c.set_ylim(0, c.get_ylim()[1]*1.3)
    c.set_title("Same data, better processing", fontsize=12); c.legend(frameon=False, fontsize=8.5, loc="upper left"); despine(c)
    fig.tight_layout(w_pad=3.0)
    save(fig, "fig-od-lsq")


if __name__ == "__main__":
    fig_scenario()
    fig_gibbs()
    fig_hg()
    dl, dg = fig_laplace()
    fig_gauss()
    fig_noise()
    fig_lsq()