#!/usr/bin/env python3
"""
orbital-elements-3d.py
======================
From the in-plane orbit (a, e, f) to the six classical orbital elements.

The two-body problem fixes the *shape* of the orbit in its own plane:

    r(f) = p / (1 + e cos f),        p = a(1 - e^2),

which needs only (a, e) for the conic and f to say where on it the body is.
Three more numbers are needed to say how that plane is *oriented* with respect
to an inertial reference frame (X = vernal equinox, XY = equatorial/reference
plane, Z = reference normal):

    Omega   right ascension of the ascending node -- swings the line of nodes
            in the reference plane, measured from X about Z.
    i       inclination -- tilts the orbital plane about the line of nodes.
    omega   argument of periapsis -- rotates the apse line inside the orbital
            plane, measured from the ascending node about h.

The transformation from the perifocal (PQW) frame to the inertial (IJK) frame
is the 3-1-3 Euler sequence (Battin Sec. 3.3, Vallado Eq. 2-74)

    Q = R3(-Omega) R1(-i) R3(-omega),      r_IJK = Q r_PQW.

Run it and drag the sliders: each of Omega, i, omega moves the orbit rigidly,
none of them changes its size or shape.

Usage
-----
    python codes/orbital-elements-3d.py                    # interactive sliders
    python codes/orbital-elements-3d.py --verify           # element round-trip
    python codes/orbital-elements-3d.py --png img/fig-orbital-elements-3d.png
    python codes/orbital-elements-3d.py --save img/orbital-elements-3d.gif

Non-dimensional units throughout: mu = 1, distances in units of the
reference-plane grid (nothing depends on the scale).
"""

from __future__ import annotations

import argparse
import os

import numpy as np

MU = 1.0

# ── palette (shared with the other figures in this deck) ─────────────────────
CREAM, INK, DIM = "#fcfbf7", "#2f2a24", "#8b857c"
MAROON, BLUE, GREEN, AMBER, PURPLE = ("#6b1f1f", "#2f5ea8", "#2d6a2d",
                                      "#b45309", "#7e22ce")
# the three reference-frame axes are drawn in one neutral colour: they are the
# backdrop, not one of the elements, so X gets no more emphasis than Y or Z.
# Darker than DIM (which rules the reference plane and its grid) so the arrows
# read cleanly against it.
AXIS = "#6b655c"

# default elements (angles in degrees for everything the user touches)
DEFAULTS = dict(a=1.0, e=0.55, inc=35.0, Om=40.0, om=60.0, f=110.0)
A_MAX = 1.6          # top of the `a` slider; also the fixed frame size, so that
                     # dragging `a` grows the orbit instead of rescaling the view


# ── rotations ────────────────────────────────────────────────────────────────

def R1(t):
    """Rotation of the *frame* about the 1-axis by angle t (radians)."""
    c, s = np.cos(t), np.sin(t)
    return np.array([[1.0, 0.0, 0.0],
                     [0.0, c, s],
                     [0.0, -s, c]])


def R3(t):
    """Rotation of the *frame* about the 3-axis by angle t (radians)."""
    c, s = np.cos(t), np.sin(t)
    return np.array([[c, s, 0.0],
                     [-s, c, 0.0],
                     [0.0, 0.0, 1.0]])


def pqw_to_ijk(Om, inc, om):
    """Perifocal -> inertial direction cosine matrix, angles in radians."""
    return R3(-Om) @ R1(-inc) @ R3(-om)


def rodrigues(axis, ang):
    """Rotation matrix about a unit `axis` by `ang` (active rotation)."""
    k = np.asarray(axis, dtype=float)
    k = k / np.linalg.norm(k)
    K = np.array([[0.0, -k[2], k[1]],
                  [k[2], 0.0, -k[0]],
                  [-k[1], k[0], 0.0]])
    return np.eye(3) + np.sin(ang) * K + (1.0 - np.cos(ang)) * (K @ K)


def state_from_elements(a, e, inc, Om, om, f):
    """(a, e, i, Omega, omega, f) -> (r, v) in the inertial frame.

    Angles in radians.  Perifocal state (Battin Eq. 3.36):
        r_PQW = r (cos f, sin f, 0),
        v_PQW = sqrt(mu/p) (-sin f, e + cos f, 0).
    """
    p = a * (1.0 - e**2)
    r = p / (1.0 + e * np.cos(f))
    r_pqw = r * np.array([np.cos(f), np.sin(f), 0.0])
    v_pqw = np.sqrt(MU / p) * np.array([-np.sin(f), e + np.cos(f), 0.0])
    Q = pqw_to_ijk(Om, inc, om)
    return Q @ r_pqw, Q @ v_pqw


def elements_from_state(r_vec, v_vec):
    """Inverse map -- the textbook algorithm, used only by --verify."""
    r = np.linalg.norm(r_vec)
    h_vec = np.cross(r_vec, v_vec)
    h = np.linalg.norm(h_vec)
    n_vec = np.cross([0.0, 0.0, 1.0], h_vec)          # line of nodes
    n = np.linalg.norm(n_vec)
    e_vec = np.cross(v_vec, h_vec) / MU - r_vec / r
    e = np.linalg.norm(e_vec)
    eps = 0.5 * v_vec @ v_vec - MU / r
    a = -MU / (2.0 * eps)

    inc = np.arccos(h_vec[2] / h)
    if n < 1e-12:            # equatorial orbit: the node line is undefined, so
        Om = 0.0             # Omega is conventionally set to zero and omega
        om = np.arctan2(e_vec[1], e_vec[0]) % (2 * np.pi)   # measured from X
        if h_vec[2] < 0.0:                             # retrograde equatorial
            om = (2 * np.pi - om) % (2 * np.pi)
    else:
        Om = np.arctan2(n_vec[1], n_vec[0]) % (2 * np.pi)
        om = np.arccos(np.clip((n_vec @ e_vec) / (n * e), -1.0, 1.0))
        if e_vec[2] < 0.0:                             # periapsis below equator
            om = 2 * np.pi - om
    f = np.arccos(np.clip((e_vec @ r_vec) / (e * r), -1.0, 1.0))
    if r_vec @ v_vec < 0.0:                            # inbound half
        f = 2 * np.pi - f
    return dict(a=a, e=e, inc=inc, Om=Om, om=om, f=f)


# ── one frame of the picture ─────────────────────────────────────────────────

def _basis(nhat, prefer=None):
    """An orthonormal pair spanning the plane through the origin normal nhat."""
    nhat = np.asarray(nhat, dtype=float)
    u = np.array([1.0, 0.0, 0.0]) if prefer is None else np.asarray(prefer,
                                                                   float)
    if abs(u @ nhat) > 0.9:
        u = np.array([0.0, 1.0, 0.0])
    u = u - (u @ nhat) * nhat
    u /= np.linalg.norm(u)
    return u, np.cross(nhat, u)


def _square(nhat, half, prefer=None):
    """Corners of a square patch of the plane through the origin, normal nhat."""
    u, w = _basis(nhat, prefer)
    return np.array([half * (su * u + sw * w)
                     for su, sw in ((1, 1), (-1, 1), (-1, -1), (1, -1))])


def _grid(nhat, half, n=4, prefer=None):
    """Line segments of a square grid ruled on that plane (for depth cues)."""
    u, w = _basis(nhat, prefer)
    segs = []
    for t in np.linspace(-half, half, n + 1):
        segs.append(np.array([t * u - half * w, t * u + half * w]))
        segs.append(np.array([-half * u + t * w, half * u + t * w]))
    return segs


def _arc(v_from, axis, ang, radius, n=80):
    """Points of an arc: `v_from` swept about `axis` through `ang`, scaled."""
    v = np.asarray(v_from, dtype=float)
    v = v / np.linalg.norm(v)
    ts = np.linspace(0.0, ang, max(n, 2))
    return np.array([radius * (rodrigues(axis, t) @ v) for t in ts])


def draw_scene(ax, prm, highlight=None, scale=None):
    """Redraw the whole 3D scene for the element set `prm` (degrees in prm).

    Every length in the scene -- axis limits, planes, arrows, arc radii -- is a
    multiple of the half-size L.  Tie L to `a` and the figure is scale-invariant:
    the orbit and the frame grow together and the picture never changes.  Pass
    `scale` to pin L to a fixed semi-major axis so that varying `a` is visible.
    """
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection

    elev, azim = ax.elev, ax.azim
    ax.clear()

    a, e = prm["a"], prm["e"]
    inc, Om, om, f = (np.deg2rad(prm[k]) for k in ("inc", "Om", "om", "f"))

    p = a * (1.0 - e**2)
    Q = pqw_to_ijk(Om, inc, om)
    ehat = Q @ np.array([1.0, 0.0, 0.0])           # periapsis direction
    hhat = Q @ np.array([0.0, 0.0, 1.0])           # orbit normal
    nhat = np.array([np.cos(Om), np.sin(Om), 0.0])  # line of nodes
    zhat = np.array([0.0, 0.0, 1.0])

    L = (a if scale is None else scale) * (1.0 + e) * 1.15   # scene half-size
    hl = 1.0 if highlight is None else 0.35         # dimming of the rest

    def w(key, base=1.0):
        """Line-width/alpha weight: emphasise the element being varied."""
        return base if (highlight is None or highlight == key) else base * hl

    # ── the orbit, in the plane and in space ────────────────────────────────
    fs = np.linspace(0.0, 2 * np.pi, 720)
    rs = p / (1.0 + e * np.cos(fs))
    orb_pqw = np.stack([rs * np.cos(fs), rs * np.sin(fs), np.zeros_like(fs)], 1)
    orb = orb_pqw @ Q.T

    # ── reference (equatorial) plane and its axes ───────────────────────────
    sq = L * np.array([[1, 1, 0], [-1, 1, 0], [-1, -1, 0], [1, -1, 0]], float)
    ax.add_collection3d(Poly3DCollection([sq], facecolor=DIM, alpha=0.09,
                                         edgecolor=DIM, linewidths=0.8))
    for seg in _grid(zhat, L, n=6):
        ax.plot(*seg.T, color=DIM, lw=0.45, alpha=0.35)
    ax.text(-0.55 * L, -0.92 * L, 0.0, "reference plane", color=DIM,
            fontsize=9.5, ha="center")

    ax.quiver(0, 0, 0, *(1.02 * L * np.array([1.0, 0.0, 0.0])), color=AXIS,
              lw=1.2, arrow_length_ratio=0.10)
    ax.text(1.10 * L, 0.0, -0.04 * L, r"$\hat{X}$ (vernal equinox)",
            color=AXIS, fontsize=10, ha="center", va="top")
    ax.quiver(0, 0, 0, *(0.95 * L * np.array([0.0, 1.0, 0.0])), color=AXIS,
              lw=1.2, arrow_length_ratio=0.09)
    ax.text(*(1.02 * L * np.array([0.0, 1.0, 0.0])), r"$\hat{Y}$", color=AXIS,
            fontsize=10)
    ax.quiver(0, 0, 0, *(0.80 * L * zhat), color=AXIS, lw=1.2,
              arrow_length_ratio=0.11)
    ax.text(*(0.87 * L * zhat), r"$\hat{Z}$", color=AXIS, fontsize=10,
            ha="center")

    # ── orbital plane ───────────────────────────────────────────────────────
    quad = _square(hhat, L * 0.80, prefer=nhat)
    ax.add_collection3d(Poly3DCollection([quad], facecolor=BLUE,
                                         alpha=0.13 * w("inc", 1.0),
                                         edgecolor=BLUE, linewidths=0.9))
    for seg in _grid(hhat, L * 0.80, n=2, prefer=nhat):
        ax.plot(*seg.T, color=BLUE, lw=0.4, alpha=0.30 * w("inc", 1.0))
    ax.text(*(_square(hhat, L * 0.80, prefer=nhat)[1] * 0.99), "orbital plane",
            color=BLUE, fontsize=9.5, ha="center", alpha=w("inc"))
    ax.plot(*orb.T, color=INK, lw=2.6)
    shadow = orb.copy()
    shadow[:, 2] = 0.0                              # projection onto the equator
    ax.plot(*shadow.T, color=INK, lw=0.9, ls=(0, (3, 4)), alpha=0.45)

    # ── central body ────────────────────────────────────────────────────────
    ax.plot([0], [0], [0], "o", color=AMBER, ms=11, zorder=8)

    # ── line of nodes, ascending / descending node ──────────────────────────
    ax.plot(*np.vstack([-L * nhat, L * nhat]).T, color=GREEN,
            lw=2.0 * w("Om"), ls=(0, (7, 4)))
    r_asc = p / (1.0 + e * np.cos(-om)) * nhat
    r_dsc = -p / (1.0 - e * np.cos(-om)) * nhat
    ax.plot(*r_asc.reshape(3, 1), "o", color=GREEN, ms=7, zorder=8)
    ax.plot(*r_dsc.reshape(3, 1), "o", mfc=CREAM, mec=GREEN, mew=1.4, ms=6,
            zorder=8)
    ax.text(*(1.05 * L * nhat), "ascending\nnode", color=GREEN, fontsize=9.5,
            ha="center", va="top")

    # ── periapsis direction and angular momentum ────────────────────────────
    ax.quiver(0, 0, 0, *(0.9 * L * ehat), color=PURPLE, lw=2.2,
              arrow_length_ratio=0.10, alpha=w("om"))
    ax.text(*(0.97 * L * ehat), r"$\mathbf{e}$ (periapsis)", color=PURPLE,
            fontsize=10, alpha=w("om"))
    ax.quiver(0, 0, 0, *(0.78 * L * hhat), color=BLUE, lw=2.2,
              arrow_length_ratio=0.11, alpha=w("inc"))
    ax.text(*(0.86 * L * hhat), r"$\mathbf{h}$", color=BLUE, fontsize=12,
            alpha=w("inc"))

    # ── the body itself ─────────────────────────────────────────────────────
    r_vec, _ = state_from_elements(a, e, inc, Om, om, f)
    ax.quiver(0, 0, 0, *r_vec, color=MAROON, lw=2.0, arrow_length_ratio=0.09,
              alpha=w("f"))
    ax.plot(*r_vec.reshape(3, 1), "o", color=MAROON, ms=9, zorder=9)
    ax.plot(*np.vstack([r_vec, [r_vec[0], r_vec[1], 0.0]]).T, color=MAROON,
            lw=0.8, ls=":", alpha=0.6)
    ax.text(*(0.62 * r_vec + 0.11 * L * hhat), r"$\mathbf{r}$",
            color=MAROON, fontsize=12, ha="center", va="center", alpha=w("f"))

    # ── the four angles, drawn as arcs where they are actually measured ─────
    def arc(v_from, axis, ang, radius, color, label, key):
        pts = _arc(v_from, axis, ang, radius)
        ax.plot(*pts.T, color=color, lw=2.0 * w(key), alpha=w(key))
        mid = pts[len(pts) // 2] * 1.22
        ax.text(*mid, label, color=color, fontsize=13, ha="center",
                va="center", alpha=w(key))

    arc(np.array([1.0, 0.0, 0.0]), zhat, Om % (2 * np.pi), 0.42 * L, AMBER,
        r"$\Omega$", "Om")
    arc(zhat, nhat, inc, 0.46 * L, BLUE, r"$i$", "inc")
    arc(nhat, hhat, om % (2 * np.pi), 0.36 * L, PURPLE, r"$\omega$", "om")
    arc(ehat, hhat, f % (2 * np.pi), 0.22 * L, MAROON, r"$f$", "f")

    # ── cosmetics ───────────────────────────────────────────────────────────
    zlo, zhi = -0.92 * L, 0.96 * L
    ax.set_xlim(-L, L)
    ax.set_ylim(-L, L)
    ax.set_zlim(zlo, zhi)
    ax.set_box_aspect((1.0, 1.0, (zhi - zlo) / (2.0 * L)), zoom=1.32)
    ax.set_facecolor(CREAM)
    ax.axis("off")
    ax.view_init(elev=elev, azim=azim)

    title = (rf"$a = {a:.2f}$,   $e = {e:.2f}$,   "
             rf"$i = {prm['inc']:.0f}^\circ$,   "
             rf"$\Omega = {prm['Om']:.0f}^\circ$,   "
             rf"$\omega = {prm['om']:.0f}^\circ$,   "
             rf"$f = {prm['f']:.0f}^\circ$")
    ax.set_title(title, color=INK, fontsize=13, pad=-2)


# ── interactive front end ────────────────────────────────────────────────────

def interactive():
    import matplotlib.pyplot as plt
    from matplotlib.widgets import Button, Slider

    plt.rcParams.update({
        "figure.facecolor": CREAM, "savefig.facecolor": CREAM,
        "axes.facecolor": CREAM, "text.color": INK, "font.size": 11,
        "axes.labelcolor": INK, "xtick.color": INK, "ytick.color": INK,
    })

    fig = plt.figure(figsize=(10.6, 8.0))
    ax = fig.add_axes([0.09, 0.225, 0.82, 0.70], projection="3d")
    ax.view_init(elev=24, azim=22)

    prm = dict(DEFAULTS)

    # sliders: shape on the left, orientation on the right
    specs = [  # (key, label, lo, hi, x, y)
        ("a", "$a$", 0.5, A_MAX, 0.10, 0.170),
        ("e", "$e$", 0.0, 0.85, 0.10, 0.110),
        ("f", r"$f\;[^\circ]$", 0.0, 360.0, 0.10, 0.050),
        ("inc", r"$i\;[^\circ]$", 0.0, 180.0, 0.62, 0.170),
        ("Om", r"$\Omega\;[^\circ]$", 0.0, 360.0, 0.62, 0.110),
        ("om", r"$\omega\;[^\circ]$", 0.0, 360.0, 0.62, 0.050),
    ]
    colors = dict(a=INK, e=INK, f=MAROON, inc=BLUE, Om=AMBER, om=PURPLE)

    sliders = {}
    for key, label, lo, hi, x, y in specs:
        sax = fig.add_axes([x, y, 0.27, 0.030], facecolor="#eee9dd")
        sliders[key] = Slider(sax, label, lo, hi, valinit=prm[key],
                              color=colors[key],
                              valfmt="%4.2f" if key in ("a", "e") else "%5.1f")
        sliders[key].label.set_fontsize(13)
        sliders[key].label.set_color(colors[key])
        sliders[key].valtext.set_fontsize(11)

    fig.text(0.235, 0.213, "shape and position in the plane  ($a$, $e$, $f$)",
             ha="center", fontsize=11.5, color=INK)
    fig.text(0.755, 0.213,
             "orientation of the plane  ($i$, $\\Omega$, $\\omega$)",
             ha="center", fontsize=11.5, color=INK)

    def redraw(_=None):
        for k, s in sliders.items():
            prm[k] = s.val
        draw_scene(ax, prm, scale=A_MAX)      # frame fixed: `a` really zooms
        fig.canvas.draw_idle()

    for s in sliders.values():
        s.on_changed(redraw)

    # play / pause: advances f so the body runs around the orbit
    state = {"running": False}
    bax = fig.add_axes([0.445, 0.050, 0.078, 0.042])
    button = Button(bax, "Play", color="#eee9dd", hovercolor="#e2dbc9")
    rax = fig.add_axes([0.445, 0.110, 0.078, 0.042])
    reset = Button(rax, "Reset", color="#eee9dd", hovercolor="#e2dbc9")

    timer = fig.canvas.new_timer(interval=40)

    def tick():
        # equal steps in *time*, not in f: h = r^2 fdot keeps the body slow at
        # apoapsis and fast at periapsis, which is the whole point of Kepler II
        a_, e_ = prm["a"], prm["e"]
        p_ = a_ * (1.0 - e_**2)
        f_ = np.deg2rad(prm["f"])
        r_ = p_ / (1.0 + e_ * np.cos(f_))
        dt = 0.010 * a_**1.5
        df = np.sqrt(MU * p_) / r_**2 * dt
        sliders["f"].set_val(np.rad2deg((f_ + df) % (2 * np.pi)))

    timer.add_callback(tick)

    def toggle(_):
        state["running"] = not state["running"]
        button.label.set_text("Pause" if state["running"] else "Play")
        (timer.start if state["running"] else timer.stop)()
        fig.canvas.draw_idle()

    button.on_clicked(toggle)

    def do_reset(_):
        for k, s in sliders.items():
            s.reset()

    reset.on_clicked(do_reset)

    fig.text(0.5, 0.012,
             "drag with the mouse to rotate the view  ·  "
             r"$\Omega$, $i$, $\omega$ move the orbit rigidly: "
             "size and shape never change",
             ha="center", fontsize=10.5, color=DIM)

    redraw()
    plt.show()


# ── static frame and animation for the slides ────────────────────────────────

def save_png(path, **over):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update({"figure.facecolor": CREAM, "savefig.facecolor": CREAM,
                         "axes.facecolor": CREAM, "text.color": INK})
    prm = dict(DEFAULTS)
    prm.update(over)
    fig = plt.figure(figsize=(9.0, 7.4))
    ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], projection="3d")
    ax.view_init(elev=24, azim=22)
    draw_scene(ax, prm)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    fig.savefig(path, dpi=160, bbox_inches="tight", facecolor=CREAM)
    print("wrote", path)


def _schedule(n_per=70):
    """Frames of the tour: each orientation angle in turn, f always running.

    Every leg starts and ends at the default element set, so the loop closes and
    the eye can compare each sweep against the same reference orbit.
    """
    prm = dict(DEFAULTS)
    prm["f"] = 0.0
    legs = [
        ("Om", [40.0, 220.0, 400.0],
         r"varying $\Omega$ -- the line of nodes swings in the reference plane"),
        ("inc", [35.0, 150.0, 5.0, 35.0],
         "varying $i$ -- the plane tilts about the line of nodes"),
        ("om", [60.0, 240.0, 420.0],
         r"varying $\omega$ -- the apse line turns inside the orbital plane"),
    ]

    frames = []
    for key, waypoints, caption in legs:
        n_seg = max(2, n_per // (len(waypoints) - 1))
        for lo, hi in zip(waypoints[:-1], waypoints[1:]):
            for k in range(n_seg):
                t = 0.5 * (1.0 - np.cos(np.pi * k / n_seg))   # ease in/out
                prm[key] = lo + (hi - lo) * t
                prm["f"] = (prm["f"] + 360.0 / n_per) % 360.0
                frames.append((dict(prm), key, caption))
        prm[key] = DEFAULTS[key]
    return frames


def render(path, fps=20, dpi=100, n_per=70):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib import animation

    plt.rcParams.update({"figure.facecolor": CREAM, "savefig.facecolor": CREAM,
                         "axes.facecolor": CREAM, "text.color": INK})

    frames = _schedule(n_per=n_per)
    fig = plt.figure(figsize=(8.4, 7.2))
    ax = fig.add_axes([0.0, 0.03, 1.0, 0.92], projection="3d")
    ax.view_init(elev=24, azim=22)
    caption = fig.text(0.5, 0.025, "", ha="center", fontsize=13, color=INK)

    def frame(k):
        prm, key, text = frames[k]
        draw_scene(ax, prm, highlight=key)
        caption.set_text(text)
        return ()

    ani = animation.FuncAnimation(fig, frame, frames=len(frames), blit=False)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    writer = "ffmpeg" if path.endswith(".mp4") else "pillow"
    print(f"writing {path} ({len(frames)} frames) ...")
    ani.save(path, writer=writer, fps=fps, dpi=dpi,
             savefig_kwargs=dict(facecolor=CREAM))
    print("done:", path, f"{os.path.getsize(path)/1e6:.2f} MB")


# ── verification ─────────────────────────────────────────────────────────────

def verify():
    """Elements -> state -> elements, over a spread of test cases."""
    print("round trip  (a, e, i, Omega, omega, f) -> (r, v) -> elements\n")
    rng = np.random.default_rng(7)
    worst = 0.0
    print(f"{'a':>6}{'e':>7}{'i':>7}{'Om':>7}{'om':>7}{'f':>7}"
          f"{'max abs error':>16}")
    cases = [(1.0, 0.55, 35.0, 40.0, 60.0, 110.0),
             (1.3, 0.05, 98.0, 250.0, 300.0, 12.0),
             (0.7, 0.80, 5.0, 12.0, 190.0, 349.0)]
    cases += [tuple(v) for v in np.column_stack([
        rng.uniform(0.5, 1.6, 5), rng.uniform(0.02, 0.85, 5),
        rng.uniform(2.0, 178.0, 5), rng.uniform(0.0, 360.0, 5),
        rng.uniform(0.0, 360.0, 5), rng.uniform(0.0, 360.0, 5)])]

    for a, e, inc, Om, om, f in cases:
        ang = [np.deg2rad(x) for x in (inc, Om, om, f)]
        r_vec, v_vec = state_from_elements(a, e, *ang)
        el = elements_from_state(r_vec, v_vec)
        err = max(abs(el["a"] - a), abs(el["e"] - e),
                  *(abs((el[k] - t + np.pi) % (2 * np.pi) - np.pi)
                    for k, t in zip(("inc", "Om", "om", "f"), ang)))
        worst = max(worst, err)
        print(f"{a:6.2f}{e:7.3f}{inc:7.1f}{Om:7.1f}{om:7.1f}{f:7.1f}"
              f"{err:16.3e}")

    print(f"\nworst error over all cases: {worst:.3e}")

    # the three orientation angles must not touch the shape
    print("\ninvariance of the conic under the orientation angles\n")
    base = state_from_elements(1.0, 0.55, *map(np.deg2rad, (35.0, 40.0, 60.0,
                                                            110.0)))
    r0 = np.linalg.norm(base[0])
    v0 = np.linalg.norm(base[1])
    for inc, Om, om in [(0.0, 0.0, 60.0), (120.0, 210.0, 60.0),
                        (35.0, 40.0, 60.0), (88.0, 355.0, 60.0)]:
        r_vec, v_vec = state_from_elements(1.0, 0.55,
                                           *map(np.deg2rad, (inc, Om, om,
                                                             110.0)))
        el = elements_from_state(r_vec, v_vec)
        print(f"  i={inc:6.1f}  Om={Om:6.1f}  om={om:6.1f}   "
              f"|r| = {np.linalg.norm(r_vec):.10f}  "
              f"|v| = {np.linalg.norm(v_vec):.10f}  "
              f"a = {el['a']:.10f}  e = {el['e']:.10f}")
    print(f"\n  reference (same f):                  |r| = {r0:.10f}  "
          f"|v| = {v0:.10f}")


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[3])
    ap.add_argument("--verify", action="store_true",
                    help="element round-trip and shape-invariance checks")
    ap.add_argument("--png", metavar="PATH", default=None,
                    help="save a single static frame")
    ap.add_argument("--save", metavar="PATH", default=None,
                    help="render the Omega/i/omega tour to PATH (.gif or .mp4)")
    ap.add_argument("--fps", type=int, default=20)
    ap.add_argument("--dpi", type=int, default=100)
    ap.add_argument("--frames", type=int, default=70,
                    help="frames per leg of the tour (3 legs)")
    args = ap.parse_args()

    if args.verify:
        verify()
    if args.png:
        save_png(args.png)
    if args.save:
        render(args.save, fps=args.fps, dpi=args.dpi, n_per=args.frames)
    if not (args.verify or args.png or args.save):
        interactive()


if __name__ == "__main__":
    main()