from __future__ import annotations

from pathlib import Path

from scripts.consolidation_ablation import RESULT_COLUMNS, make_param_grid, write_csv


def test_make_param_grid_cartesian_count() -> None:
    seeds = [0]
    rws = [0.0, 0.5, 1.0]
    temps = [0.01, 0.05, 0.1]
    intervals = [500, 2000, 5000]
    steps = [10, 50, 200]
    grid = make_param_grid("corridor", seeds, rws, temps, intervals, steps)
    # 3 * 3 * 3 * 3 = 81 per seed
    assert len(grid) == 81
    # Basic sanity: first and last entries reflect the ranges
    assert grid[0].reward_weight == rws[0]
    assert grid[-1].cons_steps == steps[-1]


def test_write_csv_emits_expected_header(tmp_path: Path) -> None:
    rows = [
        {
            "env": "corridor",
            "seed": 0,
            "reward_weight": 0.5,
            "temperature": 0.05,
            "cons_interval": 2000,
            "cons_steps": 50,
            "loss": 0.0,
            "retrieval_sharpness": 0.0,
            "reward_alignment": 0.0,
            "eval_return": 0.0,
        }
    ]
    out = tmp_path / "ablation_unit.csv"
    out_path = write_csv(out, rows)
    assert out_path.exists()
    content = out_path.read_text().strip().splitlines()
    assert content[0].strip().replace(" ", "") == ",".join(RESULT_COLUMNS)
    assert ",".join(str(v) for v in rows[0].values()) in content[1]
