"""Evaluation utilities for PeptideGym agents."""
from __future__ import annotations

from typing import Any

import numpy as np


def evaluate_agent(
    agent: Any,
    env: Any,
    n_episodes: int = 100,
) -> dict[str, Any]:
    """Run *agent* for *n_episodes* on *env* and return summary statistics.

    Returns a dict with:
        mean_reward, std_reward, mean_length, sequences
    """
    rewards: list[float] = []
    lengths: list[int] = []
    sequences: list[str] = []

    for _ in range(n_episodes):
        obs, info = env.reset()
        if hasattr(agent, "reset"):
            agent.reset()

        total_reward = 0.0
        done = False
        steps = 0

        while not done:
            action, _ = agent.predict(obs, deterministic=False)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += float(reward)
            done = terminated or truncated
            steps += 1

        rewards.append(total_reward)
        lengths.append(steps)
        sequences.append(info.get("sequence", ""))

    return {
        "mean_reward": float(np.mean(rewards)),
        "std_reward": float(np.std(rewards)),
        "mean_length": float(np.mean(lengths)),
        "sequences": sequences,
    }


def compare_baselines(
    env_id: str,
    n_episodes: int = 100,
) -> dict[str, Any]:
    """Compare random and heuristic baselines on the given environment.

    Optionally includes PPO if stable-baselines3 is installed.
    """
    import gymnasium as gym
    import peptidegym  # noqa: F401

    from peptidegym.agents.random_agent import RandomAgent

    results: dict[str, Any] = {"env_id": env_id}

    env = gym.make(env_id)

    # Random baseline
    random_agent = RandomAgent(env, seed=42)
    results["random"] = evaluate_agent(random_agent, env, n_episodes)

    # Heuristic baseline — pick the right agent for the env family
    heuristic_result = None
    if "AMP" in env_id:
        from peptidegym.agents.heuristic_agent import HeuristicAMPAgent
        heuristic_agent = HeuristicAMPAgent(env, seed=42)
        heuristic_result = evaluate_agent(heuristic_agent, env, n_episodes)
    elif "Cyclic" in env_id:
        from peptidegym.agents.heuristic_agent import HeuristicCyclicAgent
        heuristic_agent = HeuristicCyclicAgent(env, seed=42)
        heuristic_result = evaluate_agent(heuristic_agent, env, n_episodes)
    elif "Epitope" in env_id:
        from peptidegym.agents.heuristic_agent import HeuristicEpitopeAgent
        heuristic_agent = HeuristicEpitopeAgent(env, seed=42)
        heuristic_result = evaluate_agent(heuristic_agent, env, n_episodes)

    if heuristic_result is not None:
        results["heuristic"] = heuristic_result

    env.close()
    return results
