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  # noqa: E402
from giant.analysis.reduced import Reduced  # noqa: E402


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",
            "router",
            "router_gating",
            "Router gating",
            "pre-step energy [MeV]",
            {
                "n_experts": 2,
                "log_x": True,
                "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]],
                },
            },
        ),
        Reduced(
            "rs",
            "router",
            "router_share",
            "Router share",
            "species",
            {
                "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(
            "ru",
            "router",
            "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: {"rollout": [1, 2], "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], "rollout": [5, 1], "log_x": True, "log_y": True},
        ),
    ]
    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):
    calls = []
    monkeypatch.setattr(
        render_mod.subprocess, "run", lambda *a, **k: calls.append((a, k))
    )
    reduced = [
        Reduced(
            "s",
            "species",
            "single_hist",
            "Single",
            "x",
            {"edges": [0, 1, 2], "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(
        rollout="rollout.parquet",
        reference="reference.parquet",
        run_dir=str(run_dir),
        title="my-run",
        plot_meta={"checkpoint": "ckpt/best.pt"},
    )
    monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))

    Reduced(
        "s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "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],
                "rollout": [1, 2, 3],
                "reference": [3, 2, 1],
                "log_y": False,
            },
        ),
        Reduced(
            "g",
            "marginals",
            "grouped_hist",
            "Grouped",
            "x",
            {
                "edges": [0, 1, 2],
                "groups": {"a": {"rollout": [1, 2], "reference": [2, 1]}},
                "log_y": False,
            },
        ),
        Reduced(
            "p",
            "shower",
            "profile",
            "Profile",
            "depth",
            {
                "edges": [0, 1, 2],
                "rollout_mean": [1, 2],
                "rollout_std": [0.1, 0.2],
                "reference_mean": [1.1, 1.9],
                "reference_std": [0.1, 0.1],
                "ylabel": "e",
            },
        ),
        Reduced(
            "b",
            "species",
            "bar",
            "Bar",
            "species",
            {
                "labels": ["e-", "gamma"],
                "rollout": [0.6, 0.4],
                "reference": [0.5, 0.5],
                "ylabel": "frac",
            },
        ),
        Reduced(
            "s",
            "species",
            "single_hist",
            "Single",
            "x",
            {"edges": [0, 1, 2], "rollout": [5, 1], "log_y": True},
        ),
    ]
    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"}},
    }
    run_meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10}
    params = render_mod._figure_params(run_meta | {"model_config": mc})
    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},
        },
    }
    run_meta = {"model_config": mc, "steps": 10}
    params = render_mod._figure_params(run_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({"model_config": same})

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


def test_figure_params_old_shape_basics():
    run_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(run_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():
    run_meta = {
        "model_config": {"mode": "wgan", "noise_dim": 16},
        "steps": 20,
    }
    params = render_mod._figure_params(run_meta)
    assert params["noise_dim"] == 16
    assert "steps" not in params


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", "checkpoint": "ckpt.pt"})
    assert meta["note"] == "no router data"
    assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
    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