import vaft.diagram

d = vaft.diagram.magnetic_island(m=3, n=2, width=0.16, phase=0.0, r_s=0.55, elongation=1.7, triangularity=0.4, projection='poloidal')
d.tikz  # the TikZ document as text
d.save("magnetic_island_poloidal.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("magnetic_island_poloidal.tex")  # the complete TikZ document, the text source
d.save("magnetic_island_poloidal.pdf")  # vector PDF; needs pdflatex

def peeling_ballooning(*, labels: bool = True) -> Diagram:
    r"""Schematic edge-stability (peeling-ballooning) diagram.

    Axes are the pedestal's peak normalised pressure gradient
    $\alpha_\mathrm{max}$ and peak bootstrap (edge) current density
    $J_{B,\mathrm{max}}$, in arbitrary units: the boundary's shape is
    illustrative, not computed from an equilibrium. Its corner, where the
    peeling and ballooning limits meet (the star), is where ELM-limited
    pedestals typically sit; it is computed from the model, not placed.
    """
    model = PeelingBallooningModel()
    chart = Chart(x_range=(0.0, 1.25), y_range=(0.0, 1.0))
    chart.curves["boundary"] = model.boundary()
    chart.points["corner"] = model.corner()
    chart.labels.update({"peeling": (0.4, 0.86), "stable": (0.3, 0.28), "ballooning": (1.0, 0.28)})
    chart.parameters.update({k: getattr(model, k) for k in ("j0", "j_slope", "a0", "a_slope", "width", "tau")})
    scene = _render_chart(
        chart,
        x_label="$\\alpha_\\mathrm{max}$",
        y_label="$J_{B,\\mathrm{max}}$",
        curve_styles={"boundary": "boundary"},
        region_text={"peeling": "Peeling unstable", "stable": "Stable",
                     "ballooning": "\\begin{tabular}{l}Ballooning\\\\unstable\\end{tabular}"} if labels else {},
        note="schematic: boundary shape illustrative (Connor et al. 1998; Snyder et al. 2002)" if labels else "",
        star="corner",
    )
    return Diagram("peeling_ballooning", scene, model=chart)

import vaft.diagram

d = vaft.diagram.peeling_ballooning()
d.tikz  # the TikZ document as text
d.save("peeling_ballooning.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("peeling_ballooning.tex")  # the complete TikZ document, the text source
d.save("peeling_ballooning.pdf")  # vector PDF; needs pdflatex

def s_alpha_ballooning(*, s_max: float = 1.5, alpha_max: float = 3.5, labels: bool = True) -> Diagram:
    r"""First and second ballooning stability in the circular $s$-$\alpha$ model.

    Both boundaries come from ``vaft.formula.s_alpha_marginal_alpha``; the
    dashed line is the $\alpha = 0.6\,s$ approximation
    ``ballooning_stability_criterion`` uses.
    """
    if not 0.1 < s_max <= 2.5 or not 0.5 < alpha_max <= 6.0:
        raise ValueError("s_max must lie in (0.1, 2.5] and alpha_max in (0.5, 6]")
    s, first, second = _s_alpha_boundaries(0.05, float(s_max), 36)
    ok = np.isfinite(first)
    chart = Chart(x_range=(0.0, float(alpha_max)), y_range=(0.0, float(s_max)))
    # two open branches: below s ~ 0.05 the model's boundaries are not resolved
    chart.curves["first"] = np.stack([first[ok], s[ok]], axis=-1)
    second_ok = ok & np.isfinite(second)
    chart.curves["second"] = np.stack([second[second_ok], s[second_ok]], axis=-1)
    chart.curves["approximation"] = np.array([[0.0, 0.0], [0.6 * s_max, s_max]])
    # labels are placed from the computed boundaries, so they stay in their
    # regions whatever the axis ranges; a region off the chart gets none
    def at(row, curve):
        return float(np.interp(row, s[ok], curve[ok]))

    s_first, s_unstable, s_second = 0.85 * s_max, 0.6 * s_max, max(0.25 * s_max, 0.1)
    a1 = at(s_first, first)
    chart.labels["first"] = (0.45 * a1, s_first)
    a1, a2 = at(s_unstable, first), at(s_unstable, np.where(np.isfinite(second), second, alpha_max))
    if a1 < 0.9 * alpha_max:
        chart.labels["unstable"] = (0.5 * (a1 + min(a2, alpha_max)), s_unstable)
    a2 = at(s_second, np.where(np.isfinite(second), second, np.inf))
    if a2 < 0.75 * alpha_max:
        chart.labels["second"] = (0.5 * (a2 + alpha_max), s_second)
    chart.parameters.update({"s_min": 0.05, "s_max": float(s_max), "alpha_max": float(alpha_max)})
    scene = _render_chart(
        chart,
        x_label="$\\alpha$",
        y_label="$s$",
        curve_styles={"approximation": "approx", "first": "boundary", "second": "boundary"},
        region_text={k: v for k, v in {
            "first": "\\begin{tabular}{c}First\\\\stable\\end{tabular}", "unstable": "Unstable",
            "second": "\\begin{tabular}{c}Second\\\\stable\\end{tabular}"}.items() if k in chart.labels} if labels else {},
        x_ticks=_nice_ticks(alpha_max + 1e-9), y_ticks=_nice_ticks(s_max + 1e-9),
        note="circular $s$-$\\alpha$ model (Connor, Hastie \\& Taylor 1978); dashed: $\\alpha = 0.6\\,s$" if labels else "",
    )
    return Diagram("s_alpha_ballooning", scene, model=chart)

import vaft.diagram

d = vaft.diagram.s_alpha_ballooning()
d.tikz  # the TikZ document as text
d.save("s_alpha_ballooning.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("s_alpha_ballooning.tex")  # the complete TikZ document, the text source
d.save("s_alpha_ballooning.pdf")  # vector PDF; needs pdflatex

def hugill(*, elongation: float = 1.0, q_limit: Optional[float] = None, labels: bool = True) -> Diagram:
    r"""Hugill diagram: $1/q_\mathrm{cyl}$ against the Murakami parameter $\bar n R/B$.

    The lines are the registered boundaries of the ``hugill`` projection
    (:mod:`vaft.formula.boundaries`, #1425): ``greenwald_hugill``, the
    Greenwald density in these coordinates, a straight line through the
    origin $1/q_\mathrm{cyl} = (\pi/50\kappa_a)\,\bar nR/B$ ($\bar n$ in
    $10^{19}\,\mathrm{m^{-3}}$) whose slope depends only on the area
    elongation; and ``murakami_hugill``, the vertical line $\bar nR/B = 1$.

    The registered current limit ``low_q`` is $q_\psi > 2$, the equilibrium
    edge q, not the $q_\mathrm{cyl}$ on this axis, so it is not drawn. Passing
    ``q_limit`` adds a horizontal *reference* line $q_\mathrm{cyl} = q_\mathrm{limit}$,
    labelled as such; it is not a registered boundary.
    """
    _validate_shape(_HUGILL_ASPECT_RATIO, elongation, q_limit)
    q_size = q_limit if q_limit is not None else _SIZING_Q
    y_max = 1.4 / q_size
    x_q = float(_b.boundary_value(_b.get_boundary("greenwald_hugill"),
                                  inverse_cylindrical_q=1.0 / q_size, area_elongation=elongation))
    x_max = 1.45 * x_q
    plan = overlay_plan("hugill", x_range=(0.0, x_max), y_range=(0.0, y_max),
                        fixed={"area_elongation": elongation})
    curves = {c.key: c for c in plan.curves}
    chart = Chart(x_range=(0.0, x_max), y_range=(0.0, y_max))
    chart.curves["greenwald"] = curves["greenwald_hugill"].xy
    chart.curves["murakami"] = curves["murakami_hugill"].xy
    # Murakami is a soft, Ohmic-era scaling: dashed, so the solid line stays the Greenwald limit
    styles = {"greenwald": "boundary", "murakami": "approx"}
    x_m = float(_b.boundary_value(_b.get_boundary("murakami_hugill")))
    chart.labels.update({
        "accessible": (0.45 * x_q, 0.85 / q_size),
        "density": (1.2 * x_q, 0.45 / q_size),
        "murakami": (x_m + 0.42 * x_q, 1.25 / q_size),
    })
    text = {"accessible": "\\begin{tabular}{c}Below\\\\Greenwald\\end{tabular}",
            "density": "\\begin{tabular}{c}Density limit\\\\($\\bar n_e > n_G$)\\end{tabular}",
            "murakami": "Murakami (Ohmic)"}
    if q_limit is not None:
        chart.curves["low_q"] = np.array([[0.0, 1.0 / q_limit], [x_max, 1.0 / q_limit]])
        styles["low_q"] = "approx"
        chart.labels["low_q"] = (0.33 * x_max, 1.2 / q_limit)
        text["low_q"] = f"Reference $q_\\mathrm{{cyl}} = {q_limit:g}$ (not $q_\\psi$)"
    chart.parameters.update({"elongation": elongation, "q_limit": q_limit, "murakami_at_q_limit": x_q,
                             "boundaries": plan.keys})
    scene = _render_chart(
        chart,
        x_label="$\\bar n_e R/B_T\\ [10^{19}\\,\\mathrm{m^{-2}\\,T^{-1}}]$",
        y_label="$1/q_\\mathrm{cyl}$",
        curve_styles=styles,
        region_text=text if labels else {},
        x_ticks=_nice_ticks(x_max),
        y_ticks=_nice_ticks(y_max),
        note=(f"Greenwald and Murakami limits in Murakami coordinates, $\\kappa_a = {elongation:g}$"
              if labels else ""),
    )
    return Diagram("hugill", scene, model=chart)

import vaft.diagram

d = vaft.diagram.hugill()
d.tikz  # the TikZ document as text
d.save("hugill.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("hugill.tex")  # the complete TikZ document, the text source
d.save("hugill.pdf")  # vector PDF; needs pdflatex

def hugill_st(*, elongation: float = 1.8, labels: bool = True) -> Diagram:
    r"""Spherical-tokamak Hugill diagram: $1/q^{ST}_\mathrm{cyl}$ against $\bar n R/B$ (Sykes et al. 2000, MAST).

    The vertical coordinate is the spherical-tokamak cylindrical safety factor
    $q^{ST}_\mathrm{cyl} = 2.5\,a^2(1+\kappa^2)B_T/(R I_p)$ of Sykes et al. (MAST,
    IAEA FEC 2000, Fig. 12), not the conventional $q_\mathrm{cyl}$ of
    :func:`hugill`. The lines are the registered boundaries of the
    ``hugill_st`` projection: ``sykes_2000_st_hugill``, the Hugill limit as
    Sykes et al. define it ($n_H = I_p/\pi a^2\kappa$), and
    ``greenwald_hugill_st``, the unchanged Greenwald limit on these axes, both
    lines through the origin whose slopes depend on $\kappa$ (MAST's figure
    uses 1.8); and ``murakami_hugill``, drawn dashed as a historical
    conventional-tokamak reference only.
    """
    _validate_shape(_HUGILL_ASPECT_RATIO, elongation, None)
    y_max = 1.4 / _SIZING_Q
    x_q = float(_b.boundary_value(_b.get_boundary("greenwald_hugill_st"),
                                  inverse_cylindrical_q_st=1.0 / _SIZING_Q, elongation=elongation))
    x_max = 1.25 * x_q
    plan = overlay_plan("hugill_st", x_range=(0.0, x_max), y_range=(0.0, y_max), fixed={"elongation": elongation})
    curves = {c.key: c for c in plan.curves}
    chart = Chart(x_range=(0.0, x_max), y_range=(0.0, y_max))
    chart.curves["greenwald"] = curves["greenwald_hugill_st"].xy
    chart.curves["hugill"] = curves["sykes_2000_st_hugill"].xy
    chart.curves["murakami"] = curves["murakami_hugill"].xy
    styles = {"greenwald": "boundary", "hugill": "boundary", "murakami": "approx"}
    x_m = float(_b.boundary_value(_b.get_boundary("murakami_hugill")))
    # each label in its own zone, clear of the three lines: below Hugill between Murakami and the Hugill line,
    # the Hugill name between the Hugill and Greenwald lines, Greenwald beyond its line, Murakami right of x = 1
    chart.labels.update({
        "accessible": (0.26 * x_q, 0.84 / _SIZING_Q),
        "hugill": (0.95 * x_q, 1.32 / _SIZING_Q),   # past the top end of the Hugill line, left of Greenwald
        "density": (0.80 * x_q, 0.30 / _SIZING_Q),
        "murakami": (x_m + 0.20 * x_q, 1.31 / _SIZING_Q),
    })
    text = {"accessible": "\\begin{tabular}{c}Below\\\\Hugill\\end{tabular}",
            "hugill": "Hugill",
            "density": "Beyond Greenwald",
            "murakami": "\\begin{tabular}{c}Murakami\\\\(reference)\\end{tabular}"}
    chart.parameters.update({"elongation": elongation, "boundaries": plan.keys})
    scene = _render_chart(
        chart,
        x_label="$\\bar n_e R/B_T\\ [10^{19}\\,\\mathrm{m^{-2}\\,T^{-1}}]$",
        y_label="$1/q^{ST}_\\mathrm{cyl}$",
        curve_styles=styles,
        region_text=text if labels else {},
        x_ticks=_nice_ticks(x_max),
        y_ticks=_nice_ticks(y_max),
        note=(f"Spherical-tokamak Hugill diagram (Sykes et al. 2000, MAST), $\\kappa = {elongation:g}$"
              if labels else ""),
    )
    return Diagram("hugill_st", scene, model=chart)

import vaft.diagram

d = vaft.diagram.hugill_st()
d.tikz  # the TikZ document as text
d.save("hugill_st.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("hugill_st.tex")  # the complete TikZ document, the text source
d.save("hugill_st.pdf")  # vector PDF; needs pdflatex

def troyon(*, beta_N_max: Optional[float] = None, aspect_ratio: float = 3.0, elongation: float = 1.7,
           q_limit: Optional[float] = None, labels: bool = True) -> Diagram:
    r"""Troyon diagram: toroidal beta against the normalised current $I_p/(aB_T)$.

    The registered ``troyon`` stability reference
    ($\beta_N = 2.2\,\mu_0 \cdot 10^6 \approx 2.76$, :mod:`vaft.formula.boundaries`) is
    drawn through the definition of $\beta_N$, $\beta_T[\%] = \beta_N\,I_p/(aB_T)$,
    with the projection's experimental references (#1691: DIII-D, ST, NSTX levels)
    dashed beside it. None is a limit, so no side is marked stable or unstable.
    ``beta_N_max`` overrides the Troyon value for a what-if picture, and the note
    then says so.

    No current limit is drawn by default: the registered ``low_q`` is on
    $q_\psi$, and this axis is $q_\mathrm{cyl}$-like. ``q_limit`` adds a
    *reference* cutoff where ``q_cyl_from_B_R_epsilon_kappa_I`` reaches it,
    at $I_p/(aB_T) = 5\varepsilon\kappa_a/q_\mathrm{limit}$.
    """
    _validate_shape(aspect_ratio, elongation, q_limit)
    if beta_N_max is not None and not beta_N_max > 0.0:
        raise ValueError(f"beta_N_max must be positive, not {beta_N_max!r}")
    registered = float(_b.boundary_value(_b.get_boundary("troyon")))
    level = registered if beta_N_max is None else float(beta_N_max)
    a = _R0 / aspect_ratio
    x_q = _current_at_q(elongation, aspect_ratio, q_limit if q_limit is not None else _SIZING_Q) / (a * _B0)
    x_max = 1.35 * x_q
    y_max = 1.3 * level * x_q
    chart = Chart(x_range=(0.0, float(x_max)), y_range=(0.0, float(y_max)))
    plan = overlay_plan("troyon", x_range=(0.0, x_max), y_range=(0.0, y_max), samples=101)
    if beta_N_max is None:
        chart.curves["beta_limit"] = next(c for c in plan.curves if c.key == "troyon").xy
    else:
        x = np.linspace(0.0, x_max, 101)
        chart.curves["beta_limit"] = np.stack([x, level * x], axis=-1)
    styles = {"beta_limit": "boundary"}
    # the Troyon name below its line, the experimental references' names where each meets the top of the frame
    chart.labels["beta"] = (0.95 * x_q, 0.45 * level * x_q)
    text = {"beta": f"Troyon $\\beta_N = {level:.3g}$"}
    for k, curve in enumerate(c for c in plan.curves if c.key != "troyon"):
        entry = _b.get_boundary(curve.key)
        name = f"ref{k}"
        chart.curves[name] = curve.xy
        styles[name] = "approx"
        # just right of its line, alternating heights so neighbouring steep lines keep their names apart
        y_label = (0.92 - 0.12 * (k % 2)) * y_max
        x_label = y_label / float(entry.coefficient) + 0.06 * x_max
        if x_label < 0.97 * x_max:
            chart.labels[name] = (x_label, y_label)
            text[name] = f"{entry.coefficient:.3g}"
    if q_limit is not None:
        chart.curves["low_q"] = np.array([[x_q, 0.0], [x_q, y_max]])
        styles["low_q"] = "approx"
        chart.labels["low_q"] = (1.18 * x_q, 0.45 * level * x_q)
        text["low_q"] = "\\begin{tabular}{c}Reference\\\\$q_\\mathrm{cyl}$\\end{tabular}"
    chart.parameters.update({"beta_N_max": level, "beta_N_registered": registered, "aspect_ratio": aspect_ratio,
                             "elongation": elongation, "q_limit": q_limit, "current_at_q_limit": x_q})
    source = "registered Troyon stability reference" if beta_N_max is None else "user value, not the registered one"
    levels = ", ".join(f"{_b.get_boundary(c.key).coefficient:.3g}" for c in plan.curves if c.key != "troyon")
    note = f"$\\beta_N = {level:.3g}$ ({source})"
    if levels:   # the values from the registry, never restated here
        note += f"; dashed: experimental $\\beta_N$ levels {levels} (references, not limits)"
    if q_limit is not None:
        note += (f"; $q_\\mathrm{{cyl}} = {q_limit:g}$ at $R_0/a = {aspect_ratio:g}$, "
                 f"$\\kappa_a = {elongation:g}$")
    scene = _render_chart(
        chart,
        x_label="$I_p/(aB_T)\\ [\\mathrm{MA\\,m^{-1}\\,T^{-1}}]$",
        y_label="$\\beta_T\\ [\\%]$",
        curve_styles=styles,
        region_text=text if labels else {},
        x_ticks=_nice_ticks(float(x_max)),
        y_ticks=_nice_ticks(float(y_max)),
        note=note if labels else "",
    )
    return Diagram("troyon", scene, model=chart)

import vaft.diagram

d = vaft.diagram.troyon()
d.tikz  # the TikZ document as text
d.save("troyon.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("troyon.tex")  # the complete TikZ document, the text source
d.save("troyon.pdf")  # vector PDF; needs pdflatex

def stability_diagnostic_taxonomy(*, labels: bool = True) -> Diagram:
    r"""Reduced MHD stability diagnostics by physical problem and logical status.

    Rows are the physical problem (current driven, pressure / curvature,
    resistive layer, axisymmetric, trigger / nonlinear, operational space),
    columns the logical status: exact definition, reduced model, empirical
    or semi-empirical boundary, heuristic or approximate, and solver-derived
    (DCON / RDCON / STRIDE quantities VAFT reads but does not compute). Each
    cell names what ``vaft.formula`` holds there; the gaps are listed. The
    reduced columns are compared with the solver column -- interpretation
    and limiting-case validation, not a screening gate.
    """
    labels = _check_labels(labels)
    items: List = []
    width = len(STATUSES) * (_CELL_W + 0.3)
    left = _X0 - 0.5 * _CELL_W - 3.0
    right = _X0 - 0.5 * _CELL_W + width
    for j, (key, text) in enumerate(STATUSES):
        x, _ = _cell_center(0, j)
        items.append(Label((x, _Y0 + 0.5 * _CELL_H + 0.2), text, "concept group title,text width=4.1cm,align=center",
                           anchor="south", role=f"status:{key}"))
    for i, (key, text) in enumerate(PROBLEMS):
        _, y = _cell_center(i, 0)
        items += band(left, right, y - 0.5 * _CELL_H - 0.12, y + 0.5 * _CELL_H + 0.12, role=f"problem:{key}")
        items.append(Label((left + 0.15, y), text, "concept band label,text width=2.6cm,align=left,execute at begin node={\\hyphenpenalty=10000}", anchor="west",
                           role=f"problem:{key}"))
    cells = {}
    for i, (problem, _) in enumerate(PROBLEMS):
        for j, (status, _) in enumerate(STATUSES):
            text = CELLS.get((problem, status))
            if text is None:
                continue
            x, y = _cell_center(i, j)
            style = "concept leaf" if status in ("heuristic", "solver") else "concept box"
            cells[(problem, status)] = box(x, y, _CELL_W, _CELL_H, text, style=style,
                                           role=f"cell:{problem}:{status}", latex=True)
            items += list(cells[(problem, status)].items)
    bottom = _cell_center(len(PROBLEMS) - 1, 0)[1] - 0.5 * _CELL_H - 0.4
    # the reduced and definition columns are compared with the solver column, not gated by it
    y_cmp = bottom - 0.45
    x_from = _cell_center(0, 1)[0]
    x_to = _cell_center(0, 4)[0]
    items.append(Arrow((x_from, y_cmp), (x_to, y_cmp), "connector both", role="comparison"))
    items.append(Label((0.5 * (x_from + x_to), y_cmp - 0.12),
                       "comparison and limiting-case validation where assumptions overlap -- not a screening gate",
                       "small label", anchor="north", role="comparison"))
    if labels:
        items.append(Label((0.5 * (left + right), y_cmp - 1.0),
                           "Not yet in VAFT: " + "; ".join(ABSENT), "note", anchor="north", role="note"))
        items.append(Label((0.5 * (left + right), y_cmp - 1.65),
                           "theory $\\to$ formula $\\to$ equilibrium-facing process $\\to$ plot: these cells are the "
                           "formula layer; evaluating them on an equilibrium is vaft.process, running solvers vaft.code",
                           "note", anchor="north", role="note"))
    model = {"statuses": tuple(k for k, _ in STATUSES), "problems": tuple(k for k, _ in PROBLEMS),
             "cells": dict(CELLS), "absent": ABSENT,
             "centers": {k: (b.x, b.y) for k, b in cells.items()}}
    return Diagram("stability_diagnostic_taxonomy", Scene(tuple(items)), model=model)

import vaft.diagram

d = vaft.diagram.stability_diagnostic_taxonomy()
d.tikz  # the TikZ document as text
d.save("stability_diagnostic_taxonomy.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("stability_diagnostic_taxonomy.tex")  # the complete TikZ document, the text source
d.save("stability_diagnostic_taxonomy.pdf")  # vector PDF; needs pdflatex

def interchange_criteria(*, labels: bool = True) -> Diagram:
    r"""Suydam's and Mercier's local interchange criteria across one schematic profile.

    $q(r)$ is Wesson's peaked-current profile (``peaked_current_safety_factor``,
    $q_a = 3$, $\nu = 3$, so $q_0 = 0.75$) and $p \propto (1 - r^2/a^2)^2$. Both
    criteria (``suydam_criterion``, ``mercier_criterion_circular``) are
    normalised by $p_0/a$; positive means the criterion holds. Mercier is
    Suydam plus $-p'q^2$: outward pressure fall stabilises outside $q = 1$,
    so its violated region stays inside $r_1$ while Suydam's extends past it.
    Both are necessary conditions for local interchanges, not global
    stability.
    """
    labels = _check_labels(labels)
    x = np.linspace(0.01, 1.0, 397)
    r = _A * x
    q = peaked_current_safety_factor(x, _Q_A, _NU)
    p = _P0 * (1.0 - x ** 2) ** 2
    dq, dp = gradient(r, q), gradient(r, p)
    scale = _A / _P0
    suydam = suydam_criterion(r, _B, q, dq, dp) * scale
    mercier = mercier_criterion_circular(r, _B, q, dq, dp) * scale
    r1 = float(np.interp(1.0, q, x))
    lo, hi = -1.2, 2.4
    chart = Chart(x_range=(0.0, 1.08), y_range=(lo, hi))
    chart.curves.update({
        "suydam": np.stack([x, suydam], -1), "mercier": np.stack([x, mercier], -1),
        "zero": np.array([[0.0, 0.0], [1.05, 0.0]]), "q_one": np.array([[r1, lo], [r1, hi]]),
    })

    def exit_x(values):
        return float(x[np.argmax(values >= hi)])

    def last_violation(values):
        return float(x[values < 0.0].max()) if np.any(values < 0.0) else 0.0

    chart.parameters.update({"q_a": _Q_A, "nu": _NU, "q0": float(q[0]), "r1": r1,
                             "suydam_violated_to": last_violation(suydam),
                             "mercier_violated_to": last_violation(mercier)})
    scene = render_chart(chart, x_label="$r/a$", y_label="criterion $\\times\\, a/p_0$", y_ticks=(0.0,),
                         curve_styles={"zero": "approx", "q_one": "rational", "suydam": "slope plus",
                                       "mercier": "outer solution"},
                         region_text={}, x_ticks=(r1, 1.0), x_tick_text=("$r_1$", "$1$"))
    items: List = []
    if labels:
        at = chart.to_cm
        items += [
            # each curve is named where it leaves the top of the frame
            Label(tuple(at(np.array([exit_x(mercier), hi]))), "Mercier (torus)", "small label",
                  anchor="south", role="mercier"),
            Label(tuple(at(np.array([exit_x(suydam), hi]))), "Suydam (cylinder)", "small label",
                  anchor="south", role="suydam"),
            Label(tuple(at(np.array([r1 - 0.02, hi - 0.1]))), "$q = 1$", "small label", anchor="north east",
                  role="q_one"),
            Label(tuple(at(np.array([0.02, 0.12]))), "holds", "small label", anchor="south west", role="sign"),
            Label(tuple(at(np.array([0.02, -0.3]))), "violated", "small label", anchor="north west", role="sign"),
            Label((CHART_WIDTH / 2, -1.4), f"$\\displaystyle {formula_equation(suydam_criterion)}$", "formula box",
                  anchor="north", role="equations"),
            Label((CHART_WIDTH / 2, -2.75), f"$\\displaystyle {formula_equation(mercier_criterion_circular)}$",
                  "formula box", anchor="north", role="equations"),
            Label((CHART_WIDTH / 2, -4.1), "Schematic $q$ and $p$; both are necessary conditions for local "
                  "interchanges, not global stability", "note", anchor="north", role="note"),
        ]
    return Diagram("interchange_criteria", scene + Scene(tuple(items)), model=chart)

import vaft.diagram

d = vaft.diagram.interchange_criteria()
d.tikz  # the TikZ document as text
d.save("interchange_criteria.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("interchange_criteria.tex")  # the complete TikZ document, the text source
d.save("interchange_criteria.pdf")  # vector PDF; needs pdflatex

def ballooning_formulation_hierarchy(*, labels: bool = True) -> Diagram:
    r"""How the reduced $s$-$\alpha$ model and the full-geometry ballooning codes relate.

    The general local high-$n$ ideal ballooning equation (field-line bending
    against the pressure-curvature drive) either loses its geometry step by
    step -- large aspect ratio, shifted circles, the $1 + \Lambda^2$ metric,
    the $\cos\theta + \Lambda\sin\theta$ curvature -- to become the
    Connor-Hastie-Taylor equation of ``s_alpha_ballooning_stable``, or keeps
    it, as Fortran DCON (asymptotic matching, $C_A$, stable when positive)
    and GPEC.jl (Dirichlet matching, ballooning $\Delta'$, stable when
    negative, with poles) do. Their marginal boundaries are comparable only
    after a common normalisation: ``shear_from_volume`` and
    ``ballooning_alpha_from_volume``, which reduce exactly to $\hat s$ and the
    CHT $\alpha$ for circles.
    """
    labels = _check_labels(labels)
    items: List = []
    nodes = {
        "general": box(2.9, 0.0, 13.5, 1.3, "general local high-$n$ ideal ballooning equation:\\\\ field-line bending "
                       "$|\\nabla\\beta|^2/B^2$ against the drive $P'\\kappa_w$ along an extended field line",
                       style="concept strong", role="node:general", latex=True),
    }
    y = -2.2
    for key, text in REDUCTION_STEPS:
        nodes[key] = box(_X_LEFT, y, 6.2, 0.9, text, role=f"node:{key}", latex=True)
        y -= 1.3
    nodes["vaft"] = box(_X_LEFT, y - 0.35, 6.2, 1.6, "VAFT reduced $s$--$\\alpha$ (CHT):\\\\ Newcomb test, "
                        "$\\alpha_1$ and $\\alpha_2$;\\\\ any $\\hat s$, $\\alpha$; no geometry", style="concept leaf",
                        role="node:vaft", latex=True)
    nodes["dcon"] = box(_X_DCON, -3.6, 5.6, 2.6, "Fortran DCON\\\\ full geometry, $\\theta_0 = 0$\\\\ asymptotic "
                        "small solution\\\\ index $C_A$, stable when $> 0$;\\\\ only where its own $D_I < 0$",
                        style="concept leaf", role="node:dcon", latex=True)
    nodes["gpec_jl"] = box(_X_JL, -3.6, 5.6, 2.6, "GPEC.jl local ballooning\\\\ full geometry, any $\\theta_k$\\\\ "
                           "Dirichlet at $\\pm\\theta_{max}$\\\\ $\\Delta'$, stable when $< 0$;\\\\ poles; "
                           "$\\alpha_{crit,1}$, $\\alpha_{crit,2}$ scans", style="concept leaf",
                           role="node:gpec_jl", latex=True)
    compare_y = nodes["vaft"].y - 2.0
    nodes["compare"] = box(2.9, compare_y, 12.0, 1.3, "compare marginal boundaries on one normalisation, not Boolean labels:\\\\ "
                           "$C_A$ and $\\Delta'$ zeros agree only as $\\theta_{max} \\to \\infty$ with $D_I < 0$", style="concept strong", role="node:compare",
                           latex=True)
    for b in nodes.values():
        items += list(b.items)
    edges = [("general", "large_aspect_ratio", "reduce geometry"), ("large_aspect_ratio", "shifted_circles", ""),
             ("shifted_circles", "metric", ""), ("metric", "curvature", ""), ("curvature", "vaft", ""),
             ("general", "dcon", ""), ("general", "gpec_jl", ""),
             ("vaft", "compare", ""), ("dcon", "compare", ""), ("gpec_jl", "compare", "")]
    for a, b, text in edges:
        arrow = connector(nodes[a], nodes[b], role=f"edge:{a}->{b}")
        items.append(arrow)
        if text and labels:
            mx, my = 0.5 * (arrow.start[0] + arrow.end[0]), 0.5 * (arrow.start[1] + arrow.end[1])
            side = "east" if mx < nodes["general"].x - 2.0 else "west"
            items.append(Label((mx + (-0.15 if side == "east" else 0.15), my), text, "small label", anchor=side,
                               role=f"edge:{a}->{b}"))
    if labels:
        # one label for both full-geometry arrows, right of the GPEC.jl one
        items.append(Label((8.5, -1.35), "both retain the geometry", "small label", anchor="west",
                           role="edge:general->full"))
    # the equations: what VAFT solves, and the normalisation all three must share
    eq_y = compare_y - 1.6
    items += [
        Label((-3.6, eq_y), f"$\\displaystyle {formula_equation(s_alpha_ballooning_stable)}$",
              "formula box", anchor="north", role="equation:cht"),
        Label((5.3, eq_y), f"$\\displaystyle {formula_equation(shear_from_volume)}$", "formula box", anchor="north",
              role="equation:shear"),
        Label((10.9, eq_y), f"$\\displaystyle {formula_equation(ballooning_alpha_from_volume)}$", "formula box",
              anchor="north", role="equation:alpha"),
    ]
    if labels:
        items += [Label((-3.6, eq_y + 0.1), "what VAFT solves", "small label", anchor="south", role="equation:cht"),
                  Label((8.1, eq_y + 0.1), "the normalisation all three must share", "small label", anchor="south",
                        role="equation:shear")]
    if labels:
        items.append(Label((2.9, eq_y - 1.9), "Common normalisation: both reduce exactly to $\\hat s$ and the CHT "
                           "$\\alpha$ for circular large-aspect-ratio surfaces ($\\psi$ per radian)", "note",
                           anchor="north", role="note"))
    model = {"steps": tuple(k for k, _ in REDUCTION_STEPS), "implementations": dict(IMPLEMENTATIONS),
             "nodes": {k: (b.x, b.y) for k, b in nodes.items()}, "edges": tuple((a, b) for a, b, _ in edges)}
    return Diagram("ballooning_formulation_hierarchy", Scene(tuple(items)), model=model)

import vaft.diagram

d = vaft.diagram.ballooning_formulation_hierarchy()
d.tikz  # the TikZ document as text
d.save("ballooning_formulation_hierarchy.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("ballooning_formulation_hierarchy.tex")  # the complete TikZ document, the text source
d.save("ballooning_formulation_hierarchy.pdf")  # vector PDF; needs pdflatex

def li_qa(*, reference: str = "wesson_1989", labels: bool = True) -> Diagram:
    r"""The $l_i$-$q$ diagram of one source: Wesson 1989 (JET, empirical) or Cheng 1987 (theory).

    Parameters
    ----------
    reference : {"wesson_1989", "cheng_1987"}
        Which published diagram. The two use different quantities: Wesson's
        $q_\psi$ and $l_i(3)$, Cheng's straight-cylinder $q(a)$ and $l_i$.
    labels : bool
        Draw region labels and the source note.

    Returns
    -------
    Diagram
        ``Diagram.model`` is a :class:`Chart` with the ``lower`` and ``upper``
        boundary curves in data coordinates and the projection key and
        boundary keys in ``parameters``.
    """
    try:
        spec = _REFERENCES[reference]
        spec = {**spec, "labels": dict(spec["labels"]), "text": dict(spec["text"])}   # never edit the table
    except KeyError:
        raise ValueError(f"reference must be one of {sorted(_REFERENCES)}, not {reference!r}") from None
    projection = get_projection(spec["projection"])
    chart = Chart(x_range=spec["x_range"], y_range=spec["y_range"])
    keys = []
    edges = []
    for key in projection.default_boundaries:
        boundary = _b.get_boundary(key)
        keys.append(key)
        if boundary.form == "threshold":  # a vertical q edge, drawn between the two branches below
            edges.append(boundary)
            continue
        lo, hi = boundary.applicability.ranges[boundary.input_names[0]]
        curve = _b.boundary_curve(boundary, boundary.input_names[0], _samples(lo, hi))
        xy = curve.xy[np.isfinite(curve.xy).all(axis=1)]
        # 9 decimals: the TikZ source (and the committed SVG's hash) must not depend on the platform's
        # last-bit floating-point differences, which flipped one printed coordinate (5.029 vs 5.0291)
        chart.curves[boundary.branch] = np.round(xy, 9)
    if reference == "cheng_1987":
        # the theoretical permissible domain, filled between the two bounds over the range Fig. 4 draws (#1603)
        lower, upper = chart.curves["lower"], chart.curves["upper"]
        top = np.interp(lower[:, 0], upper[:, 0], upper[:, 1], left=np.nan, right=np.nan)
        keep = np.isfinite(top)
        chart.curves["domain"] = np.round(np.vstack([lower[keep], np.c_[lower[keep, 0], top[keep]][::-1]]), 9)
        # the jig-saw edges at integer q(a) = m, where an m/1 resonance enters the plasma (comparison theorem, p. 354)
        for m in (3, 4, 5, 6):
            spec["labels"][f"m{m}"] = (m, 0.30)
            spec["text"][f"m{m}"] = f"${m}/1$"
    for boundary in edges:
        q_edge = float(_b.boundary_value(boundary))
        span = sorted(float(np.interp(q_edge, *chart.curves[b].T)) for b in ("lower", "upper"))
        chart.curves["edge"] = _b.threshold_curve(boundary, projection.y, span, target_axis="x").xy
    chart.labels.update(spec["labels"])
    chart.parameters.update({"reference": reference, "projection": projection.key, "boundaries": tuple(keys)})
    x_hi, y_hi = chart.x_range[1], chart.y_range[1]
    scene = _render_chart(
        chart,
        x_label=spec["x_label"],
        y_label=spec["y_label"],
        curve_styles={**({"domain": "region plasma"} if "domain" in chart.curves else {}),
                      **{name: "boundary" for name in ("lower", "upper", "edge") if name in chart.curves}},
        region_text=spec["text"] if labels else {},
        x_ticks=[t for t in _nice_ticks(x_hi) if t >= chart.x_range[0]],
        y_ticks=[t for t in _nice_ticks(y_hi) if t >= chart.y_range[0]],
        note=spec["note"] if labels else "",
    )
    return Diagram("li_qa", scene, model=chart)

import vaft.diagram

d = vaft.diagram.li_qa(reference='wesson_1989')
d.tikz  # the TikZ document as text
d.save("li_qa_wesson_1989.svg")  # rendered vector image, the canonical artifact; needs latex, dvisvgm
d.save("li_qa_wesson_1989.tex")  # the complete TikZ document, the text source
d.save("li_qa_wesson_1989.pdf")  # vector PDF; needs pdflatex

def exb_drift(*, mass_ratio: float = 4.0, projection: str = "perpendicular", labels: bool = True) -> Diagram:
    r"""E-cross-B drift of an ion and an electron in uniform crossed fields.

    $\mathbf{B} = B\hat z$ (out of the page) and $\mathbf{E} = E\hat x$. Both
    particles have the same kinetic energy in their guiding-centre frames, so
    the ion's orbit is $\sqrt{m_i/m_e}$ times larger; both guiding centres
    move at the same ``exb_drift_velocity`` -- along $-\hat y$ -- so no
    current flows.

    ``projection="perpendicular"`` is the plane normal to $\mathbf{B}$;
    ``"3d"`` shows the same orbits with a small parallel velocity, as helices
    along $\mathbf{B}$ drifting sideways.
    """
    _check_projection(projection, ("perpendicular", "3d"))
    labels = _check_labels(labels)
    mass_ratio = _validate_mass_ratio(mass_ratio)
    E, B = np.array([0.2, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])
    v_E = exb_drift_velocity(E, B)
    duration = 2.0 * math.pi * mass_ratio * 3.0  # three ion gyroperiods
    figure = ParticleFigure(parameters={"mass_ratio": mass_ratio, "E": float(E[0]), "B": float(B[2]),
                                        "duration": duration})
    rho_i = mass_ratio ** 0.5  # equal energy: u_i = u_e / sqrt(mass_ratio)
    gap = 2.2 * rho_i + 2.5
    for name, q, m, x_start in (("ion", 1.0, mass_ratio, 0.0), ("electron", -1.0, 1.0, gap)):
        u = 1.0 / math.sqrt(m)  # speed in the drift frame
        # start at the top of the orbit: gyration velocity along +x for an ion, -x for an electron
        u_vec = np.array([u * np.sign(q), 0.0, 0.0])
        x0 = np.array([x_start, 0.0, 0.0])
        v_par = np.array([0.0, 0.0, _EXB_V_PAR])
        x, _, dt = _orbit(q, m, x0, v_E + u_vec + v_par, _uniform(E), _uniform(B), 1.0,
                          duration / (2 * math.pi * m))
        figure.orbits[name] = x
        gc0 = _guiding_centre(q, m, x0, u_vec, B)
        figure.vectors[f"{name}_guiding_centre"] = gc0
        figure.vectors[f"{name}_drift"] = v_E
        figure.parameters[f"{name}_dt"] = dt
    figure.parameters["v_par"] = _EXB_V_PAR
    if projection == "3d":
        return _diagram("exb_drift_3d", _exb_3d(figure, labels), figure, labels)

    S = 0.45  # cm per unit length
    items: List = []
    for name, style in (("ion", "orbit ion"), ("electron", "orbit electron")):
        x = figure.orbits[name]
        items.append(Polyline.of(S * x[:, :2], style, role=f"{name}_orbit"))
        gc0 = figure.vectors[f"{name}_guiding_centre"]
        items.append(Arrow(tuple(S * gc0[:2]), tuple(S * (gc0 + v_E * duration)[:2]), "drift",
                           role=f"{name}_guiding_centre"))
        items.append(Label(tuple(S * x[0, :2]), "$+$" if name == "ion" else "$-$", "charge", role=f"{name}_start"))

    xs = np.concatenate([o[:, 0] for o in figure.orbits.values()])
    ys = np.concatenate([o[:, 1] for o in figure.orbits.values()])
    right, top, bottom = S * xs.max() + 1.3, S * ys.max(), S * ys.min()
    center = 0.5 * S * (xs.min() + xs.max()) + 1.2
    if labels:
        items += [
            Label((right + 0.25, top - 0.4), "$\\odot$", "legend symbol", role="legend"),
            Label((right + 0.8, top - 0.4), "$\\mathbf{B}$ ($+z$)", "label", anchor="west", role="legend"),
            Arrow((right, top - 1.5), (right + 1.2, top - 1.5), "field arrow", role="legend"),
            Label((right + 1.4, top - 1.5), "$\\mathbf{E}$ ($+x$)", "label", anchor="west", role="legend"),
            Label((right + 0.9, top - 2.6), "$\\mathbf{v}_E$ (guiding centre)", "small label", anchor="west",
                  role="legend"),
            Arrow((right, top - 2.6), (right + 0.8, top - 2.6), "drift", role="legend"),
            Arrow((right + 0.1, bottom + 0.4), (right + 1.1, bottom + 0.4), "axis", role="axes"),
            Label((right + 1.15, bottom + 0.4), "$x$", anchor="west", role="axes"),
            Arrow((right + 0.1, bottom + 0.4), (right + 0.1, bottom + 1.4), "axis", role="axes"),
            Label((right + 0.1, bottom + 1.45), "$y$", anchor="south", role="axes"),
        ]
        items += _title("$\\mathbf{E}\\times\\mathbf{B}$ drift",
                        "$\\mathbf{v}_E = \\mathbf{E}\\times\\mathbf{B}/B^2$: the same for ions and electrons, "
                        "so no current", top + 1.8, center)
        items.append(Label((center, bottom - 0.8),
                           f"orbits integrated from the Lorentz force at equal energy; "
                           f"$m_i/m_e = {mass_ratio:g}$ for visibility", "note", role="note"))
    return _diagram("exb_drift", Scene(tuple(items)), figure, labels)