#!/usr/bin/env python3
"""
Benchmark SYNAPSE-X agents across all task difficulties.
"""

import json
import random
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from agents.baseline import run_episode as run_baseline_episode
from agents.random_agent import run_episode as run_random_episode
from env.environment import SynapseXEnvironment
from env.grader import TASK_REGISTRY, TASK_SEEDS, grade, grade_with_variance


def evaluate_agent(agent_name: str, runner):
    results = {}
    variance = {}
    for task_name in TASK_REGISTRY:
        task_env = SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=TASK_SEEDS[task_name])
        actions = runner(task_env)
        results[task_name] = grade(task_name, actions).model_dump()
        variance[task_name] = grade_with_variance(
            task_name,
            lambda _task_name, _obs, env: runner(env),
        ).model_dump()
    return {"agent": agent_name, "results": results, "variance": variance}


def main():
    baseline = evaluate_agent("baseline", run_baseline_episode)
    random_agent = evaluate_agent(
        "random",
        lambda env: run_random_episode(env, rng=random.Random(42)),
    )
    summary = {"baseline": baseline["results"], "random": random_agent["results"]}
    variance_summary = {"baseline": baseline["variance"], "random": random_agent["variance"]}

    task_names = list(TASK_REGISTRY.keys())
    print("Agent      " + "  ".join(f"{task_name:>7}" for task_name in task_names))
    print("-" * (11 + 9 * len(task_names)))
    print(
        "Random     "
        + "  ".join(f"{summary['random'][task_name]['score']:.4f}" for task_name in task_names)
    )
    print(
        "Baseline   "
        + "  ".join(f"{summary['baseline'][task_name]['score']:.4f}" for task_name in task_names)
    )
    print()
    print("Variance mean scores")
    print(
        "Random     "
        + "  ".join(f"{variance_summary['random'][task_name]['score']:.4f}" for task_name in task_names)
    )
    print(
        "Baseline   "
        + "  ".join(f"{variance_summary['baseline'][task_name]['score']:.4f}" for task_name in task_names)
    )
    print()
    print(json.dumps({"canonical": summary, "variance_mean": variance_summary}, indent=2))
    with open(PROJECT_ROOT / "benchmark_results.json", "w", encoding="utf-8") as handle:
        json.dump({"canonical": summary, "variance_mean": variance_summary}, handle, indent=2)


if __name__ == "__main__":
    main()