tests/test_baselines.py

from __future__ import annotations

import subprocess
import sys
from pathlib import Path

from dovla_cil.data.datasets import CILDataset
from dovla_cil.experiments.baselines import (
    BaselineConfig,
    loss_weights_for_baseline,
    prepare_dataset_for_baseline,
)
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.utils.io import read_json


def _make_dataset(tmp_path: Path) -> Path:
    dataset_dir = tmp_path / "cil"
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=dataset_dir,
        num_states_per_task=1,
        k=4,
        seed=9,
        shard_size=8,
        inline_observations=True,
    )
    return dataset_dir


def test_expert_only_dataset_has_one_record_per_group(tmp_path: Path) -> None:
    dataset_dir = _make_dataset(tmp_path)
    prepared = prepare_dataset_for_baseline(
        dataset_dir, "expert_only_bc", tmp_path / "expert_only"
    )
    dataset = CILDataset(prepared)
    assert len(dataset.records) == len(dataset.group_ids)
    assert all(len(dataset.get_group(group_id)) == 1 for group_id in dataset.group_ids)


def test_random_negatives_mode_can_generate(tmp_path: Path) -> None:
    dataset_dir = _make_dataset(tmp_path)
    prepared = prepare_dataset_for_baseline(
        dataset_dir, "random_negatives", tmp_path / "random_negatives"
    )
    dataset = CILDataset(prepared)
    assert any(record.candidate_type == "random_negative" for record in dataset.records)
    assert (prepared / "baseline_metadata.json").exists()


def test_world_model_auxiliary_sets_loss_weights() -> None:
    weights = loss_weights_for_baseline("world_model_auxiliary")
    assert weights.weight("effect") == 1.0
    assert weights.weight("progress") == 1.0
    assert weights.weight("rank") == 0.0
    assert weights.weight("regret") == 0.0


def test_cross_state_baseline_is_measured_not_placeholder(tmp_path: Path) -> None:
    dataset_dir = _make_dataset(tmp_path)
    prepared = prepare_dataset_for_baseline(
        dataset_dir, "cross_state_negatives", tmp_path / "cross_state"
    )

    metadata = read_json(prepared / "baseline_metadata.json")
    assert metadata["approximate"] is False


def test_baseline_config_model_dump(tmp_path: Path) -> None:
    config = BaselineConfig(
        baseline="expert_only_bc",
        dataset=tmp_path / "dataset",
        out=tmp_path / "out",
    )
    payload = config.model_dump()
    assert payload["baseline"] == "expert_only_bc"


def test_baseline_cli_smoke_runs(tmp_path: Path) -> None:
    dataset_dir = _make_dataset(tmp_path)
    out_dir = tmp_path / "run"
    subprocess.run(
        [
            sys.executable,
            "scripts/run_baseline.py",
            "--baseline",
            "expert_only_bc",
            "--dataset",
            str(dataset_dir),
            "--out",
            str(out_dir),
            "--epochs",
            "1",
            "--batch-groups",
            "1",
            "--records-per-group",
            "1",
            "--hidden-dim",
            "32",
            "--eval-num-tasks",
            "2",
            "--eval-k",
            "4",
        ],
        check=True,
        capture_output=True,
        text=True,
    )
    assert (out_dir / "train" / "best.pt").exists()
    metrics = read_json(out_dir / "metrics.json")
    assert metrics["baseline"] == "expert_only_bc"
    assert "pairwise_ranking_accuracy" in metrics["eval"]