"""Benchmark runner — evaluate baselines on all PeptideGym environments."""
from __future__ import annotations

import json
import os
from typing import Any


def run_benchmarks(
    configs: list[dict[str, str]] | None = None,
    n_episodes: int = 100,
) -> list[dict[str, Any]]:
    """Run random + heuristic baselines on all configs.

    Returns a list of result dicts, one per environment.
    """
    from peptidegym.benchmarks.environments import BENCHMARK_CONFIGS
    from peptidegym.training.evaluate import compare_baselines

    if configs is None:
        configs = BENCHMARK_CONFIGS

    all_results: list[dict[str, Any]] = []
    for cfg in configs:
        env_id = cfg["env_id"]
        name = cfg.get("name", env_id)
        print(f"Benchmarking {name} ({env_id}) ...")
        try:
            result = compare_baselines(env_id, n_episodes=n_episodes)
            result["name"] = name
            all_results.append(result)
        except Exception as e:
            print(f"  SKIP {name}: {e}")
            all_results.append({"env_id": env_id, "name": name, "error": str(e)})

    return all_results


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="Run PeptideGym benchmarks")
    parser.add_argument("--episodes", type=int, default=100)
    parser.add_argument("--output", default="results/benchmark_results.json")
    args = parser.parse_args()

    results = run_benchmarks(n_episodes=args.episodes)

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)

    print(f"\nResults saved to {args.output}")
    for r in results:
        if "error" not in r:
            rand_mr = r.get("random", {}).get("mean_reward", "N/A")
            heur_mr = r.get("heuristic", {}).get("mean_reward", "N/A")
            print(f"  {r['name']}: random={rand_mr:.4f}, heuristic={heur_mr:.4f}")
        else:
            print(f"  {r['name']}: ERROR - {r['error']}")


if __name__ == "__main__":
    main()
