source/src/bgc_retrieval/reporting.py

"""Generate paper-consumable summaries from locked evaluation records."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import pandas as pd

from .artifacts import write_json_immutable
from .statistics import aggregate_draws_by_group, hierarchical_bootstrap_ci


def summarize_results(
    results: pd.DataFrame,
    metrics: list[str],
    bootstrap_samples: int,
    confidence: float,
    seed: int,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    aggregated = aggregate_draws_by_group(results, metrics)
    rows: list[dict[str, Any]] = []
    for method, method_rows in aggregated.groupby("method"):
        for metric in metrics:
            lower, upper = hierarchical_bootstrap_ci(
                method_rows[metric], bootstrap_samples, confidence, seed
            )
            rows.append(
                {
                    "method": method,
                    "metric": metric,
                    "mean": float(method_rows[metric].mean()),
                    "ci_lower": lower,
                    "ci_upper": upper,
                    "groups": int(method_rows["group_id"].nunique()),
                }
            )
    return aggregated, pd.DataFrame(rows)


def write_paper_outputs(
    output_dir: str | Path,
    results: pd.DataFrame,
    metrics: list[str],
    bootstrap_samples: int,
    confidence: float,
    seed: int,
    metadata: dict[str, Any],
) -> None:
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    for name in ("query_results.csv", "group_results.csv", "summary.csv", "metadata.json"):
        if (output / name).exists():
            raise FileExistsError(f"Refusing to overwrite paper output: {output / name}")
    aggregated, summary = summarize_results(
        results, metrics, bootstrap_samples, confidence, seed
    )
    results.to_csv(output / "query_results.csv", index=False)
    aggregated.to_csv(output / "group_results.csv", index=False)
    summary.to_csv(output / "summary.csv", index=False)
    write_json_immutable(output / "metadata.json", metadata)