giant/tests/test_training_plots.py

"""Tests for giant.training.plots (gitea #75) — render smoke tests skipped
where plotstyle/LaTeX is unavailable, plus pure-function column-classification
coverage that needs neither."""

from __future__ import annotations

import csv
from pathlib import Path

import pytest

pytest.importorskip("plotstyle")

from giant.training import plots as plots_mod  # noqa: E402
from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics  # noqa: E402

# --- fixtures ----------------------------------------------------------


_RICH_HEADER = [
    "epoch",
    "stage1/train/loss",
    "stage1/train/loss_gen",
    "stage1/train/nsec_acc",
    "stage1/train/grad_norm",
    "stage1/val/loss",
    "stage1/val/loss_gen",
    "stage1/val/nsec_acc",
    "stage1/lr",
    "stage1/router/entropy",
    "stage1/router/util_min",
    "stage1/router/util_max",
    "stage1/router/util_std",
    "stage2/train/d_loss",
    "stage2/train/g_loss",
    "stage2/train/wasserstein",
    "stage2/train/gp_loss",
    "stage2/train/loss_nsec",
    "stage2/train/nsec_acc",
    "stage2/train/grad_norm_d",
    "stage2/train/grad_norm_g",
    "stage2/lr",
    "stage2/critic_lr",
    "val/loss",
    "val/marginal_kl",
    "grad_norm",
    "gpu_mem_mb",
    "samples_per_sec",
    "is_best",
    "epoch_time_s",
]

_RICH_ROWS = [
    [
        1,
        1.0,
        0.8,
        0.5,
        1.2,
        0.9,
        0.7,
        0.6,
        3e-4,
        1.5,
        0.05,
        0.3,
        0.1,
        -0.2,
        0.3,
        0.5,
        0.1,
        0.4,
        0.4,
        0.9,
        1.1,
        3e-4,
        1e-4,
        0.85,
        0.4,
        2.1,
        512.0,
        100.0,
        1,
        5.0,
    ],
    [
        2,
        0.8,
        0.6,
        0.6,
        1.0,
        0.7,
        0.5,
        0.7,
        2e-4,
        1.6,
        0.06,
        0.28,
        0.09,
        -0.1,
        0.25,
        0.4,
        0.09,
        0.3,
        0.5,
        0.8,
        1.0,
        2e-4,
        8e-5,
        0.7,
        0.35,
        1.9,
        520.0,
        105.0,
        0,
        5.1,
    ],
]

_MINIMAL_HEADER = [
    "epoch",
    "stage1/train/loss",
    "stage1/train/loss_gen",
    "stage1/val/loss",
    "stage1/val/loss_gen",
    "stage1/lr",
    "val/loss",
    "grad_norm",
    "gpu_mem_mb",
    "samples_per_sec",
    "is_best",
    "epoch_time_s",
]

_MINIMAL_ROWS = [
    [1, 1.0, 0.8, 0.9, 0.7, 3e-4, 0.85, 0.4, 0.0, 100.0, 0, 5.0],
    [2, 0.8, 0.6, 0.7, 0.5, 2e-4, 0.7, 0.35, 0.0, 105.0, 1, 5.1],
]


def _write_csv(path: Path, header: list[str], rows: list[list]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(header)
        writer.writerows(rows)


# --- MetricsTable --------------------------------------------------------


def test_metrics_table_load_round_trips(tmp_path: Path):
    csv_path = tmp_path / "metrics.csv"
    _write_csv(csv_path, _MINIMAL_HEADER, _MINIMAL_ROWS)
    table = MetricsTable.load(csv_path)
    assert table.epochs == [1, 2]
    assert table.columns["stage1/train/loss"] == [1.0, 0.8]
    assert "epoch" not in table.columns
    assert table.best_epochs() == [2]


def test_metrics_table_best_epochs_empty_without_is_best_column():
    table = MetricsTable(epochs=[1, 2], columns={"stage1/train/loss": [1.0, 0.5]})
    assert table.best_epochs() == []


# --- column classification (pure functions, no matplotlib) --------------


def _rich_columns() -> dict[str, list]:
    return {name: [0.0] for name in _RICH_HEADER if name != "epoch"}


def test_stages_detects_only_stages_present():
    assert plots_mod._stages(_rich_columns()) == ["stage1", "stage2"]
    assert plots_mod._stages({"stage2/train/loss": [0.0]}) == ["stage2"]
    assert plots_mod._stages({"val/loss": [0.0]}) == []


def test_split_matches_stage_and_split_prefix_only():
    cols = _rich_columns()
    train = plots_mod._split(cols, "stage1", "train")
    assert train == {
        "loss": "stage1/train/loss",
        "loss_gen": "stage1/train/loss_gen",
        "nsec_acc": "stage1/train/nsec_acc",
        "grad_norm": "stage1/train/grad_norm",
    }
    assert plots_mod._split(cols, "stage2", "val") == {}


def test_point_in_time_excludes_train_val_router():
    cols = _rich_columns()
    pit = plots_mod._point_in_time(cols, "stage1")
    assert pit == {"lr": "stage1/lr"}
    pit2 = plots_mod._point_in_time(cols, "stage2")
    assert pit2 == {"lr": "stage2/lr", "critic_lr": "stage2/critic_lr"}


def test_router_columns():
    cols = _rich_columns()
    assert plots_mod._router(cols, "stage1") == {
        "entropy": "stage1/router/entropy",
        "util_min": "stage1/router/util_min",
        "util_max": "stage1/router/util_max",
        "util_std": "stage1/router/util_std",
    }
    assert plots_mod._router(cols, "stage2") == {}


def test_run_level_excludes_stage_prefixed_columns_including_val_loss_lookalike():
    cols = _rich_columns()
    run_level = plots_mod._run_level(cols)
    assert set(run_level) == {
        "val/loss",
        "val/marginal_kl",
        "grad_norm",
        "gpu_mem_mb",
        "samples_per_sec",
        "is_best",
        "epoch_time_s",
    }
    # stage-prefixed "val/loss" lookalike (stage1/val/loss) must not leak in
    assert "stage1/val/loss" not in run_level


def test_loss_keys_excludes_acc_and_wgan_and_grad_norm():
    train = {"loss": "x", "loss_gen": "x", "nsec_acc": "x", "grad_norm": "x", "d_loss": "x"}
    val = {"loss": "x", "loss_gen": "x"}
    assert plots_mod._loss_keys(train, val) == ["loss", "loss_gen"]


# --- derive_metrics_dir ---------------------------------------------------


def test_derive_metrics_dir_explicit_out_dir_wins():
    assert derive_metrics_dir("runs/my-run", out_dir="/somewhere") == Path("/somewhere")


def test_derive_metrics_dir_default_base():
    assert derive_metrics_dir("runs/my-run", default_base="/data/analysis_runs") == Path(
        "/data/analysis_runs/metrics_my-run"
    )


def test_derive_metrics_dir_falls_back_to_cwd_analysis_runs(monkeypatch, tmp_path):
    monkeypatch.chdir(tmp_path)
    assert derive_metrics_dir("runs/my-run") == tmp_path / "analysis_runs" / "metrics_my-run"


# --- render_metrics end to end -------------------------------------------


def _try_render(run_dir: Path, out_dir: Path) -> list[Path]:
    try:
        return render_metrics(run_dir, out_dir)
    except RuntimeError as e:  # LaTeX missing at render time
        pytest.skip(f"LaTeX rendering unavailable: {e}")


def test_render_metrics_rich_run_produces_expected_plots_outside_run_dir(tmp_path: Path):
    run_dir = tmp_path / "run"
    out_dir = tmp_path / "out"
    _write_csv(run_dir / "metrics.csv", _RICH_HEADER, _RICH_ROWS)

    paths = _try_render(run_dir, out_dir)

    names = {p.stem for p in paths}
    assert names == {
        "overview",
        "stage1_loss",
        "stage2_loss",
        "lr",
        "stage1_accuracy",
        "stage2_accuracy",
        "grad_norm",
        "stage1_router",
        "stage2_wgan_balance",
        "throughput",
    }
    assert all(p.exists() for p in paths)
    assert all(p.is_relative_to(out_dir) for p in paths)
    # nothing written into the training run directory itself
    assert not any(run_dir.rglob("*.pdf"))


def test_render_metrics_minimal_run_omits_router_wgan_accuracy(tmp_path: Path):
    run_dir = tmp_path / "run"
    out_dir = tmp_path / "out"
    _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)

    paths = _try_render(run_dir, out_dir)

    names = {p.stem for p in paths}
    assert names == {"overview", "stage1_loss", "lr", "grad_norm", "throughput"}
    assert "stage1_accuracy" not in names
    assert "stage1_router" not in names
    assert "stage1_wgan_balance" not in names


def test_render_metrics_default_out_dir_uses_analysis_runs_convention(tmp_path: Path):
    run_dir = tmp_path / "runs" / "my-run"
    _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
    default_base = tmp_path / "analysis_runs"

    try:
        paths = render_metrics(run_dir, default_base=default_base)
    except RuntimeError as e:
        pytest.skip(f"LaTeX rendering unavailable: {e}")

    assert paths
    assert all(p.is_relative_to(default_base / "metrics_my-run") for p in paths)


# --- CLI -------------------------------------------------------------------


def test_cli_analyze_metrics_smoke(tmp_path: Path):
    from typer.testing import CliRunner

    from giant.cli import app

    run_dir = tmp_path / "run"
    out_dir = tmp_path / "out"
    _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)

    runner = CliRunner()
    result = runner.invoke(app, ["analyze", "metrics", str(run_dir), "--out", str(out_dir)])

    if result.exit_code != 0 and "LaTeX" in str(result.output):
        pytest.skip("LaTeX rendering unavailable")
    assert result.exit_code == 0, result.output
    assert any(out_dir.glob("*.pdf"))