giant/giant/analysis/render.py

"""Render reduced artifacts to styled PDFs + gallery metadata (the local step).

This is the *only* module that imports ``plotstyle`` (ETPlot's KIT matplotlib
theme), which renders through a real LaTeX toolchain — so it runs on the
submit/login node, never on a compute worker. It reads nothing but the small
``Reduced`` JSON files a run produced, so it is fully decoupled from the heavy
streaming compute.

For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
Optionally runs ``gallery generate`` to build the static HTML site.

Every rollout series gets a stable color via ``ps.get_color(i)``, ``i`` being
its position in ``payload["series"]`` — that position is fixed by the run's
YAML/``--label`` order (threaded unchanged from ``condor.RunMeta.rollouts``
through every ``PlotSpec``), so a given rollout keeps the same color across
every plot in a run. The reference, where a plot has one, always draws in one
fixed, distinct style (dark ink, dashed) instead of taking a slot in that
cycle.
"""

from __future__ import annotations

import dataclasses
import subprocess
from pathlib import Path

import numpy as np
import plotstyle as ps
from matplotlib.colors import LogNorm
import yaml

from giant.analysis.reduced import Reduced

_REFERENCE_LABEL = "reference (Geant4)"

_TEX_ESCAPE_MAP = {
    "\\": r"\textbackslash{}",
    "%": r"\%",
    "&": r"\&",
    "#": r"\#",
    "$": r"\$",
    "_": r"\_",
    "{": r"\{",
    "}": r"\}",
}


def _tex_escape(text: str) -> str:
    """Escape characters LaTeX treats specially in catalog-authored title/xlabel
    text (e.g. a literal ``%`` in a "90% of deposited energy" title, which
    ``usetex`` otherwise reads as a comment marker and aborts the whole figure —
    see gitea #81). A single pass over the *original* characters, so the
    backslashes an escape itself introduces (e.g. ``\textbackslash{}``) are
    never re-escaped."""
    return "".join(_TEX_ESCAPE_MAP.get(c, c) for c in text)


def _ref_color() -> str:
    return ps.colors.INK["primary"]


def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
    counts = np.asarray(counts, dtype=np.float64)
    total = counts.sum()
    if total == 0:
        return counts
    return counts / (total * (edges[1] - edges[0]))


def _overlay(ax, edges: np.ndarray, payload: dict, log_y: bool) -> None:
    if "reference" in payload:
        ax.stairs(
            _density(payload["reference"], edges), edges, label=_REFERENCE_LABEL, color=_ref_color(), linestyle="--"
        )
    for i, (name, counts) in enumerate(payload.get("series", {}).items()):
        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
    if log_y:
        ax.set_yscale("log")


def _router_summary(router_cfg: dict) -> str:
    if not router_cfg.get("enabled"):
        return "off"
    return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}"


def _figure_params_v2(mc: dict, meta: dict) -> dict:
    """`_figure_params_single` for a new-shape (nested) `model_config` — has a
    `stage1_model` key. Reports stage 1's architecture (the headline
    generator); stage 2's generator is only added (`mode_s2`) when it
    differs from stage 1's, since a mixed run (the `stage1=flow` +
    `stage2=wgan` case) is the interesting exception, not
    the common case."""
    s1 = mc["stage1_model"]
    s2 = mc.get("stage2_model") or {}
    mode = s1.get("generator")
    params: dict = {}
    if s1.get("hidden_dim") is not None:
        params["hidden_dim"] = s1["hidden_dim"]
    if s1.get("n_res_blocks") is not None:
        params["n_res_blocks"] = s1["n_res_blocks"]
    if mode is not None:
        params["mode"] = mode
    s2_mode = s2.get("generator")
    if s2_mode is not None and s2_mode != mode:
        params["mode_s2"] = s2_mode
    particle_type = ((mc.get("conditioning") or {}).get("particle") or {}).get("type")
    if particle_type is not None:
        params["conditioning"] = particle_type
    params["router"] = _router_summary(s1.get("router") or {})
    if meta.get("training_epoch") is not None:
        params["epoch"] = meta["training_epoch"]
    if meta.get("best_val_loss") is not None:
        params["best_val_loss"] = round(meta["best_val_loss"], 4)
    if mode == "wgan":
        noise_dim = (s1.get("wgan") or {}).get("noise_dim")
        if noise_dim is not None:
            params["noise_dim"] = noise_dim
    elif meta.get("steps") is not None:
        params["steps"] = meta["steps"]
    return params


def _figure_params_single(meta: dict) -> dict:
    """Curated run identity for the figure subtitle (``new_figure(params=...)``),
    for exactly one rollout's ``plot_meta``.

    ``meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
    carry every threaded model/training/rollout/dataset parameter for
    after-the-fact lookup — this picks only the handful that matter for
    telling figures apart at a glance while flipping through a gallery, since
    the subtitle is one unwrapped line of text. The last slot is
    architecture-conditional: flow/ddpm runs show the ODE ``steps`` used for
    this rollout, wgan runs show ``noise_dim`` instead since wgan sampling is
    single-pass and has no ODE step count.

    Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0
    nested one (has a ``stage1_model`` key — see ``_figure_params_v2``).
    """
    mc = meta.get("model_config") or {}
    if "stage1_model" in mc:
        return _figure_params_v2(mc, meta)

    mode = mc.get("mode")
    params: dict = {}
    if mc.get("hidden_dim") is not None:
        params["hidden_dim"] = mc["hidden_dim"]
    if mc.get("n_blocks") is not None:
        params["n_blocks"] = mc["n_blocks"]
    if mode is not None:
        params["mode"] = mode
    if mc.get("conditioning") is not None:
        params["conditioning"] = mc["conditioning"]
    params["router"] = _router_summary(mc.get("router") or {})
    if meta.get("training_epoch") is not None:
        params["epoch"] = meta["training_epoch"]
    if meta.get("best_val_loss") is not None:
        params["best_val_loss"] = round(meta["best_val_loss"], 4)
    if mode == "wgan":
        if mc.get("noise_dim") is not None:
            params["noise_dim"] = mc["noise_dim"]
    elif meta.get("steps") is not None:
        params["steps"] = meta["steps"]
    return params


def _figure_params(run_meta: dict) -> dict:
    """Curated run identity for the figure subtitle.

    A single-rollout run reuses that rollout's ``plot_meta`` (same curated
    model/training/rollout subset as always — see ``_figure_params_single``);
    a multi-rollout run instead names the series being compared, since no
    single ``model_config`` applies to the figure as a whole (each plot's own
    gallery YAML still carries every rollout's full ``plot_meta`` for
    after-the-fact lookup, via ``_plot_metadata``).
    """
    rollouts = run_meta.get("rollouts") or {}
    if len(rollouts) == 1:
        ((_, meta),) = rollouts.items()
        return _figure_params_single(meta)
    return {"rollouts": ", ".join(rollouts)} if rollouts else {}


def _render_overlay(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    _overlay(ax, edges, r.payload, r.payload.get("log_y", False))
    ax.set_xlabel(r.xlabel)
    ax.set_ylabel("density")
    ps.style_legend(ax, title="source")
    return fig


def _render_single(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    for i, (name, counts) in enumerate(r.payload.get("series", {}).items()):
        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
    if r.payload.get("log_y"):
        ax.set_yscale("log")
    if r.payload.get("log_x"):
        ax.set_xscale("log")
    ax.set_xlabel(r.xlabel)
    ax.set_ylabel("density")
    ps.style_legend(ax, title="source")
    return fig


def _render_grouped(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    groups = r.payload["groups"]
    labels = list(groups)
    n = len(labels)
    ncols = min(3, n) or 1
    nrows = (n + ncols - 1) // ncols
    fig, axes = ps.new_figure(
        "slide-16x9",
        title=r.title,
        params=params,
        nrows=nrows,
        ncols=ncols,
        squeeze=False,
    )
    flat = axes.ravel()
    for i, lbl in enumerate(labels):
        ax = flat[i]
        _overlay(ax, edges, groups[lbl], r.payload.get("log_y", False))
        ax.set_title(lbl, fontsize=8)
        ax.set_xlabel(r.xlabel)
    for j in range(n, len(flat)):
        flat[j].set_visible(False)
    ps.style_legend(flat[0], title="source")
    return fig


def _render_profile(r: Reduced, params: dict):
    edges = np.asarray(r.payload["edges"])
    centers = 0.5 * (edges[:-1] + edges[1:])
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    if "reference" in r.payload:
        ref = r.payload["reference"]
        mean, std = np.asarray(ref["mean"]), np.asarray(ref["std"])
        color = _ref_color()
        ax.plot(centers, mean, label=_REFERENCE_LABEL, color=color, linestyle="--")
        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
    for i, (name, side) in enumerate(r.payload.get("series", {}).items()):
        mean, std = np.asarray(side["mean"]), np.asarray(side["std"])
        color = ps.get_color(i)
        ax.plot(centers, mean, label=name, color=color)
        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
    ax.set_xlabel(r.xlabel)
    ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
    ps.style_legend(ax, title="source")
    return fig


def _render_bar(r: Reduced, params: dict):
    labels = r.payload["labels"]
    x = np.arange(len(labels))
    series = r.payload.get("series", {})
    has_ref = "reference" in r.payload
    n_bars = len(series) + (1 if has_ref else 0)
    width = 0.8 / max(n_bars, 1)
    offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, n_bars)
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    idx = 0
    if has_ref:
        ax.bar(x + offsets[idx], r.payload["reference"], width, label=_REFERENCE_LABEL, color=_ref_color())
        idx += 1
    for i, (name, vals) in enumerate(series.items()):
        ax.bar(x + offsets[idx], vals, width, label=name, color=ps.get_color(i))
        idx += 1
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.set_ylabel(r.payload.get("ylabel", "value"))
    ps.style_legend(ax, title="source")
    return fig


def _render_router_gating(r: Reduced, params: dict):
    series = r.payload.get("series", {})
    names = list(series)
    log_x = r.payload.get("log_x", False)
    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=2, squeeze=False)
    for row, name in enumerate(names):
        entry = series[name]
        n_experts = entry["n_experts"]
        for col, key in enumerate(("rollout", "reference")):
            ax = axes[row, col]
            side = entry.get(key, {})
            centers = np.asarray(side.get("centers", []))
            means = np.asarray(side.get("means", []))
            if len(centers) and means.size:
                cum = np.zeros(len(centers))
                for i in range(n_experts):
                    ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
                    cum = cum + means[:, i]
            if log_x:
                ax.set_xscale("log")
            ax.set_ylim(0, 1)
            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
            ax.set_title(f"{name} — {panel_label}", fontsize=8)
            if row == len(names) - 1:
                ax.set_xlabel(r.xlabel)
        axes[row, 0].set_ylabel("mean gate weight")
    if names:
        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
    return fig


def _render_router_share(r: Reduced, params: dict):
    series = r.payload.get("series", {})
    names = list(series)
    present: tuple[str, ...] = ("rollout", "reference")
    if names:
        present = tuple(k for k in ("rollout", "reference") if k in series[names[0]])
    ncols = max(len(present), 1)
    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=ncols, squeeze=False)
    for row, name in enumerate(names):
        entry = series[name]
        n_experts = entry["n_experts"]
        cats = entry["categories"]
        x = np.arange(len(cats))
        for col, key in enumerate(present):
            ax = axes[row, col]
            side = entry.get(key)
            if side is not None:
                shares = np.array([side[c] for c in cats])  # (n_cat, n_experts)
                bottom = np.zeros(len(cats))
                for i in range(n_experts):
                    ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
                    bottom += shares[:, i]
            ax.set_xticks(x)
            ax.set_xticklabels(cats, rotation=45, ha="right")
            ax.set_ylim(0, 1)
            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
            ax.set_title(f"{name} — {panel_label}", fontsize=8)
        axes[row, 0].set_ylabel("share of rows dispatched to expert")
    if names:
        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
    return fig


def _render_router_specialization(r: Reduced, params: dict):
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    series = r.payload.get("series", {})
    chance_levels: set[float] = set()
    for i, (name, entry) in enumerate(series.items()):
        color = ps.get_color(i)
        if entry.get("chance_level") is not None:
            chance_levels.add(entry["chance_level"])
        for key, linestyle, label in (
            ("rollout", "-", name),
            ("reference", "--", f"{name} ({_REFERENCE_LABEL})"),
        ):
            side = entry.get(key)
            if side and side["centers"]:
                ax.plot(
                    side["centers"],
                    side["score"],
                    label=label,
                    color=color,
                    linestyle=linestyle,
                    marker="o",
                    markersize=3,
                )
    for lvl in sorted(chance_levels):
        ax.axhline(lvl, linestyle=":", color="gray")
    if r.payload.get("log_x"):
        ax.set_xscale("log")
    ax.set_ylim(0, 1)
    ax.set_xlabel(r.xlabel)
    ax.set_ylabel("max gate weight")
    ps.style_legend(ax, title="router")
    return fig


def _render_heatmap(r: Reduced, params: dict):
    series = dict(r.payload["series"])
    row_labels = r.payload["row_labels"]
    col_labels = r.payload["col_labels"]
    # A heatmap-shaped plot is one matrix per rollout, so the reference (when the
    # comparison has one — the distance scorecard doesn't) becomes one more panel
    # rather than another line.
    if r.payload.get("reference") is not None:
        series["reference"] = r.payload["reference"]
    names = list(series)
    norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
    fig, axes = ps.new_figure(
        "slide-16x9" if len(names) > 1 else "thesis-single",
        title=r.title,
        params=params,
        nrows=1,
        ncols=len(names),
        squeeze=False,
    )
    flat = axes.ravel()
    im = None
    for ax, name in zip(flat, names):
        mat = np.asarray(series[name], dtype=float)
        im = ax.imshow(
            mat,
            origin="upper",
            aspect="auto",
            cmap=r.payload.get("cmap", "viridis"),
            norm=norm,
            vmin=None if norm else r.payload.get("vmin"),
            vmax=None if norm else r.payload.get("vmax"),
        )
        ax.set_xticks(range(len(col_labels)))
        ax.set_xticklabels(col_labels, rotation=45, ha="right")
        ax.set_yticks(range(len(row_labels)))
        ax.set_yticklabels(row_labels)
        ax.set_xlabel(r.xlabel)
        if len(names) > 1:
            ax.set_title(name, fontsize=8)
    flat[0].set_ylabel(r.payload.get("ylabel", ""))
    fig.colorbar(im, ax=list(flat), label=r.payload.get("cbar_label", "value"))
    return fig


def _render_unavailable(r: Reduced, params: dict):
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    ax.axis("off")
    ax.text(
        0.5,
        0.5,
        r.payload.get("note", "not available"),
        ha="center",
        va="center",
        wrap=True,
        fontsize=10,
        transform=ax.transAxes,
    )
    return fig


_RENDERERS = {
    "overlay_hist": _render_overlay,
    "single_hist": _render_single,
    "grouped_hist": _render_grouped,
    "profile": _render_profile,
    "bar": _render_bar,
    "router_gating": _render_router_gating,
    "router_share": _render_router_share,
    "router_specialization": _render_router_specialization,
    "heatmap": _render_heatmap,
    "unavailable": _render_unavailable,
}


def render(r: Reduced, run_meta: dict | None = None):
    """Build the matplotlib figure for one reduced artifact (dispatch on kind).

    ``title``/``xlabel`` are LaTeX-escaped here, at the one point every kind's
    renderer draws them from — ``_plot_metadata`` deliberately keeps using the
    unescaped ``r`` for the gallery YAML, which isn't LaTeX.
    """
    escaped = dataclasses.replace(r, title=_tex_escape(r.title), xlabel=_tex_escape(r.xlabel))
    return _RENDERERS[r.kind](escaped, _figure_params(run_meta or {}))


def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
    meta = {
        "title": r.title,
        "description": f"Rollout vs reference: {r.title}.",
        "plot_type": r.kind,
        "family": r.family,
    }
    meta.update(r.meta)
    if "note" in r.payload:
        meta["note"] = r.payload["note"]
    if run_meta:
        # Every threaded model/training/rollout/dataset parameter, so a
        # single plot's metadata is self-contained for later comparison
        # without cross-referencing the run's root metadata.yaml.
        meta["parameters"] = {k: v for k, v in run_meta.items() if k != "title"}
    return meta


def render_all(
    reduced_dir: str | Path,
    out_dir: str | Path,
    run_meta: dict | None = None,
    *,
    run_gallery: bool = False,
) -> list[Path]:
    """Render every reduced artifact under ``reduced_dir`` to a PDF tree.

    Writes ``<out>/<family>/<id>.pdf`` + ``<id>.yaml`` and a per-family
    ``metadata.yaml`` (carrying the run's checkpoint/paths as gallery params).
    Returns the list of PDF paths written.
    """
    ps.use()
    run_meta = run_meta or {}
    reduced_dir, out_dir = Path(reduced_dir), Path(out_dir)
    pdfs: list[Path] = []
    families: set[str] = set()

    for jf in sorted(reduced_dir.glob("*.json")):
        r = Reduced.load(jf)
        family_dir = out_dir / r.family
        family_dir.mkdir(parents=True, exist_ok=True)
        families.add(r.family)
        fig = render(r, run_meta)
        ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
        (family_dir / f"{r.id}.yaml").write_text(yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False))
        pdfs.append(family_dir / f"{r.id}.pdf")
        import matplotlib.pyplot as plt

        plt.close(fig)

    # Root + per-family gallery metadata.
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "metadata.yaml").write_text(
        yaml.safe_dump(
            {
                "title": run_meta.get("title", "GIANT rollout analysis"),
                "description": "Autoregressive rollout(s) compared against a held-out Geant4 reference steps file.",
                "experiment": "GIANT",
                "parameters": {k: v for k, v in run_meta.items() if k != "title"},
            },
            sort_keys=False,
        )
    )
    for fam in families:
        (out_dir / fam / "metadata.yaml").write_text(
            yaml.safe_dump({"title": fam, "description": f"{fam} plots."}, sort_keys=False)
        )

    if run_gallery:
        subprocess.run(["gallery", "generate", "--source", str(out_dir)], check=True)
    return pdfs


def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
    """Render a prepped run directory: ``<run_dir>/reduced`` → ``<run_dir>/plots``.

    First joins every plot's chunk partials (``reduced_partial/<id>__*.json``)
    into ``reduced/<id>.json`` via ``merge_all`` — a no-op merge when the run
    wasn't chunked (``n_chunks=1``) — then pulls the rollout provenance
    (checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's
    gallery metadata and renders.
    """
    from giant.analysis.condor import RunMeta, merge_all

    run_dir = Path(run_dir)
    merge_all(run_dir)
    meta = RunMeta.load(run_dir / "run_meta.json")
    run_meta = {
        "title": meta.title,
        "reference": meta.reference,
        "rollouts": {ro["name"]: ro["plot_meta"] for ro in meta.rollouts},
    }
    return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)