import abc
import typing as t

import numpy as np
from gymnasium.vector import VectorEnv

from rlbcore.api.gym_ import GymVectorEnv
from rlbcore.api.memory import ExperienceT, Memory


class Policy(t.Protocol):
    """A Policy is a function that maps observations to actions.

    Used to generate actions to take for a given state in an environment.

    See also:
        [ContinuousRandomPolicy](../policies.md#continuousrandompolicy) for an example
            implementation of this interface.

    """

    def __call__(
        self,
        observations: np.ndarray[t.Any, np.dtype[t.Any]],
        *,
        t: int | None = None,
        eval: bool = False,
        rng: np.random.Generator | None = None,
        **kwargs: t.Any,
    ) -> np.ndarray[t.Any, np.dtype[t.Any]]:
        raise NotImplementedError


class EvalMetrics(t.NamedTuple):
    """The metrics returned by an Evaluator.

    Attributes:
        returns: The returns of the episodes generated by the evaluator.
        n_transitions: The number of transitions that occurred in the episodes generated
            by the evaluator.
        time_taken: The time taken to generate the episodes.
    """

    returns: list[float]
    n_transitions: int
    time_taken: float


class Evaluator(abc.ABC):
    """The Evaluator evaluates the performance of an agent.

    Attributes:
        env: The environment to evaluate the agent in.

    See also:
        [SimpleEvaluator](../evaluators.md#simpleevaluator) for an example
            implementation of this interface.
    """

    env: VectorEnv

    @abc.abstractmethod
    def step(
        self, policy: Policy, *, rng: np.random.Generator | None = None
    ) -> EvalMetrics:
        """Evaluate the performance of the policy on the given env.

        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: The returns of the episodes generated by the evaluator.
            - n_transitions: The number of transitions that occurred in the episodes
                generated by the evaluator.
            - time_taken: The time taken to generate the episodes.
        """
        raise NotImplementedError


class Observer(abc.ABC, t.Generic[ExperienceT]):
    """An Observer takes a step in the environment and stores the experience in the
    memory.

    See also:
        [TransitionObserver](../observers.md#transitionobserver) for an example
            implementation of this interface.
    """

    env: GymVectorEnv
    memory: Memory[ExperienceT]
    policy: Policy

    def step(
        self,
        *,
        rng: np.random.Generator | None = None,
    ) -> None:
        """Step the environment and store the experience in the memory.

        Args:
            device: The device on which the policy model exists.
            rng: The random number generator to use for sampling random numbers from
                numpy.
            generator: The generator used for sampling random numbers.
        """
        self.step_async(rng=rng)
        self.step_wait()

    @abc.abstractmethod
    def step_async(
        self,
        *,
        rng: np.random.Generator | None = None,
    ) -> None:
        """Step the environment and store the experience in the memory asynchronously.

        Args:
            device: The device on which the policy model exists.
            rng: The random number generator to use for sampling random numbers from
                numpy.
            generator: The generator used for sampling random numbers.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def step_wait(self) -> None:
        """Wait for the asynchronous step to complete."""
        raise NotImplementedError
