giant/tests/test_render.py

"""Render smoke test — skipped where plotstyle / LaTeX is unavailable."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

pytest.importorskip("plotstyle")

from giant.analysis import render as render_mod
from giant.analysis.reduced import Reduced


def _try_render(reduced: list[Reduced], out: Path):
    from giant.analysis.render import render_all

    for r in reduced:
        r.save(out / "reduced" / f"{r.id}.json")
    return render_all(out / "reduced", out / "plots")


def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
    reduced = [
        Reduced(
            "rg",
            "model",
            "router_gating",
            "Router gating",
            "pre-step energy [MeV]",
            {
                "log_x": True,
                "series": {
                    "flow": {
                        "n_experts": 2,
                        "router_type": "energy",
                        "rollout": {
                            "centers": [1.0, 10.0, 100.0],
                            "means": [[0.6, 0.4], [0.5, 0.5], [0.4, 0.6]],
                        },
                        "reference": {
                            "centers": [1.0, 10.0, 100.0],
                            "means": [[0.55, 0.45], [0.5, 0.5], [0.45, 0.55]],
                        },
                    },
                    "wgan": {
                        "n_experts": 2,
                        "router_type": "energy",
                        "rollout": {"centers": [1.0], "means": [[0.5, 0.5]]},
                        "reference": {"centers": [1.0], "means": [[0.5, 0.5]]},
                    },
                },
            },
        ),
        Reduced(
            "rs",
            "model",
            "router_share",
            "Router share",
            "species",
            {
                "series": {
                    "flow": {
                        "categories": ["e-", "gamma"],
                        "n_experts": 2,
                        "router_type": "energy",
                        "rollout": {"e-": [0.7, 0.3], "gamma": [0.2, 0.8]},
                        "reference": {"e-": [0.6, 0.4], "gamma": [0.3, 0.7]},
                    },
                },
            },
        ),
        Reduced(
            "rp",
            "model",
            "router_share",
            "Router share by process (reference-only)",
            "process",
            {
                "series": {
                    "flow": {
                        "categories": ["compt", "phot"],
                        "n_experts": 2,
                        "router_type": "energy",
                        "reference": {"compt": [0.4, 0.6], "phot": [0.9, 0.1]},
                    },
                },
            },
        ),
        Reduced(
            "rz",
            "model",
            "router_specialization",
            "Router specialization",
            "pre-step energy [MeV]",
            {
                "log_x": True,
                "series": {
                    "flow": {
                        "n_experts": 2,
                        "chance_level": 0.5,
                        "rollout": {"centers": [1.0, 10.0], "score": [0.6, 0.7]},
                        "reference": {"centers": [1.0, 10.0], "score": [0.55, 0.65]},
                    },
                    "wgan": {
                        "n_experts": 4,
                        "chance_level": 0.25,
                        "rollout": {"centers": [1.0, 10.0], "score": [0.3, 0.4]},
                        "reference": {"centers": [], "score": []},
                    },
                },
            },
        ),
        Reduced(
            "ru",
            "model",
            "unavailable",
            "Router unavailable",
            "x",
            {"note": "router diagnostics unavailable: no router in this run"},
        ),
        Reduced(
            "g4",
            "marginals",
            "grouped_hist",
            "Grouped (4)",
            "x",
            {
                "edges": [0, 1, 2],
                "groups": {
                    lbl: {"series": {"flow": [1, 2], "wgan": [2, 1]}, "reference": [2, 1]}
                    for lbl in ("a", "b", "c", "d")
                },
                "log_y": True,
            },
        ),
        Reduced(
            "sl",
            "species",
            "single_hist",
            "Single (log-x)",
            "x",
            {"edges": [1, 10, 100], "series": {"flow": [5, 1], "wgan": [3, 2]}, "log_x": True, "log_y": True},
        ),
        Reduced(
            "hm",
            "quality",
            "heatmap",
            "Distance summary (2 rollouts)",
            "grouping axis",
            {
                "series": {"flow": [[0.1, 0.2], [0.3, 0.4]], "wgan": [[0.5, 0.6], [0.7, 0.8]]},
                "row_labels": ["step_length", "edep"],
                "col_labels": ["overall", "energy"],
                "cbar_label": "KS statistic",
                "vmin": 0.0,
                "vmax": 1.0,
            },
        ),
    ]
    try:
        pdfs = _try_render(reduced, tmp_path)
    except RuntimeError as e:  # LaTeX missing at render time
        pytest.skip(f"LaTeX rendering unavailable: {e}")
    assert len(pdfs) == len(reduced)
    assert all(p.exists() for p in pdfs)


def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
    # render_mod.subprocess *is* the stdlib subprocess module, so a blanket
    # patch of .run would also swallow the real subprocess.run calls
    # matplotlib's texmanager makes to compile LaTeX during savefig — only
    # intercept the "gallery generate" call itself and pass everything else
    # (LaTeX included) through to the real subprocess.run.
    calls = []
    real_run = render_mod.subprocess.run

    def fake_run(*a, **k):
        if a and a[0] and a[0][0] == "gallery":
            calls.append((a, k))
            return None
        return real_run(*a, **k)

    monkeypatch.setattr(render_mod.subprocess, "run", fake_run)
    reduced = [
        Reduced(
            "s",
            "species",
            "single_hist",
            "Single",
            "x",
            {"edges": [0, 1, 2], "series": {"rollout": [5, 1]}},
        )
    ]
    for r in reduced:
        r.save(tmp_path / "reduced" / f"{r.id}.json")
    try:
        render_mod.render_all(tmp_path / "reduced", tmp_path / "plots", run_gallery=True)
    except RuntimeError as e:
        pytest.skip(f"LaTeX rendering unavailable: {e}")

    assert len(calls) == 1
    args, kwargs = calls[0]
    assert args[0] == ["gallery", "generate", "--source", str(tmp_path / "plots")]
    assert kwargs == {"check": True}


def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkeypatch):
    from giant.analysis import condor as condor_mod

    run_dir = tmp_path / "run"
    (run_dir / "reduced").mkdir(parents=True)

    merge_calls = []
    monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
    meta = condor_mod.RunMeta(
        rollouts=[{"name": "rollout", "path": "rollout.parquet", "plot_meta": {"checkpoint": "ckpt/best.pt"}}],
        reference="reference.parquet",
        run_dir=str(run_dir),
        title="my-run",
    )
    monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))

    Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "series": {"rollout": [1]}}).save(
        run_dir / "reduced" / "s.json"
    )

    try:
        pdfs = render_mod.render_run(run_dir)
    except RuntimeError as e:
        pytest.skip(f"LaTeX rendering unavailable: {e}")

    assert merge_calls == [run_dir]
    assert len(pdfs) == 1
    plot_meta = (run_dir / "plots" / "species" / "s.yaml").read_text()
    assert "checkpoint" in plot_meta
    root_meta = (run_dir / "plots" / "metadata.yaml").read_text()
    assert "my-run" in root_meta


def test_render_one_of_each_kind(tmp_path: Path):
    reduced = [
        Reduced(
            "m",
            "marginals",
            "overlay_hist",
            "Overlay",
            "x",
            {
                "edges": [0, 1, 2, 3],
                "series": {"flow": [1, 2, 3], "wgan": [2, 2, 2]},
                "reference": [3, 2, 1],
                "log_y": False,
            },
        ),
        Reduced(
            "g",
            "marginals",
            "grouped_hist",
            "Grouped",
            "x",
            {
                "edges": [0, 1, 2],
                "groups": {"a": {"series": {"flow": [1, 2]}, "reference": [2, 1]}},
                "log_y": False,
            },
        ),
        Reduced(
            "p",
            "shower",
            "profile",
            "Profile",
            "depth",
            {
                "edges": [0, 1, 2],
                "series": {"flow": {"mean": [1, 2], "std": [0.1, 0.2]}},
                "reference": {"mean": [1.1, 1.9], "std": [0.1, 0.1]},
                "ylabel": "e",
            },
        ),
        Reduced(
            "b",
            "species",
            "bar",
            "Bar",
            "species",
            {
                "labels": ["e-", "gamma"],
                "series": {"flow": [0.6, 0.4], "wgan": [0.55, 0.45]},
                "reference": [0.5, 0.5],
                "ylabel": "frac",
            },
        ),
        Reduced(
            "cost",
            "cost",
            "bar",
            "Cost",
            "phase",
            {
                "labels": ["sample", "write", "total"],
                "series": {"flow": [10.0, 1.0, 11.0]},
                "reference": [5.0, 0.5, 5.5],
                "ylabel": "us/step",
                "log_y": True,
            },
        ),
        Reduced(
            "s",
            "species",
            "single_hist",
            "Single",
            "x",
            {"edges": [0, 1, 2], "series": {"flow": [5, 1]}, "log_y": True},
        ),
        Reduced(
            "hm1",
            "secondaries",
            "heatmap",
            "Heatmap (single rollout)",
            "predicted",
            {
                "series": {"flow": [[1, 0], [0, 1]]},
                "reference": [[2, 0], [0, 1]],
                "row_labels": ["0", "1+"],
                "col_labels": ["0", "1+"],
                "cbar_label": "count",
                "log_color": True,
            },
        ),
        Reduced(
            "ph1",
            "prediction",
            "paired_hist",
            "Paired hist (single prediction)",
            "x",
            {"edges": [0, 1, 2, 3], "series": {"pred": {"pred": [1, 2, 3], "true": [2, 2, 2]}}, "log_y": False},
        ),
        Reduced(
            "ph2",
            "prediction",
            "paired_hist",
            "Paired hist (two predictions)",
            "x",
            {
                "edges": [0, 1, 2, 3],
                "series": {"a": {"pred": [1, 2, 3], "true": [2, 2, 2]}, "b": {"pred": [3, 2, 1]}},
                "log_y": False,
            },
        ),
        Reduced(
            "hm2d",
            "prediction",
            "heatmap2d",
            "Scatter (truth vs pred)",
            "true x",
            {
                "x_edges": [0, 1, 2],
                "y_edges": [0, 1, 2],
                "series": {"pred": [[2, 0], [1, 3]]},
                "ylabel": "predicted x",
                "cbar_label": "count",
                "log_color": True,
                "diagonal": True,
            },
        ),
        Reduced(
            "profile_noref",
            "prediction",
            "profile",
            "Residual profile (no reference)",
            "true x",
            {"edges": [0, 1, 2], "series": {"pred": {"mean": [0.1, -0.1], "std": [0.2, 0.2]}}},
        ),
        Reduced(
            "bar_noref",
            "prediction",
            "bar",
            "Constraint violations (no reference)",
            "check",
            {"labels": ["a", "b"], "series": {"pred": [0.01, 0.0]}, "ylabel": "rate"},
        ),
    ]
    try:
        pdfs = _try_render(reduced, tmp_path)
    except RuntimeError as e:  # LaTeX missing at render time
        pytest.skip(f"LaTeX rendering unavailable: {e}")
    assert len(pdfs) == len(reduced)
    assert all(p.exists() for p in pdfs)
    assert (tmp_path / "plots" / "metadata.yaml").exists()


# ── pure-function helpers: no matplotlib figure needed ──────────────────


def test_density_zero_total_returns_counts_unchanged():
    counts = np.array([0.0, 0.0, 0.0])
    out = render_mod._density(counts, np.array([0.0, 1.0, 2.0, 3.0]))
    np.testing.assert_array_equal(out, counts)


def test_density_normalizes_by_total_and_bin_width():
    counts = [1, 3]
    edges = np.array([0.0, 2.0, 4.0])  # bin width 2
    out = render_mod._density(counts, edges)
    np.testing.assert_allclose(out, np.array([1, 3]) / (4 * 2))


def test_router_summary_disabled_is_off():
    assert render_mod._router_summary({"enabled": False, "type": "energy"}) == "off"
    assert render_mod._router_summary({}) == "off"


def test_router_summary_enabled_formats_type_and_n_experts():
    cfg = {"enabled": True, "type": "energy", "n_experts": 8}
    assert render_mod._router_summary(cfg) == "energy×8"


def test_figure_params_v2_basics_and_router_and_epoch():
    mc = {
        "stage1_model": {
            "hidden_dim": 256,
            "n_res_blocks": 4,
            "generator": "flow",
            "router": {"enabled": True, "type": "energy", "n_experts": 4},
        },
        "conditioning": {"particle": {"type": "physical"}},
    }
    meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10, "model_config": mc}
    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
    assert params == {
        "hidden_dim": 256,
        "n_res_blocks": 4,
        "mode": "flow",
        "conditioning": "physical",
        "router": "energy×4",
        "epoch": 12,
        "best_val_loss": 0.1235,
        "steps": 10,
    }


def test_figure_params_v2_wgan_reports_noise_dim_not_steps():
    mc = {
        "stage1_model": {
            "generator": "wgan",
            "wgan": {"noise_dim": 32},
        },
    }
    meta = {"model_config": mc, "steps": 10}
    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
    assert params["mode"] == "wgan"
    assert params["noise_dim"] == 32
    assert "steps" not in params


def test_figure_params_v2_reports_mode_s2_only_when_it_differs():
    same = {
        "stage1_model": {"generator": "flow"},
        "stage2_model": {"generator": "flow"},
    }
    assert "mode_s2" not in render_mod._figure_params({"rollouts": {"rollout": {"model_config": same}}})

    mixed = {
        "stage1_model": {"generator": "flow"},
        "stage2_model": {"generator": "wgan"},
    }
    params = render_mod._figure_params({"rollouts": {"rollout": {"model_config": mixed}}})
    assert params["mode_s2"] == "wgan"


def test_figure_params_old_shape_basics():
    meta = {
        "model_config": {
            "hidden_dim": 128,
            "n_blocks": 3,
            "mode": "ddpm",
            "conditioning": "embedding",
            "router": {"enabled": False},
        },
        "training_epoch": 5,
        "best_val_loss": 0.5,
        "steps": 20,
    }
    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
    assert params == {
        "hidden_dim": 128,
        "n_blocks": 3,
        "mode": "ddpm",
        "conditioning": "embedding",
        "router": "off",
        "epoch": 5,
        "best_val_loss": 0.5,
        "steps": 20,
    }


def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
    meta = {
        "model_config": {"mode": "wgan", "noise_dim": 16},
        "steps": 20,
    }
    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
    assert params["noise_dim"] == 16
    assert "steps" not in params


def test_figure_params_multi_rollout_names_the_series():
    run_meta = {"rollouts": {"flow": {"model_config": {"mode": "flow"}}, "wgan": {"model_config": {"mode": "wgan"}}}}
    assert render_mod._figure_params(run_meta) == {"rollouts": "flow, wgan"}


def test_figure_params_empty_rollouts_is_empty():
    assert render_mod._figure_params({}) == {}
    assert render_mod._figure_params({"rollouts": {}}) == {}


def test_plot_metadata_includes_note_and_run_meta_parameters():
    r = Reduced("u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"})
    meta = render_mod._plot_metadata(r, {"title": "run-1", "reference": "ref.parquet", "rollouts": {"rollout": {}}})
    assert meta["note"] == "no router data"
    assert meta["parameters"] == {"reference": "ref.parquet", "rollouts": {"rollout": {}}}
    assert "title" not in meta["parameters"]


def test_plot_metadata_omits_parameters_when_run_meta_empty():
    r = Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1]})
    meta = render_mod._plot_metadata(r, {})
    assert "parameters" not in meta
    assert "note" not in meta


def test_tex_escape_handles_percent_and_other_special_chars():
    assert render_mod._tex_escape("90% of deposited energy") == r"90\% of deposited energy"
    assert render_mod._tex_escape(r"a_b & c#d $e {f} \bar") == r"a\_b \& c\#d \$e \{f\} \textbackslash{}bar"


def test_render_survives_title_and_xlabel_with_literal_percent(tmp_path: Path):
    # Regression test for gitea #81: a literal "%" in a catalog title (e.g.
    # "Shower containment depth (90% of deposited energy)") crashed the whole
    # LaTeX render, since usetex treats an unescaped "%" as a comment marker.
    reduced = [
        Reduced(
            "shower_containment_depth_90",
            "shower",
            "single_hist",
            "Shower containment depth (90% of deposited energy)",
            "depth containing 90% of deposited energy [mm]",
            {"edges": [0, 1, 2], "series": {"flow": [5, 1]}},
        ),
    ]
    try:
        pdfs = _try_render(reduced, tmp_path)
    except RuntimeError as e:  # LaTeX missing at render time
        pytest.skip(f"LaTeX rendering unavailable: {e}")
    assert len(pdfs) == 1
    assert pdfs[0].exists()