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.
"""

from __future__ import annotations

import subprocess
from pathlib import Path

import numpy as np
import plotstyle as ps
import yaml

from giant.analysis.reduced import Reduced

_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}


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, series: dict[str, list], log_y: bool) -> None:
    for key in ("reference", "rollout"):
        if key in series:
            ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
    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, run_meta: dict) -> dict:
    """`_figure_params` 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 run_meta.get("training_epoch") is not None:
        params["epoch"] = run_meta["training_epoch"]
    if run_meta.get("best_val_loss") is not None:
        params["best_val_loss"] = round(run_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 run_meta.get("steps") is not None:
        params["steps"] = run_meta["steps"]
    return params


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

    ``run_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 = run_meta.get("model_config") or {}
    if "stage1_model" in mc:
        return _figure_params_v2(mc, run_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 run_meta.get("training_epoch") is not None:
        params["epoch"] = run_meta["training_epoch"]
    if run_meta.get("best_val_loss") is not None:
        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
    if mode == "wgan":
        if mc.get("noise_dim") is not None:
            params["noise_dim"] = mc["noise_dim"]
    elif run_meta.get("steps") is not None:
        params["steps"] = run_meta["steps"]
    return params


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)
    ax.stairs(_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"])
    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)
    for key in ("reference", "rollout"):
        mean = np.asarray(r.payload[f"{key}_mean"])
        std = np.asarray(r.payload[f"{key}_std"])
        (line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=line.get_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))
    width = 0.4
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    ax.bar(x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"])
    ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
    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):
    n_experts = r.payload["n_experts"]
    log_x = r.payload.get("log_x", False)
    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False)
    flat = axes.ravel()
    for ax, key in zip(flat, ("rollout", "reference")):
        side = r.payload.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)
        ax.set_title(_SERIES_LABELS[key], fontsize=8)
        ax.set_xlabel(r.xlabel)
    flat[0].set_ylabel("mean gate weight")
    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
    return fig


def _render_router_share(r: Reduced, params: dict):
    categories = r.payload["categories"]
    n_experts = r.payload["n_experts"]
    x = np.arange(len(categories))
    present = [k for k in ("rollout", "reference") if k in r.payload]
    fig, axes = ps.new_figure(
        "slide-16x9",
        title=r.title,
        params=params,
        nrows=1,
        ncols=len(present),
        squeeze=False,
    )
    flat = axes.ravel()
    for ax, key in zip(flat, present):
        side = r.payload[key]
        shares = np.array([side[c] for c in categories])  # (n_cat, n_experts)
        bottom = np.zeros(len(categories))
        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(categories, rotation=45, ha="right")
        ax.set_ylim(0, 1)
        ax.set_title(_SERIES_LABELS[key], fontsize=8)
    flat[0].set_ylabel("share of rows dispatched to expert")
    ps.style_legend(flat[0], title=f"{r.payload.get('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)
    for key in ("reference", "rollout"):
        side = r.payload.get(key)
        if side and side["centers"]:
            ax.plot(side["centers"], side["score"], label=_SERIES_LABELS[key], marker="o", markersize=3)
    chance = r.payload.get("chance_level")
    if chance is not None:
        ax.axhline(chance, linestyle="--", color="gray", label="chance level (1/n_experts)")
    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=f"{r.payload.get('router_type', '')} router")
    return fig


def _render_heatmap(r: Reduced, params: dict):
    mat = np.asarray(r.payload["matrix"], dtype=float)
    row_labels = r.payload["row_labels"]
    col_labels = r.payload["col_labels"]
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    im = ax.imshow(
        mat,
        origin="upper",
        aspect="auto",
        cmap=r.payload.get("cmap", "viridis"),
        vmin=r.payload.get("vmin"),
        vmax=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)
    ax.set_ylabel(r.payload.get("ylabel", ""))
    fig.colorbar(im, ax=ax, 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)."""
    return _RENDERERS[r.kind](r, _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 compared against held-out Geant4 reference steps.",
                "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,
        "rollout": meta.rollout,
        "reference": meta.reference,
        **meta.plot_meta,
    }
    return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)