"""Contains different ways to evaluate an agent."""
import typing as t

import attrs
import numpy as np
from tqdm.rich import tqdm

from rlbcore import api
from rlbcore import external_utils as extu


@attrs.define()
class SimpleEvaluator(api.Evaluator):
    """Specifies a simple procedure to evaluate an agent's with a given environment.

    Args:
        env: The environment to evaluate the agent in.
        n_episodes: The number of episodes to evaluate the agent for.
        n_transitions: The number of transitions to evaluate the agent for.
        min_episodes: The minimum number of episodes to evaluate the agent for.

    Raises:
        ValueError: If neither n_episodes or n_transitions is specified.
    """

    env: api.GymVectorEnv
    n_episodes: int | None = None
    n_transitions: int | None = None
    min_episodes: int = 1

    _ep_return_recorder: extu.EpisodeReturnRecorder = attrs.field(init=False)
    _default_rng: np.random.Generator = attrs.field(
        init=False, factory=np.random.default_rng
    )

    def __attrs_post_init__(self) -> None:
        self._ep_return_recorder = extu.EpisodeReturnRecorder()
        if all((x is None for x in (self.n_episodes, self.n_transitions))):
            raise ValueError("Must specify either n_episodes or n_transitions.")

    def step(
        self, policy: api.Policy, *, rng: np.random.Generator | None = None
    ) -> api.EvalMetrics:
        """Evaluate performance of agent against env for requested number of
        episodes / transitions.

        Args:
            policy: The policy to evaluate.
            rng: The random number generator used to seed env and any other source of
                randomness generated from numpy.

        Returns:
            returns (list[float]): The returns of the episodes generated by the
                evaluator.
            n_transitions (int): The number of transitions that occurred in the episodes
                generated by the evaluator.
            time_taken (float): The time taken for this step.

        Important:
            The evaluator will always return at least min_episodes episodes. This means
            if you set n_transitions, and n_transitions is less than the number of
            transitions to finish an episode, the evaluator will return at least
            min_episodes episodes (which defaults to 1)

        Example: Evaluate a random policy on Pendulum-v1 for 100 transitions.
            ```python
            >>> import numpy as np
            ... import pytest
            ... from rlbcore.evaluators import SimpleEvaluator
            ... from rlbcore.policies import ContinuousRandomPolicy
            ... from rlbcore.api import vector_make
            ... env = vector_make("Pendulum-v1", 1)
            ... evaluator = SimpleEvaluator(env, n_transitions=100)
            ... policy = ContinuousRandomPolicy.from_env(env)
            ... metrics = evaluator.step(policy, rng=np.random.default_rng(0))
            ... new_rng = np.random.default_rng(0)
            ... obs, _ = env.reset(seed=new_rng.integers(0, 2**32 - 1).item())
            ... ep_returns = []
            ... curr_ep_return = 0
            ... done = False
            ... # Note that if you don't change the min_episodes to 0, the evaluator
            ... # will always return at least 1 episode.
            ... while not done:
            ...     obs, rews, truncs, terms, _ = env.step(policy(obs, rng=new_rng))
            ...     curr_ep_return += rews.item()
            ...     done = (truncs | terms).item()
            ...     if done:
            ...         ep_returns.append(curr_ep_return)
            ...         curr_ep_return = 0
            ... assert metrics.returns == pytest.approx(ep_returns), (
            ...     metrics.returns, ep_returns
            ... )

            ```
        """
        with extu.Timer() as timer:
            ep_returns, n_transitions = self._evaluate(policy, rng or self._default_rng)
        return api.EvalMetrics(
            returns=ep_returns,
            n_transitions=n_transitions,
            time_taken=timer.seconds,
        )

    def _evaluate(
        self, policy: api.Policy, rng: np.random.Generator
    ) -> tuple[list[float], int]:
        n_transitions = 0
        total = self.n_episodes or self.n_transitions
        assert total is not None
        pbar = tqdm(total=total, leave=False, colour="green")
        obs, _ = self.env.reset(seed=int(rng.integers(0, 2**32 - 1)))
        ep_returns: list[float] = []
        while not self._done(len(ep_returns), n_transitions):
            obs, rewards, dones = self._transition(
                observations=obs, policy=policy, pbar=pbar, rng=rng
            )
            new_ep_returns = self._ep_return_recorder.track(
                rewards=rewards, dones=dones
            )
            if new_ep_returns is not None:
                assert dones.any()
                ep_returns.extend(new_ep_returns)
            # Number of transitions is number of envs which is the same as batch size
            # of dones
            n_transitions += len(dones)
        while len(ep_returns) < self.min_episodes:
            # Ensure that we have at least self.min_episodes number of episodes.
            obs, rewards, dones = self._transition(
                observations=obs, policy=policy, pbar=pbar, rng=rng
            )
            new_ep_returns = self._ep_return_recorder.track(
                rewards=rewards, dones=dones
            )
            if new_ep_returns is not None:
                assert dones.any()
                ep_returns.extend(new_ep_returns)
            n_transitions += len(dones)
        pbar.close()

        return ep_returns, n_transitions

    def _done(
        self,
        n_finished_episodes: int,
        n_transitions: int,
    ) -> bool:
        if self.n_episodes:
            return n_finished_episodes >= self.n_episodes
        if self.n_transitions:
            return n_transitions >= self.n_transitions
        raise ValueError("n_episodes or n_transitions must be specified.")

    def _transition(
        self,
        observations: np.ndarray[t.Any, np.dtype[t.Any]],
        policy: api.Policy,
        pbar: tqdm,  # type: ignore
        rng: np.random.Generator,
    ) -> tuple[
        np.ndarray[t.Any, np.dtype[t.Any]],
        np.ndarray[t.Any, np.dtype[np.float32]],
        np.ndarray[t.Any, np.dtype[np.bool_]],
    ]:
        actions = policy(observations=observations, eval=True, rng=rng)
        next_obs, rewards, terminals, truncateds, _ = self.env.step(actions)
        dones = terminals | truncateds
        pbar.update()
        return next_obs, rewards, dones
