from abc import ABC, abstractmethod
from typing import Any, Tuple, Sequence

import gym.spaces as spaces


class MultiAgentSafetyEnv(ABC):
    """
    An environment for a multi-agent gridworld.
    IMPORTANT: Unlike a gym.env, this is stateless.
    The step function accepts the current state and returns the next state; the current state must be stored elsewhere
    """

    @abstractmethod
    def agent_obs_spaces(self) -> Sequence[spaces.Space]:
        """
        :return: The observation spaces of each agent
        """
        pass

    @abstractmethod
    def agent_actions_spaces(self) -> Sequence[spaces.Space]:
        """
        :return: The action space for each of the agents
        """
        pass

    @abstractmethod
    def state_space(self) -> spaces.Space:
        """
        :return: The true state space of the environment
        """
        pass

    @abstractmethod
    def initial_state(self) -> Tuple[Any, Sequence[Any]]:
        """
        :return: A (possibly randomized) initial environment state, and the joint observations from that state
        """
        pass

    @abstractmethod
    def step(self, environment_state, joint_action: Sequence[Any]) -> Tuple[
        Any, Sequence[Any], Sequence[float], bool, bool]:
        """
        Take a step in the environment
        NOTE This does not modify environment_state, but returns a new state
        :return: A tuple (new environment state, [observations for each agent], [rewards for each agent], done, safety)
        """
        pass
