dataforge/cli/bench.py

"""CLI subcommand: ``dataforge bench``."""

from __future__ import annotations

import json
from collections.abc import Callable
from pathlib import Path
from typing import Annotated, Any

import typer
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

_console = Console(stderr=True)
run_agent_comparison: Callable[..., Any] | None = None


def _parse_csv_list(raw_value: str) -> list[str]:
    """Parse a comma-separated CLI option into a list of strings."""
    values = [item.strip() for item in raw_value.split(",")]
    return [value for value in values if value]


def _parse_seed_list(raw_value: str | None) -> list[int] | None:
    """Parse an optional comma-separated seed list."""
    if raw_value is None:
        return None
    seeds = [item.strip() for item in raw_value.split(",") if item.strip()]
    return [int(seed) for seed in seeds]


def _runner() -> Callable[..., Any]:
    """Load the benchmark runner lazily so core CLI imports stay lightweight."""
    global run_agent_comparison
    if run_agent_comparison is None:
        from dataforge.bench.runner import run_agent_comparison as loaded_runner

        run_agent_comparison = loaded_runner
    return run_agent_comparison


def _render_summary(output: Any) -> None:
    """Render the aggregate F1/quota summary table."""
    table = Table(title="DataForge Benchmark Summary")
    table.add_column("Method")
    table.add_column("Dataset")
    table.add_column("Status")
    table.add_column("F1")
    table.add_column("Avg Steps")
    table.add_column("Quota")
    for aggregate in output.aggregates:
        table.add_row(
            aggregate.method,
            aggregate.dataset,
            aggregate.status,
            "Skipped" if aggregate.f1_mean is None else f"{aggregate.f1_mean:.4f}",
            "Skipped" if aggregate.avg_steps_mean is None else f"{aggregate.avg_steps_mean:.2f}",
            "Skipped"
            if aggregate.quota_units_mean is None
            else f"{aggregate.quota_units_mean:.4f}",
        )
    Console().print(table)
    if any(aggregate.status == "skipped" for aggregate in output.aggregates):
        Console().print(
            Panel(
                "Some LLM baselines were skipped. Set DATAFORGE_LLM_PROVIDER=groq and "
                "GROQ_API_KEY to enable them.",
                title="Benchmark Warning",
                style="yellow",
            )
        )


def _render_coverage_matrix(output: Any) -> None:
    """Render the per-error-class recall matrix (honest coverage view)."""
    from dataforge.bench.error_classes import BENCH_ERROR_CLASSES, class_coverage_matrix

    matrix = class_coverage_matrix(list(output.records))
    if not matrix:
        return
    table = Table(title="Per-Error-Class Detection Recall (coverage)")
    table.add_column("Method/Dataset")
    for error_class in BENCH_ERROR_CLASSES:
        table.add_column(error_class)
    for (method, dataset), scores in sorted(matrix.items()):
        row = [f"{method}/{dataset}"]
        for error_class in BENCH_ERROR_CLASSES:
            score = scores.get(error_class)
            if score is None or score.support == 0:
                row.append("-")
            else:
                row.append(f"{score.detection_recall:.2f}/{score.recall:.2f} (n={score.support})")
        table.add_row(*row)
    Console().print(table)
    Console().print(
        Panel(
            "Each cell shows detection_recall/correction_recall per error class on "
            "the full RAHA datasets ('-' = no ground-truth cells of that class). "
            "Detection credits flagging the error; correction credits producing the "
            "exact right value. The honest split: a class can be well-detected yet "
            "not auto-correctable (no derivable value), which is reported, not hidden.",
            title="Coverage",
            style="cyan",
        )
    )


def bench(
    methods: Annotated[
        str,
        typer.Option(
            "--methods",
            help="Comma-separated benchmark methods.",
        ),
    ] = "heuristic,llm_zeroshot",
    datasets: Annotated[
        str,
        typer.Option(
            "--datasets",
            help="Comma-separated benchmark datasets.",
        ),
    ] = "hospital",
    seeds: Annotated[
        int,
        typer.Option("--seeds", help="Number of seeds per method/dataset pair."),
    ] = 3,
    seed_list: Annotated[
        str | None,
        typer.Option(
            "--seed-list",
            help="Explicit comma-separated seed list. Overrides --seeds for reproducibility.",
        ),
    ] = None,
    really_run_big_bench: Annotated[
        bool,
        typer.Option(
            "--really-run-big-bench",
            help="Override the free-tier benchmark quota guard when estimated calls exceed 500.",
        ),
    ] = False,
    output_json: Annotated[
        Path,
        typer.Option(
            "--output-json",
            help="Where to write eval/results/agent_comparison.json.",
        ),
    ] = Path("eval/results/agent_comparison.json"),
    cache_root: Annotated[
        Path | None,
        typer.Option(
            "--cache-root",
            help="Benchmark dataset cache root. Defaults to the user DataForge cache.",
        ),
    ] = None,
    verify_dataset_hashes: Annotated[
        bool,
        typer.Option(
            "--verify-dataset-hashes/--no-verify-dataset-hashes",
            help="Verify cached benchmark bytes against pinned upstream hashes.",
        ),
    ] = True,
    json_output: Annotated[
        bool,
        typer.Option("--json", help="Print benchmark results as JSON."),
    ] = False,
    quick: Annotated[
        bool,
        typer.Option(
            "--quick",
            help="Offline coverage check: run random,heuristic on all datasets (1 seed) "
            "and print the per-error-class coverage matrix. No API keys required.",
        ),
    ] = False,
) -> None:
    """Run real-world benchmark methods across cached benchmark datasets."""
    if quick:
        methods = "random,heuristic"
        datasets = datasets if datasets != "hospital" else "hospital,flights,beers"
        seeds = 1
        seed_list = None
    try:
        output = _runner()(
            methods=_parse_csv_list(methods),
            datasets=_parse_csv_list(datasets),
            seeds=seeds,
            seed_list=_parse_seed_list(seed_list),
            output_json=output_json,
            really_run_big_bench=really_run_big_bench,
            cache_root=cache_root,
            verify_dataset_hashes=verify_dataset_hashes,
        )
    except Exception as exc:
        _console.print(
            Panel(
                f"[bold red]{exc}[/bold red]",
                title="Benchmark Error",
                style="red",
            )
        )
        raise typer.Exit(code=2) from exc

    if json_output:
        typer.echo(json.dumps(output.model_dump(mode="json"), indent=2, sort_keys=True))
        return

    _render_summary(output)
    if quick:
        _render_coverage_matrix(output)