from gym import spaces
import numpy as np
from pettingzoo.utils.env import ParallelEnv
from pettingzoo.utils.agent_selector import agent_selector
from gym.utils import seeding
from gym import spaces
from pettingzoo.utils import wrappers


def make_env(raw_env):
    def env(**kwargs):
        env = raw_env(**kwargs)
        env = wrappers.AssertOutOfBoundsWrapper(env)
        backup_policy = "taking zero action (no movement, communication 0)"
        env = wrappers.NanNoOpWrapper(env, 4, backup_policy)
        env = wrappers.OrderEnforcingWrapper(env)
        return env
    return env


class MatrixGameEnv(ParallelEnv):
    def __init__(self, num_agents, num_actions, utility_matrix, memory_length, max_frames):
        super(MatrixGameEnv, self).__init__()

        self.num_agents = num_agents
        self.num_actions = [num_actions] * num_agents if type(num_actions) == int else list(num_actions)
        self.sum_actions = sum(self.num_actions)
        self.utility_matrix = np.array(utility_matrix)
        self.memory_length = memory_length
        self.max_frames = max_frames

        assert len(utility_matrix.shape) == num_agents + 1
        assert utility_matrix.shape[num_agents] == num_agents
        for i in range(num_agents):
            assert utility_matrix.shape[i] == self.num_actions[i]

        self.agents = list(range(self.num_agents))
        self.action_spaces = list()
        self.observation_spaces = list()
        self.input_structures = list()
        self.action_positions = list()
        position = 0
        for i in range(self.num_agents):
            self.action_spaces.append(spaces.Discrete(self.num_actions[i]))
            observation_space = spaces.Box(low=0., high=1., shape=(1 + self.memory_length * self.sum_actions,),
                                           dtype=np.float32)
            self.observation_spaces.append(observation_space)
            self.input_structures.append(self._get_input_structure_agent(i))
            self.action_positions.append(position)
            position += self.num_actions[i]

        self.action_history = np.zeros((max_frames, self.sum_actions))
        self.steps = 0

    def _get_obs_agent(self, agent_id):
        num_padding = max(0, self.memory_length - self.steps)
        padding = np.zeros((num_padding, self.sum_actions))
        history = self.action_history[max(0, self.steps - self.memory_length): self.steps]
        history = np.concatenate((padding, history), axis=0)
        l_p = self.action_positions[agent_id]
        r_p = l_p + self.num_actions[agent_id]
        history = np.concatenate((history[:, l_p: r_p], history[:, :l_p], history[:, r_p:]), axis=1).reshape(-1)
        return np.concatenate(([self.steps / self.max_frames], history))

    def _get_obs_all(self):
        obs_all = [self._get_obs_agent(i) for i in range(self.num_agents)]
        # print(obs_all)
        return obs_all

    def _get_input_structure_agent(self, agent_id):
        structure = [("step", 1)]
        for i in range(self.memory_length):
            structure.append((f"history-{i}-self", self.num_actions[agent_id]))
            for j in range(self.num_agents):
                if j != agent_id:
                    structure.append((f"history-{i}-agent-{j}", self.num_actions[j]))
        return structure

    def reset(self):
        self.action_history.fill(0.)
        self.steps = 0
        return self._get_obs_all()

    def step(self, actions):
        list_actions = [actions[i] for i in range(self.num_agents)]
        utility = self.utility_matrix[tuple(list_actions)]

        si = 0
        for i in range(self.num_agents):
            self.action_history[self.steps, si + actions[i]] = 1.
            si += self.num_actions[i]

        self.steps += 1

        rewards = (utility * 0.1).tolist()
        dones = [self.steps >= self.max_frames] * self.num_agents
        infos = [None] * self.num_agents

        return self._get_obs_all(), rewards, dones, infos

    def render(self, mode="human"):
        pass

    def seed(self, seed=None):
        pass

from gym import spaces
import numpy as np
from pettingzoo.utils.env import ParallelEnv
from pettingzoo.utils.agent_selector import agent_selector
from gym.utils import seeding
from gym import spaces
from pettingzoo.utils import wrappers


def make_env(raw_env):
    def env(**kwargs):
        env = raw_env(**kwargs)
        env = wrappers.AssertOutOfBoundsWrapper(env)
        backup_policy = "taking zero action (no movement, communication 0)"
        env = wrappers.NanNoOpWrapper(env, 4, backup_policy)
        env = wrappers.OrderEnforcingWrapper(env)
        return env
    return env


class MatrixGameEnv(ParallelEnv):
    def __init__(self, num_agents, num_actions, utility_matrix, memory_length, max_frames):
        super(MatrixGameEnv, self).__init__()

        self.num_agents = num_agents
        self.num_actions = [num_actions] * num_agents if type(num_actions) == int else list(num_actions)
        self.sum_actions = sum(self.num_actions)
        self.utility_matrix = np.array(utility_matrix)
        self.memory_length = memory_length
        self.max_frames = max_frames

        assert len(utility_matrix.shape) == num_agents + 1
        assert utility_matrix.shape[num_agents] == num_agents
        for i in range(num_agents):
            assert utility_matrix.shape[i] == self.num_actions[i]

        self.agents = list(range(self.num_agents))
        self.action_spaces = list()
        self.observation_spaces = list()
        self.input_structures = list()
        self.action_positions = list()
        position = 0
        for i in range(self.num_agents):
            self.action_spaces.append(spaces.Discrete(self.num_actions[i]))
            observation_space = spaces.Box(low=0., high=1., shape=(1 + self.memory_length * self.sum_actions,),
                                           dtype=np.float32)
            self.observation_spaces.append(observation_space)
            self.input_structures.append(self._get_input_structure_agent(i))
            self.action_positions.append(position)
            position += self.num_actions[i]

        self.action_history = np.zeros((max_frames, self.sum_actions))
        self.steps = 0

    def _get_obs_agent(self, agent_id):
        num_padding = max(0, self.memory_length - self.steps)
        padding = np.zeros((num_padding, self.sum_actions))
        history = self.action_history[max(0, self.steps - self.memory_length): self.steps]
        history = np.concatenate((padding, history), axis=0)
        l_p = self.action_positions[agent_id]
        r_p = l_p + self.num_actions[agent_id]
        history = np.concatenate((history[:, l_p: r_p], history[:, :l_p], history[:, r_p:]), axis=1).reshape(-1)
        return np.concatenate(([self.steps / self.max_frames], history))

    def _get_obs_all(self):
        obs_all = [self._get_obs_agent(i) for i in range(self.num_agents)]
        # print(obs_all)
        return obs_all

    def _get_input_structure_agent(self, agent_id):
        structure = [("step", 1)]
        for i in range(self.memory_length):
            structure.append((f"history-{i}-self", self.num_actions[agent_id]))
            for j in range(self.num_agents):
                if j != agent_id:
                    structure.append((f"history-{i}-agent-{j}", self.num_actions[j]))
        return structure

    def reset(self):
        self.action_history.fill(0.)
        self.steps = 0
        return self._get_obs_all()

    def step(self, actions):
        list_actions = [actions[i] for i in range(self.num_agents)]
        utility = self.utility_matrix[tuple(list_actions)]

        si = 0
        for i in range(self.num_agents):
            self.action_history[self.steps, si + actions[i]] = 1.
            si += self.num_actions[i]

        self.steps += 1

        rewards = (utility * 0.1).tolist()
        dones = [self.steps >= self.max_frames] * self.num_agents
        infos = [None] * self.num_agents

        return self._get_obs_all(), rewards, dones, infos

    def render(self, mode="human"):
        pass

    def seed(self, seed=None):
        pass

# print(obs_all)