evaluate.py

#!/usr/bin/env python
"""Evaluate BrainRL baselines with optional condition + subject splits.

Examples
--------

Evaluate every condition / every subject:

    python evaluate.py --episodes 32

Evaluate the test subjects on the single-male-narrator condition only:

    python evaluate.py \\
      --episodes 32 \\
      --condition single_m \\
      --participant-info configs/participant_run_info.json \\
      --train-subjects sub-01:sub-20 \\
      --test-subjects sub-21:sub-26 \\
      --split test \\
      --output-csv outputs/eval/single_m_test.csv \\
      --plot-dir outputs/eval/single_m_test
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

from baselines import (
    EpisodeResult,
    PromptPolicy,
    default_baselines,
    run_policy_episode,
    r2_curves_by_policy,
    summarize_results,
)
from data_split import DEFAULT_PARTICIPANT_INFO, build_condition_split


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Evaluate BrainRL policies")
    parser.add_argument("--episodes", type=int, default=8, help="Episodes per (subject, run) pair per policy")
    parser.add_argument("--seed", type=int, default=42, help="Base random seed")
    parser.add_argument(
        "--condition",
        choices=("single_m", "single_f", "mixed_m", "mixed_f"),
        default=None,
        help="Optional condition filter (uses participant_run_info.json).",
    )
    parser.add_argument(
        "--participant-info",
        type=str,
        default=str(DEFAULT_PARTICIPANT_INFO),
        help="Path to participant_run_info.json.",
    )
    parser.add_argument(
        "--train-subjects",
        type=str,
        default=None,
        help="Subject spec for the train split (e.g. sub-01:sub-20).",
    )
    parser.add_argument(
        "--test-subjects",
        type=str,
        default=None,
        help="Subject spec for the test split (e.g. sub-21:sub-26).",
    )
    parser.add_argument(
        "--split",
        choices=("train", "test", "all"),
        default="all",
        help="Which subject split to evaluate.",
    )
    parser.add_argument(
        "--exclude-subjects",
        type=str,
        default=None,
        help=(
            "Comma list / range of subjects to drop from both train and test "
            "splits, e.g. 'sub-03,sub-18' for corrupted recordings."
        ),
    )
    parser.add_argument(
        "--output-csv",
        type=str,
        default=None,
        help="Optional CSV path for summary rows.",
    )
    parser.add_argument(
        "--plot-dir",
        type=str,
        default=None,
        help="Optional directory for baseline_comparison.png and r2_curves.png.",
    )
    parser.add_argument(
        "--use-llm",
        action="store_true",
        help="Add prompt-based LLM policy to the OpenEnv comparison.",
    )
    return parser


def print_summary(rows: list[dict[str, object]], header: str) -> None:
    bar = "=" * 92
    print(bar)
    print(header)
    print(bar)
    for row in rows:
        print(
            f"{row['policy']:>14} | "
            f"episodes={row['episodes']} | "
            f"mean_final_r2={float(row['mean_final_r2']):.4f} | "
            f"corr={float(row['mean_priority_correlation']):.4f} | "
            f"2v2={float(row['mean_2v2_accuracy']):.4f} | "
            f"mean_total_reward={float(row['mean_total_reward']):.4f} | "
            f"order={row['example_order']}"
        )
    print(bar)


def write_summary(path: Path, rows: list[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def _episode_pairs(args: argparse.Namespace) -> tuple[list[dict[str, str | None]], str]:
    """Resolve which (subject, run, condition) pairs to roll out."""

    if args.condition is None:
        return [{"subject_id": None, "run_id": None, "condition": None}], "all-conditions"

    split = build_condition_split(
        condition=args.condition,
        participant_info_path=args.participant_info,
        train_subjects_spec=args.train_subjects,
        test_subjects_spec=args.test_subjects,
        exclude_subjects=args.exclude_subjects,
    )
    pairs = split.pairs_for(args.split)
    if not pairs:
        raise SystemExit(
            f"No (subject, run) pairs available for condition={args.condition} split={args.split}."
        )
    summary = split.summary()
    excluded = summary.get("excluded_subjects") or []
    print(
        f"[split] condition={args.condition} split={args.split} "
        f"train_subjects={summary['n_train_subjects']} test_subjects={summary['n_test_subjects']} "
        f"selected_pairs={len(pairs)}"
        + (f" excluded={excluded}" if excluded else "")
    )
    return [pair.as_dict() for pair in pairs], f"{args.condition}/{args.split}"


def main() -> None:
    args = build_parser().parse_args()
    pairs, label = _episode_pairs(args)
    policies = default_baselines(seed=int(args.seed))
    if args.use_llm:
        policies.append(PromptPolicy(use_llm=True))

    results: list[EpisodeResult] = []
    for policy in policies:
        for pair_idx, pair in enumerate(pairs):
            for episode_idx in range(int(args.episodes)):
                seed = int(args.seed) + pair_idx * 1009 + episode_idx
                results.append(
                    run_policy_episode(
                        policy=policy,
                        seed=seed,
                        subject_id=pair.get("subject_id"),
                        run_id=pair.get("run_id"),
                        condition=pair.get("condition"),
                    )
                )

    rows = summarize_results(results, split_label=label)
    print_summary(rows, header=f"BrainRL policy comparison ({label})")

    if args.output_csv:
        out_csv = Path(args.output_csv).expanduser()
        write_summary(out_csv, rows)
        print(f"Wrote summary CSV: {out_csv}")

    if args.plot_dir:
        from plotting import plot_baseline_comparison, plot_r2_curves

        plot_dir = Path(args.plot_dir).expanduser()
        plot_dir.mkdir(parents=True, exist_ok=True)
        bar_path = plot_baseline_comparison(rows, plot_dir / "baseline_comparison.png")
        curves_path = plot_r2_curves(
            r2_curves_by_policy(results),
            plot_dir / "r2_curves.png",
            title=f"BrainRL R² curves ({label})",
        )
        meta_path = plot_dir / "split_summary.json"
        with meta_path.open("w", encoding="utf-8") as handle:
            json.dump({"label": label, "n_pairs": len(pairs), "pairs": pairs}, handle, indent=2)
        print(
            f"Wrote plots to {plot_dir} "
            "(baseline_comparison.png, r2_curves.png)"
        )


if __name__ == "__main__":
    main()