giant/tests/test_analysis_run.py

"""Tests for the rollout-YAML(s) → run-directory flow, compute, and merge."""

from __future__ import annotations

from pathlib import Path

import pyarrow.parquet as pq
import pytest
import yaml

from giant.analysis import (
    RunMeta,
    compute_one,
    compute_reduced,
    derive_run_dir,
    load_rollout_yaml,
    load_rollout_yamls,
    merge_one,
    prep,
)
from giant.analysis.run import Context
from giant.analysis.reduced import Partial, Reduced
from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
from tests.test_analysis_reduce import _reference_frame, _rollout_frame


def _write_rollout(path: Path) -> None:
    tbl = _rollout_frame().collect().to_arrow()
    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
    pq.write_table(tbl, path)


def _write_inputs(tmp_path: Path) -> Path:
    """Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
    rollout = tmp_path / "rollout.parquet"
    reference = tmp_path / "reference.parquet"
    _write_rollout(rollout)
    _reference_frame().collect().write_parquet(reference)

    yaml_path = tmp_path / "run.yaml"
    yaml_path.write_text(
        yaml.safe_dump(
            {
                "prediction_id": "abcd1234ef",
                "output": str(rollout),
                "dataset": str(reference),
                "checkpoint": "/ckpt/best.pt",
                "kind": "rollout",
                "energy_cutoff": 0.1,
                "steps": 10,
            }
        )
    )
    return yaml_path


def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
    """Two rollout YAMLs (distinct output files) sharing one reference file."""
    reference = tmp_path / "reference.parquet"
    _reference_frame().collect().write_parquet(reference)

    paths = []
    for tag, pred_id in (("a", "aaaa1111ef"), ("b", "bbbb2222ef")):
        rollout = tmp_path / f"rollout_{tag}.parquet"
        _write_rollout(rollout)
        yaml_path = tmp_path / f"run_{tag}.yaml"
        yaml_path.write_text(
            yaml.safe_dump(
                {
                    "prediction_id": pred_id,
                    "output": str(rollout),
                    "dataset": str(reference),
                    "checkpoint": f"/ckpt/{tag}.pt",
                    "kind": "rollout",
                    "energy_cutoff": 0.1,
                    "steps": 10,
                }
            )
        )
        paths.append(yaml_path)
    return paths[0], paths[1]


def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
    """``prep`` with small test-sized context bins/sampling."""
    return prep(
        rollout_yamls,
        run_dir,
        n_chunks=chunks,
        labels=labels,
        n_energy_bins=2,
        n_marginal_bins=8,
        top_k_pdg=3,
        sample_rows=1000,
    )


def test_load_rollout_yaml_requires_paths(tmp_path: Path):
    bad = tmp_path / "bad.yaml"
    bad.write_text(yaml.safe_dump({"output": "x.parquet"}))  # no dataset
    with pytest.raises(ValueError):
        load_rollout_yaml(bad)


def test_load_rollout_yamls_single_defaults_to_rollout_name(tmp_path: Path):
    yaml_path = _write_inputs(tmp_path)
    loaded, reference = load_rollout_yamls([yaml_path])
    assert [lr.name for lr in loaded] == ["rollout"]
    assert reference.endswith("reference.parquet")


def test_load_rollout_yamls_multi_defaults_to_stem(tmp_path: Path):
    a, b = _write_two_inputs(tmp_path)
    loaded, _ = load_rollout_yamls([a, b])
    assert [lr.name for lr in loaded] == ["run_a", "run_b"]


def test_load_rollout_yamls_explicit_labels(tmp_path: Path):
    a, b = _write_two_inputs(tmp_path)
    loaded, _ = load_rollout_yamls([a, b], labels=["flow", "wgan"])
    assert [lr.name for lr in loaded] == ["flow", "wgan"]


def test_load_rollout_yamls_label_count_mismatch(tmp_path: Path):
    a, b = _write_two_inputs(tmp_path)
    with pytest.raises(ValueError, match="--label"):
        load_rollout_yamls([a, b], labels=["only-one"])


def test_load_rollout_yamls_rejects_duplicate_names(tmp_path: Path):
    a, b = _write_two_inputs(tmp_path)
    with pytest.raises(ValueError, match="collide"):
        load_rollout_yamls([a, b], labels=["same", "same"])


def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
    a, _ = _write_two_inputs(tmp_path)
    other_ref = tmp_path / "other_reference.parquet"
    _reference_frame().collect().write_parquet(other_ref)
    c = tmp_path / "run_c.yaml"
    c.write_text(
        yaml.safe_dump(
            {"prediction_id": "cccc3333ef", "output": str(tmp_path / "rollout_c.parquet"), "dataset": str(other_ref)}
        )
    )
    _write_rollout(tmp_path / "rollout_c.parquet")
    with pytest.raises(ValueError, match="same reference"):
        load_rollout_yamls([a, c])


def test_derive_run_dir_next_to_rollout():
    y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
    assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")
    assert derive_run_dir([y], "/somewhere") == Path("/somewhere")


def test_derive_run_dir_default_base():
    y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
    assert derive_run_dir([y], default_base="/work/lbogner/giant2/analysis_runs") == Path(
        "/work/lbogner/giant2/analysis_runs/analysis_abcd1234"
    )
    # an explicit run_dir still wins over default_base
    assert derive_run_dir([y], "/somewhere", default_base="/other") == Path("/somewhere")


def test_derive_run_dir_multi_rollout_joins_tags():
    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(2)]
    assert derive_run_dir(ys, default_base="/base") == Path("/base/analysis_tag0xxxx-tag1xxxx")


def test_derive_run_dir_many_rollouts_truncates_with_plus_count():
    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(5)]
    run_dir = derive_run_dir(ys, default_base="/base")
    assert run_dir == Path("/base/analysis_tag0xxxx-tag1xxxx-tag2xxxx-plus2")


def test_prep_lays_out_run_dir(tmp_path: Path):
    yaml_path = _write_inputs(tmp_path)
    run_dir = _prep([yaml_path])
    assert run_dir == tmp_path / "analysis_abcd1234"
    assert (run_dir / "shared.json").exists()
    ctx = Context.load(run_dir / "shared.json")
    assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
    meta = RunMeta.load(run_dir / "run_meta.json")
    assert meta.reference.endswith("reference.parquet")
    assert [ro["name"] for ro in meta.rollouts] == ["rollout"]
    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/best.pt"
    assert "best.pt" in meta.title
    assert meta.n_chunks == 1
    assert meta.rows_per_chunk == [meta.total_rows]  # single chunk holds everything
    assert meta.total_rows == 8  # 5 rollout rows + 3 reference rows


def test_prep_multi_rollout_lays_out_run_dir(tmp_path: Path):
    a, b = _write_two_inputs(tmp_path)
    run_dir = _prep([a, b], labels=["flow", "wgan"])
    meta = RunMeta.load(run_dir / "run_meta.json")
    assert [ro["name"] for ro in meta.rollouts] == ["flow", "wgan"]
    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/a.pt"
    assert meta.rollouts[1]["plot_meta"]["checkpoint"] == "/ckpt/b.pt"
    # 5 rows from each rollout + 3 from the shared reference
    assert meta.total_rows == 13


def test_prep_splits_rows_per_chunk(tmp_path: Path):
    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
    meta = RunMeta.load(run_dir / "run_meta.json")
    assert len(meta.rows_per_chunk) == 2
    assert sum(meta.rows_per_chunk) == meta.total_rows == 8


def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Path):
    """Re-prepping with a different n_chunks must not leave old chunk
    partials on disk for merge_one to silently merge against the new
    context (they'd be keyed/sized for the old n_chunks)."""
    yaml_path = _write_inputs(tmp_path)
    run_dir = _prep([yaml_path], chunks=2)
    compute_one("marginal_edep", run_dir, chunk_index=0)
    compute_one("marginal_edep", run_dir, chunk_index=1)
    stale = run_dir / "reduced_partial" / "marginal_edep__0.json"
    assert stale.exists()
    (run_dir / "reduced").mkdir(exist_ok=True)
    (run_dir / "reduced" / "marginal_edep.json").write_text("{}")

    _prep([yaml_path], run_dir, chunks=1)

    assert not stale.exists()
    assert not (run_dir / "reduced" / "marginal_edep.json").exists()
    assert (run_dir / "shared.json").exists()  # prep's own fresh output untouched


def test_compute_one_from_run_dir(tmp_path: Path):
    run_dir = _prep([_write_inputs(tmp_path)])
    out = compute_one("marginal_edep", run_dir)
    assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
    partial = Partial.load(out)
    assert partial.id == "marginal_edep" and partial.chunk == 0
    assert "r" in partial.data and "t" in partial.data
    assert list(partial.data["r"]) == ["rollout"]


def test_compute_reduced_explicit_paths(tmp_path: Path):
    run_dir = _prep([_write_inputs(tmp_path)])
    meta = RunMeta.load(run_dir / "run_meta.json")
    rollouts = [{"name": ro["name"], "path": ro["path"]} for ro in meta.rollouts]
    out = compute_reduced(
        "marginal_step_length",
        rollouts,
        meta.reference,
        run_dir / "shared.json",
        tmp_path / "r.json",
    )
    assert Partial.load(out).id == "marginal_step_length"


def test_merge_one_produces_reduced(tmp_path: Path):
    run_dir = _prep([_write_inputs(tmp_path)])
    compute_one("marginal_edep", run_dir)
    out = merge_one("marginal_edep", run_dir)
    assert out == run_dir / "reduced" / "marginal_edep.json"
    reduced = Reduced.load(out)
    assert reduced.id == "marginal_edep"
    assert len(reduced.payload["series"]["rollout"]) == len(reduced.payload["edges"]) - 1


def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
    compute_one("marginal_edep", run_dir, chunk_index=0)  # chunk 1 never computed
    with pytest.raises(FileNotFoundError, match="missing chunk"):
        merge_one("marginal_edep", run_dir)


def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
    (tmp_path / "a").mkdir()
    (tmp_path / "b").mkdir()
    unchunked_dir = _prep([_write_inputs(tmp_path / "a")])
    compute_one("marginal_step_length", unchunked_dir)
    unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))

    chunked_dir = _prep([_write_inputs(tmp_path / "b")], chunks=2)
    for k in range(2):
        compute_one("marginal_step_length", chunked_dir, chunk_index=k)
    chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))

    assert chunked.payload == unchunked.payload


def test_two_rollout_compute_and_merge_produces_both_series(tmp_path: Path):
    a, b = _write_two_inputs(tmp_path)
    run_dir = _prep([a, b], labels=["flow", "wgan"])
    compute_one("marginal_edep", run_dir)
    reduced = Reduced.load(merge_one("marginal_edep", run_dir))
    assert list(reduced.payload["series"]) == ["flow", "wgan"]
    assert "reference" in reduced.payload


def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
    run_dir = _prep([_write_inputs(tmp_path)])  # n_chunks=1 (default)
    with pytest.raises(ValueError, match="out of range"):
        compute_one("marginal_edep", run_dir, chunk_index=1)


def test_estimate_runtime_s_scales_with_rows_and_margin():
    from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
    from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S

    assert RUNTIME_SAFETY_MARGIN > 0
    small = estimate_runtime_s("marginal_edep", 1_000)
    large = estimate_runtime_s("marginal_edep", 100_000_000)
    assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S
    assert large > small  # bigger chunk -> longer estimate