workspace/tests/test_dataset_reports.py

from __future__ import annotations

import csv
import subprocess
import sys
from pathlib import Path

from dovla_cil.data.schema import (
    CIL_VERSION,
    ActionChunk,
    CILRecord,
    FailureInfo,
    RewardInfo,
    StructuredEffect,
    make_record_id,
)
from dovla_cil.data.sharding import write_cil_shards
from dovla_cil.experiments.reports import (
    compute_candidate_stats,
    compute_failure_stats,
    generate_dataset_report,
    load_dataset_summary,
)
from dovla_cil.utils.io import read_json


def test_dataset_report_script_runs_and_creates_files(tmp_path: Path) -> None:
    dataset_dir = tmp_path / "dataset"
    out_dir = tmp_path / "report"
    _write_tiny_dataset(dataset_dir)

    result = subprocess.run(
        [
            sys.executable,
            "scripts/report_dataset.py",
            "--dataset",
            str(dataset_dir),
            "--out",
            str(out_dir),
            "--sample-groups",
            "2",
        ],
        check=True,
        text=True,
        capture_output=True,
    )

    expected = {
        "summary.json",
        "candidate_type_counts.csv",
        "reward_histogram.png",
        "success_by_candidate_type.csv",
        "success_by_candidate_type.png",
        "regret_histogram.png",
        "group_size_distribution.png",
        "failure_type_counts.csv",
        "failure_type_counts.png",
        "examples.md",
    }
    assert "num records: 3" in result.stdout
    assert expected.issubset({path.name for path in out_dir.iterdir()})


def test_report_stats_match_expected_toy_dataset(tmp_path: Path) -> None:
    dataset_dir = tmp_path / "dataset"
    out_dir = tmp_path / "report"
    records = _write_tiny_dataset(dataset_dir)

    summary = generate_dataset_report(dataset_dir, out_dir, sample_groups=2, seed=4)
    candidate_stats = {
        row["candidate_type"]: row for row in compute_candidate_stats(records)
    }
    failure_stats = {row["failure_type"]: row for row in compute_failure_stats(records)}
    summary_file = read_json(out_dir / "summary.json")
    candidate_counts = _read_csv(out_dir / "candidate_type_counts.csv")

    assert summary["num_records"] == 3
    assert summary["num_groups"] == 2
    assert summary["candidate_type_counts"] == {"expert": 2, "wrong_target": 1}
    assert summary["failure_type_counts"] == {"success": 2, "wrong_target": 1}
    assert summary_file["success_rate"] == 2 / 3
    assert candidate_stats["expert"]["success_rate"] == 1.0
    assert candidate_stats["wrong_target"]["success_rate"] == 0.0
    assert failure_stats["wrong_target"]["count"] == 1
    assert candidate_counts == [
        {"candidate_type": "expert", "count": "2"},
        {"candidate_type": "wrong_target", "count": "1"},
    ]
    assert "| record_id | candidate_type | reward.progress |" in (
        out_dir / "examples.md"
    ).read_text(encoding="utf-8")


def test_load_dataset_summary_function(tmp_path: Path) -> None:
    dataset_dir = tmp_path / "dataset"
    _write_tiny_dataset(dataset_dir)

    summary = load_dataset_summary(dataset_dir)

    assert summary["num_records"] == 3
    assert summary["reward"]["min"] == 0.2
    assert summary["reward"]["max"] == 1.0
    assert summary["group_size"]["max"] == 2


def _write_tiny_dataset(dataset_dir: Path) -> list[CILRecord]:
    records = [
        _record("group-a", 0, "expert", 1.0, True, "success", regret=0.0, rank=0),
        _record("group-a", 1, "wrong_target", 0.2, False, "wrong_target", regret=0.8, rank=1),
        _record("group-b", 0, "expert", 1.0, True, "success", regret=0.0, rank=0),
    ]
    write_cil_shards(
        records,
        output_dir=dataset_dir,
        max_records_per_shard=10,
        dataset_name="report_toy",
        backend="toy",
        k=2,
        task_count=1,
        seed=5,
    )
    return records


def _record(
    group_id: str,
    index: int,
    candidate_type: str,
    reward_progress: float,
    success: bool,
    failure_type: str,
    *,
    regret: float,
    rank: int,
) -> CILRecord:
    action = ActionChunk(
        action_id=f"{candidate_type}-{index}",
        representation="semantic",
        horizon=1,
        values=[{"command": "noop" if not success else "grasp", "object": "red_mug"}],
        skill_type="grasp",
        metadata={"candidate_type": candidate_type, "intended_target": "red_mug"},
    )
    return CILRecord(
        version=CIL_VERSION,
        record_id=make_record_id(group_id, action.action_id, seed=9),
        group_id=group_id,
        state_hash=f"state-{group_id}",
        task_id="toy_report_task",
        scene_id=f"scene-{group_id}",
        instruction="Pick the red mug.",
        instruction_family={"family": "pick"},
        observation_ref=None,
        observation_inline={"symbolic_state": {"objects": {}}},
        action_chunk=action,
        next_observation_ref=None,
        next_observation_inline={"symbolic_state": {"objects": {}}},
        structured_effect=StructuredEffect(
            object_pose_delta={"red_mug": [reward_progress, 0.0, 0.0]},
            relation_after={"grasped(red_mug)": success},
            moved_objects=["red_mug"] if reward_progress > 0 else [],
        ),
        reward=RewardInfo(
            progress=reward_progress,
            success=success,
            terminal_success=success,
            dense_components={"progress": reward_progress},
        ),
        regret=regret,
        rank_within_group=rank,
        candidate_type=candidate_type,
        failure=FailureInfo(
            type=failure_type,
            symbolic_reason=failure_type,
            language_explanation=failure_type,
        ),
    )


def _read_csv(path: Path) -> list[dict[str, str]]:
    with path.open("r", encoding="utf-8", newline="") as handle:
        return list(csv.DictReader(handle))