"""Evaluation utilities for OncoSim agents."""

from __future__ import annotations

import gymnasium as gym
import numpy as np


def evaluate_agent(
    env: gym.Env,
    agent,
    n_episodes: int = 100,
    seed: int = 42,
) -> dict[str, float]:
    """Evaluate any agent with an act(obs) method.

    Args:
        env: Gymnasium environment.
        agent: Agent with act(obs) -> action method.
        n_episodes: Number of evaluation episodes.
        seed: Random seed.

    Returns:
        Dict with mean_reward, std_reward, min_reward, max_reward, success_rate.
    """
    rewards = []
    successes = 0

    for ep in range(n_episodes):
        obs, _ = env.reset(seed=seed + ep)
        if hasattr(agent, "reset"):
            agent.reset()
        total_reward = 0.0
        done = False
        final_info = {}

        while not done:
            action = agent.act(obs)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            final_info = info
            done = terminated or truncated

        rewards.append(total_reward)
        # Consider episode successful if plan quality > 0.7 or tcp > 0.9
        if final_info.get("plan_quality", 0) > 0.7 or final_info.get("tcp", 0) > 0.9:
            successes += 1

    return {
        "mean_reward": float(np.mean(rewards)),
        "std_reward": float(np.std(rewards)),
        "min_reward": float(np.min(rewards)),
        "max_reward": float(np.max(rewards)),
        "success_rate": successes / n_episodes,
    }


def compare_agents(
    env: gym.Env,
    agents: dict[str, object],
    n_episodes: int = 100,
    seed: int = 42,
) -> dict[str, dict[str, float]]:
    """Compare multiple agents on the same environment.

    Args:
        env: Gymnasium environment.
        agents: Dict mapping agent name to agent object.
        n_episodes: Number of evaluation episodes.
        seed: Random seed.

    Returns:
        Dict mapping agent name to evaluation results.
    """
    results = {}
    for name, agent in agents.items():
        results[name] = evaluate_agent(env, agent, n_episodes, seed)
    return results
