from typing import Any, Dict, List, Sequence, Tuple

from gym3.types import ValType


class Env:
    """
    An interface for reinforcement learning environments.

    :param ob_space: ValType representing the valid observations generated by the environment
    :param ac_space: ValType representing the valid actions that an agent can take in the environment
    :param num: number of simultaneous episodes
    """

    def __init__(self, ob_space: ValType, ac_space: ValType, num: int):
        self.ob_space = ob_space
        self.ac_space = ac_space
        self.num = num

    def 
(self) -> Tuple[Any, Any, Any]:
        """
        Structured data that's accessible to the agent

        Can be called zero or more times per act() call; it is idempotent and does
        not affect the environment.  Returns the observation after the last act()
        call has affected the environment or the initial state if no act()
        call has been made.

        The following snippet shows what the start of a rollout looks like, and
        the usual convention for indexing the values:

        env = Env()
        # initial reward is ignored; didn't depend on actions
        _reward, ob_0, first_0 = env.observe()
        env.act(ac_0)
        reward_0, ob_1, first_1 = env.observe()
        env.act(ac_1)

        Note that the initial reward `_reward` is ignored by algorithms, since it
        didn't depend on the actions. And that first_0 is always True since the
        environment was just created.

        :returns:
            reward: (float array) last reward, shaped (num,)
            ob: observation matching ob_space, with (num,) prepended to the shape of each leaf
            first: (bool array) whether new episode just started, shaped (num,)
        """
        raise NotImplementedError

    def get_info(self) -> List[Dict]:
        """
        Return unstructured diagnostics that aren't accessible to the agent
        Per-episode stats, rendered images, etc.
        Corresponds to same timestep as the observation from observe().

        :returns: a list of dictionaries with length `self.num`
        """
        return [{} for _ in range(self.num)]

    def act(self, ac: Any) -> None:
        """
        Apply an action

        :param ac: action matching ac_space, with (num,) prepended to the shape of each leaf
        """
        raise NotImplementedError

    def callmethod(
        self, method: str, *args: Sequence[Any], **kwargs: Sequence[Any]
    ) -> List[Any]:
        """
        Call a method on the underlying python object that offers the Gym3 interface.
        By default, this is the Env instance (self), but this can be overridden if you are
        encapsulating another object or if your wrapper wishes to handle the method call.

        :param method: name of method to call on the base object
        :param *args, **kwargs: the value of each argument must be a list with length equal to `self.num`

        :returns: A list of objects, with length equal to `self.num`
        """
        return getattr(self, method)(*args, **kwargs)

observe