"""Figures for universal-variables.qmd (img/fig-uv-*.png).

  fig-uv-stumpff     the Stumpff functions C(z), S(z), smooth through z = 0
  fig-uv-chi         r and t as functions of the universal anomaly, for an ellipse, a parabola and a hyperbola
  fig-uv-newton      the universal Kepler equation of Example 1 and Newton's iterates on it
  fig-uv-precision   position error of the classical and the universal propagator near e = 1
  fig-uv-lambert     Lambert's problem in universal variables: y(z) and the time of flight t(z)

Run from the sfm-slides directory:   python3 codes/fig-uv.py
"""
import sys
import numpy as np
import matplotlib.pyplot as plt
import mpmath as mp
sys.path.insert(0, "codes")
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _uv import stumpff, kepler_uv, propagate, lambert_uv, lambert_geometry, lambert_y, lambert_tof, coe, MU_E

mu = MU_E


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


# ------------------------------------------------------------------ Stumpff functions
z = np.linspace(-30, 60, 2000)
CS = np.array([stumpff(v) for v in z])
fig, axs = plt.subplots(1, 2, figsize=(12.5, 4.4))
for ax, k, name in ((axs[0], 0, r"$C(z)=\dfrac{1-\cos\sqrt{z}}{z}$"), (axs[1], 1, r"$S(z)=\dfrac{\sqrt{z}-\sin\sqrt{z}}{\sqrt{z}^{\,3}}$")):
    ax.axvspan(-30, 0, color=MAROON, alpha=0.06)
    ax.axvspan(0, 60, color=BLUE, alpha=0.05)
    ax.plot(z, CS[:, k], color=INK, lw=2)
    ax.axvline(0, color=INK, lw=0.7, ls=":")
    ax.axvline(4 * np.pi ** 2, color=GREEN, lw=1, ls="--")
    ax.scatter([0], [0.5 if k == 0 else 1 / 6], color=AMBER, zorder=4, s=40)
    ax.text(0.8, (0.5 if k == 0 else 1 / 6) * 1.12, ["$C(0)=1/2$", "$S(0)=1/6$"][k], color=AMBER, fontsize=11)
    ytop = 1.5 if k == 0 else 0.34
    ax.text(-15, ytop, "hyperbola\n$z<0$", ha="center", color=MAROON, fontsize=10, va="top")
    ax.text(20, ytop, "ellipse\n$z>0$", ha="center", color=BLUE, fontsize=10, va="top")
    ax.text(4 * np.pi ** 2 + 0.8, 0.02 if k == 0 else 0.05, r"$z=4\pi^2$" + "\n(one full turn)", color=GREEN, fontsize=9)
    ax.set_xlabel("$z=\\alpha\\chi^2$")
    ax.set_title(name, fontsize=13)
    ax.set_xlim(-30, 60)
    despine(ax)
axs[0].set_ylim(0, 1.6)
axs[1].set_ylim(0, 0.36)
save(fig, "fig-uv-stumpff")

# ------------------------------------------------------------------ r(chi), t(chi) for three conics
rp = 7000.0
fig, (aR, aT) = plt.subplots(1, 2, figsize=(12.5, 4.6))
chi = np.linspace(-260, 260, 800)
for e, col, lab in ((0.5, BLUE, "ellipse, $e=0.5$"), (1.0, GREEN, "parabola, $e=1$"), (1.5, MAROON, "hyperbola, $e=1.5$")):
    alpha = (1 - e) / rp                                   # 1/a, from r_p = a(1 - e)
    zz = alpha * chi ** 2
    C = np.array([stumpff(v)[0] for v in zz])
    S = np.array([stumpff(v)[1] for v in zz])
    r = chi ** 2 * C + rp * (1 - zz * C)                    # periapsis start: sigma_0 = 0
    t = (chi ** 3 * S + rp * chi * (1 - zz * S)) / np.sqrt(mu)
    aR.plot(chi, r / 1000, color=col, lw=2, label=lab)
    aT.plot(chi, t / 3600, color=col, lw=2, label=lab)
aR.axhline(rp / 1000, color=INK, lw=0.7, ls=":")
aR.text(-255, rp / 1000 + 1.2, "$r_p$", fontsize=11)
aR.text(0, 44, r"ellipse: $r=a+(r_p-a)\cos(\chi/\sqrt{a})$" + "\noscillates about $1/\\alpha=a$", color=BLUE, fontsize=9, ha="center")
aR.set_ylim(0, 78)
aR.set_xlabel(r"universal anomaly $\chi$ [km$^{1/2}$]")
aR.set_ylabel("$r$ [1000 km]")
aR.set_title(r"$d^2r/d\chi^2+\alpha r=1$: a linear equation for every conic", fontsize=12)
aR.legend(frameon=False, fontsize=10, loc="upper center")
aT.set_xlabel(r"universal anomaly $\chi$ [km$^{1/2}$]")
aT.set_ylabel("time from periapsis [h]")
aT.set_title(r"$\sqrt{\mu}\,dt=r\,d\chi$: time always increases with $\chi$", fontsize=12)
aT.axhline(0, color=INK, lw=0.6, ls=":")
aT.legend(frameon=False, fontsize=10, loc="upper left")
for ax in (aR, aT):
    despine(ax)
save(fig, "fig-uv-chi")

# ------------------------------------------------------------------ Newton on the universal Kepler equation, Example 1
r0 = np.array([-2757.4, -6465.5, -1726.8])
v0 = np.array([7.4488, -4.3468, -4.4076])
dt = 3 * 3600.0
nr0 = np.linalg.norm(r0)
vr0 = r0 @ v0 / nr0
alpha = 2 / nr0 - v0 @ v0 / mu
sm = np.sqrt(mu)


def ukep(x):
    zz = alpha * x * x
    C, S = stumpff(zz)
    return (nr0 * vr0 / sm * x * x * C + (1 - alpha * nr0) * x ** 3 * S + nr0 * x) / sm


_, hist = kepler_uv(nr0, vr0, alpha, dt, history=True)
xs = np.linspace(0, 420, 600)
fig, ax = plt.subplots(figsize=(8.4, 5.0))
ax.plot(xs, [ukep(x) / 3600 for x in xs], color=INK, lw=2, label=r"$t(\chi)$, the universal Kepler equation")
ax.axhline(dt / 3600, color=AMBER, lw=1.2, ls="--", label=r"$\Delta t=3$ h")
for k in range(len(hist) - 1):
    x0, x1 = hist[k], hist[k + 1]
    ax.plot([x0, x0], [dt / 3600, ukep(x0) / 3600], color=PURPLE, lw=0.9, ls=":")
    ax.plot([x0, x1], [ukep(x0) / 3600, dt / 3600], color=PURPLE, lw=1.3)
    ax.scatter([x0], [ukep(x0) / 3600], color=PURPLE, s=22, zorder=4)
ax.text(hist[0] - 8, ukep(hist[0]) / 3600 - 0.35, r"$\chi_0=\sqrt{\mu}\,|\alpha|\,\Delta t$", color=PURPLE, fontsize=10, ha="right")
ax.scatter([hist[-1]], [dt / 3600], color=MAROON, s=50, zorder=5)
ax.text(hist[-1] + 6, dt / 3600 - 0.35, rf"$\chi={hist[-1]:.2f}$" + r" km$^{1/2}$", color=MAROON, fontsize=10)
ax.text(40, 3.9, r"slope $dt/d\chi=r/\sqrt{\mu}>0$" + "\n(Newton's step is the residual over $r$)", fontsize=10)
ax.set_xlabel(r"$\chi$ [km$^{1/2}$]")
ax.set_ylabel("time [h]")
ax.set_xlim(0, 420)
ax.set_ylim(0, 5)
ax.legend(frameon=False, fontsize=10, loc="lower right")
ax.set_title("Example 1: Newton's method on the universal Kepler equation", fontsize=12)
despine(ax)
save(fig, "fig-uv-newton")

# ------------------------------------------------------------------ precision near e = 1
mp.mp.dps = 50


def uv_mp(r0, v0, dt):
    """Universal-variable propagation in 50-digit arithmetic: the reference."""
    r0 = [mp.mpf(x) for x in r0]
    v0 = [mp.mpf(x) for x in v0]
    m, dt = mp.mpf(mu), mp.mpf(dt)
    nr0 = mp.sqrt(sum(x * x for x in r0))
    vr0 = sum(a * b for a, b in zip(r0, v0)) / nr0
    al = 2 / nr0 - sum(x * x for x in v0) / m
    sm = mp.sqrt(m)

    def cs(z):
        if z > 0:
            s = mp.sqrt(z)
            return (1 - mp.cos(s)) / z, (s - mp.sin(s)) / s ** 3
        if z < 0:
            s = mp.sqrt(-z)
            return (mp.cosh(s) - 1) / (-z), (mp.sinh(s) - s) / s ** 3
        return mp.mpf(1) / 2, mp.mpf(1) / 6
    x = sm * dt / nr0
    for _ in range(200):
        zz = al * x * x
        C, S = cs(zz)
        res = nr0 * vr0 / sm * x * x * C + (1 - al * nr0) * x ** 3 * S + nr0 * x - sm * dt
        rr = nr0 * vr0 / sm * x * (1 - zz * S) + (1 - al * nr0) * x * x * C + nr0
        d = res / rr
        x -= d
        if abs(d) < mp.mpf(10) ** -40:
            break
    zz = al * x * x
    C, S = cs(zz)
    F, G = 1 - x * x * C / nr0, dt - x ** 3 * S / sm
    return np.array([float(F * a + G * b) for a, b in zip(r0, v0)])


def classical(r0, v0, dt):
    """The textbook route: elements, mean anomaly, Kepler's equation (elliptic or hyperbolic), back to r."""
    a, e, _, _, _, f = coe(r0, v0)
    f = np.radians(f)
    if e < 1:
        E0 = 2 * np.arctan(np.sqrt((1 - e) / (1 + e)) * np.tan(f / 2))
        M = E0 - e * np.sin(E0) + np.sqrt(mu / a ** 3) * dt
        E = M
        for _ in range(200):
            d = (E - e * np.sin(E) - M) / (1 - e * np.cos(E))
            E -= d
            if abs(d) < 1e-15:
                break
        fn = 2 * np.arctan(np.sqrt((1 + e) / (1 - e)) * np.tan(E / 2))
    else:
        H0 = 2 * np.arctanh(np.sqrt((e - 1) / (e + 1)) * np.tan(f / 2))
        M = e * np.sinh(H0) - H0 + np.sqrt(mu / (-a) ** 3) * dt
        H = np.arcsinh(M / e)
        for _ in range(200):
            d = (e * np.sinh(H) - H - M) / (e * np.cosh(H) - 1)
            H -= d
            if abs(d) < 1e-15:
                break
        fn = 2 * np.arctan(np.sqrt((e + 1) / (e - 1)) * np.tanh(H / 2))
    p = a * (1 - e * e)
    rr = p / (1 + e * np.cos(fn))
    h = np.cross(r0, v0)
    ur = r0 / np.linalg.norm(r0)
    ut = np.cross(h / np.linalg.norm(h), ur)
    return rr * (np.cos(fn - f) * ur + np.sin(fn - f) * ut)


def precision_case(e, rp=7000.0):
    """Start 10 min before periapsis, step 20 min across it; return (uv error, classical error) in km."""
    vp = np.sqrt(mu * (1 + e) / rp)
    rP, vP = np.array([rp, 0, 0.0]), np.array([0, vp, 0.0])
    r0 = uv_mp(rP, vP, -600)
    _, v0 = propagate(rP, vP, -600)
    ref = uv_mp(r0, v0, 1200)
    return np.linalg.norm(propagate(r0, v0, 1200)[0] - ref), np.linalg.norm(classical(r0, v0, 1200) - ref)


if __name__ == "__main__":
    gaps = np.logspace(-9, -1, 17)
    rows = [(g, *precision_case(1 - g)) for g in gaps] + [(g, *precision_case(1 + g)) for g in gaps]
    rows = np.maximum(np.array(rows), 1e-13)               # an exact zero cannot be drawn on a log axis
    fig, ax = plt.subplots(figsize=(8.4, 5.0))
    n = len(gaps)
    ax.loglog(rows[:n, 0], rows[:n, 2] * 1e3, "o-", color=MAROON, label="classical, ellipse ($e<1$, eccentric anomaly)")
    ax.loglog(rows[n:, 0], rows[n:, 2] * 1e3, "s--", color=MAROON, mfc="none", label="classical, hyperbola ($e>1$, hyperbolic anomaly)")
    ax.loglog(rows[:n, 0], rows[:n, 1] * 1e3, "o-", color=BLUE, label="universal variables, $e<1$")
    ax.loglog(rows[n:, 0], rows[n:, 1] * 1e3, "s--", color=BLUE, mfc="none", label="universal variables, $e>1$")
    ax.set_xlabel(r"distance from the parabola, $|1-e|$")
    ax.set_ylabel("position error after 20 min [m]")
    ax.set_title("Propagating across periapsis, $r_p=7000$ km (reference: 50-digit arithmetic)", fontsize=11)
    ax.legend(frameon=False, fontsize=9, loc="upper right")
    ax.set_ylim(1e-10, 1e2)
    despine(ax)
    save(fig, "fig-uv-precision")

    # -------------------------------------------------------------- Lambert in universal variables
    r1 = np.array([6678.137, 0, 0.0])
    th = np.radians(175)
    r2 = 42164 * np.array([np.cos(th), np.sin(th), 0.0])
    dth, A, n1, n2 = lambert_geometry(r1, r2)
    zz = np.linspace(-80, 4 * np.pi ** 2 - 0.05, 4000)
    yv = np.array([lambert_y(v, n1, n2, A) for v in zz])
    tv = np.array([lambert_tof(v, n1, n2, A) for v in zz])
    zneg = zz[yv < 0].max()
    fig, (aY, aT) = plt.subplots(1, 2, figsize=(12.5, 4.6))
    aY.plot(zz, yv / 1000, color=INK, lw=2)
    aY.axhline(0, color=INK, lw=0.6)
    aY.axvspan(-80, zneg, color=MAROON, alpha=0.08)
    aY.text(-78, 60, "$y<0$:\nno orbit", color=MAROON, fontsize=10)
    aY.set_xlabel("$z$")
    aY.set_ylabel("$y(z)$ [1000 km]")
    aY.set_title(r"$y(z)=r_1+r_2+A\,\dfrac{zS(z)-1}{\sqrt{C(z)}}$", fontsize=12)
    aY.set_ylim(-100, 160)
    aT.semilogy(zz[yv > 0], tv[yv > 0] / 3600, color=INK, lw=2)
    aT.axvspan(-80, zneg, color=MAROON, alpha=0.08)
    aT.axvline(0, color=GREEN, lw=1, ls=":")
    aT.text(1.0, 60, "parabola\n$z=0$", color=GREEN, fontsize=9)
    aT.axvline(4 * np.pi ** 2, color=GREEN, lw=1, ls="--")
    aT.text(4 * np.pi ** 2 - 0.4, 0.15, r"$z\to4\pi^2$: $\Delta t\to\infty$", color=GREEN, fontsize=9, ha="right")
    for T, col, lab in ((4, BLUE, "Example 3: 4 h, ellipse"), (2, MAROON, "2 h, hyperbola")):
        sol = lambert_uv(r1, r2, T * 3600)
        aT.axhline(T, color=col, lw=1, ls="--")
        aT.scatter([sol["z"]], [T], color=col, s=45, zorder=4)
        aT.text(-62, T * 1.12, lab + rf", $z={sol['z']:.3f}$", color=col, fontsize=9)
    aT.set_xlabel("$z$")
    aT.set_ylabel(r"$\Delta t(z)$ [h]")
    aT.set_title(r"time of flight: increasing in $z$, one root per $\Delta t$", fontsize=12)
    aT.set_ylim(0.02, 300)
    for ax in (aY, aT):
        ax.set_xlim(-80, 4 * np.pi ** 2)
        despine(ax)
    save(fig, "fig-uv-lambert")
    print("wrote five figures")