from typing import Tuple, List, Dict, Optional
import copy

import gymnasium as gym
from gymnasium import spaces
from gymnasium.utils import seeding
import numpy as np
from gymnasium.spaces import MultiDiscrete

try:
    from pymarlzooplus.envs.multiagent_particle_envs.mpe.rendering import Viewer
except Exception as e:
    print(
        "\n#####################################################################################"
        f"\nCannot import 'Viewer' from 'multiagent_particle_envs.mpe.rendering' due to the following error: \n\n{e}"
        "\n\nProbably this is a problem of 'pyglet', so rendering will not work."
        "\n#####################################################################################\n"
    )


# environment for all agents in the multiagent world
# currently code assumes that no agents will be created/destroyed at runtime!
class MultiAgentEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array']
    }

    def __init__(
            self,
            world,
            reset_callback=None,
            reward_callback=None,
            observation_callback=None,
            info_callback=None,
            done_callback=None,
            shared_viewer=True
    ):

        world = copy.deepcopy(world)
        self.world = world
        self.agents = self.world.policy_agents
        # set required vectorized gym env property
        self.n = len(world.policy_agents)
        # scenario callbacks
        self.reset_callback = reset_callback
        self.reward_callback = reward_callback
        self.observation_callback = observation_callback
        self.info_callback = info_callback
        self.done_callback = done_callback
        # environment parameters
        self.discrete_action_space = True
        # if true, action is a number 0...N, otherwise action is a one-hot N-dimensional vector
        self.discrete_action_input = False
        # if true, even the action is continuous, action will be performed discretely
        self.force_discrete_action = world.discrete_action if hasattr(world, 'discrete_action') else False
        # if true, every agent has the same reward
        self.shared_reward = world.collaborative if hasattr(world, 'collaborative') else False
        self.time = 0

        # configure spaces
        self.action_space = [] # type: ignore[override]
        self.observation_space = [] # type: ignore[override]
        for agent in self.agents:
            total_action_space = []
            # physical action space
            if self.discrete_action_space:
                u_action_space = spaces.Discrete(world.dim_p * 2 + 1)
            else:
                u_action_space = spaces.Box(low=-agent.u_range, high=+agent.u_range, shape=(world.dim_p,), dtype=np.float32)
            if agent.movable:
                total_action_space.append(u_action_space)
            # communication action space
            if self.discrete_action_space:
                c_action_space = spaces.Discrete(world.dim_c)
            else:
                c_action_space = spaces.Box(low=0.0, high=1.0, shape=(world.dim_c,), dtype=np.float32)
            if not agent.silent:
                total_action_space.append(c_action_space)
            # total action space
            if len(total_action_space) > 1:
                # all action spaces are discrete, so simplify to MultiDiscrete action space
                if all([isinstance(act_space, spaces.Discrete) for act_space in total_action_space]):
                    act_counts = [space.n for space in total_action_space]
                    act_space = MultiDiscrete(act_counts)
                    act_space.n = int(np.prod(act_space.nvec))
                else:
                    act_space = spaces.Tuple(total_action_space)
                self.action_space.append(act_space)  # type: ignore[override]
            else:
                self.action_space.append(total_action_space[0])  # type: ignore[override]
            # observation space
            obs_dim = len(observation_callback(agent, self.world))
            self.observation_space.append(  # type: ignore[override]
                spaces.Box(low=-np.inf, high=+np.inf, shape=(obs_dim,), dtype=np.float32)
            )
            agent.action.c = np.zeros(self.world.dim_c)

        self.action_space = spaces.Tuple(tuple(self.action_space))  # type: ignore[override]
        self.observation_space = spaces.Tuple(tuple(self.observation_space))  # type: ignore[override]
        self.n_agents = self.n

        # rendering
        self.render_geoms = None
        self.render_geoms_xform = None
        self.shared_viewer = shared_viewer
        if self.shared_viewer:
            self.viewers: Optional[List[Viewer], List[None]] = [None]
        else:
            self.viewers: Optional[List[Viewer], List[None]] = [None] * self.n
        self._reset_render()

    def seed(self, seed):
        if seed is not None:
            self._np_random, seed = seeding.np_random(seed)
            self.world.seed(seed)

    def step(
            self,
            action_n: list
    ) -> Tuple[
        List[np.ndarray],
        List[float],
        List[bool],
        List[bool],
        Dict
    ]:

        if hasattr(self.action_space, 'spaces'):
            action_space = self.action_space.spaces
        else:
            raise AttributeError("The observation space must be a Tuple of spaces")

        one_hot_actions = []
        for act, acsp in zip(action_n, action_space):
            one_hot = np.zeros(acsp.n)
            one_hot[act] = 1.0
            one_hot_actions.append(one_hot)
        action_n = one_hot_actions

        obs_n = []
        reward_n = []
        done_n = []
        info_n = {'n': []}
        self.agents = self.world.policy_agents
        # set action for each agent
        for i, agent in enumerate(self.agents):
            self._set_action(action_n[i], agent, self.action_space[i])
        # advance world state
        self.world.step()
        # record observation for each agent
        for agent in self.agents:
            obs_n.append(self._get_obs(agent))
            reward_n.append(self._get_reward(agent))
            done_n.append(self._get_done(agent))

            info_n['n'].append(self._get_info(agent))

        # all agents get total reward in cooperative case
        reward = np.sum(reward_n)
        if self.shared_reward:
            reward_n = [reward] * self.n

        # Just for support with gymnasium
        truncated = [False] * len(done_n)

        return obs_n, reward_n, done_n, truncated, info_n

    def reset(self, seed=None, options=None):

        # Options are used just for compatibility with gymnasium
        assert options is None, "Options must be None"

        if seed is not None:
            # setting seed
            super().reset(seed=seed, options=options)

        # reset world
        self.reset_callback(self.world)
        # reset renderer
        self._reset_render()
        # record observations for each agent
        obs_n = []
        self.agents = self.world.policy_agents
        for agent in self.agents:
            obs_n.append(self._get_obs(agent))

        # Dictionary returned just for compatibility with gymnasium
        return tuple(obs_n), {}

    # get info used for benchmarking
    def _get_info(self, agent):
        if self.info_callback is None:
            return {}
        return self.info_callback(agent, self.world)

    # get observation for a particular agent
    def _get_obs(self, agent):
        if self.observation_callback is None:
            return np.zeros(0)
        return self.observation_callback(agent, self.world).astype(np.float32)

    # get dones for a particular agent
    # unused right now -- agents are allowed to go beyond the viewing screen
    def _get_done(self, agent):
        if self.done_callback is None:
            return False
        return self.done_callback(agent, self.world)

    # get reward for a particular agent
    def _get_reward(self, agent):
        if self.reward_callback is None:
            return 0.0
        return self.reward_callback(agent, self.world)

    # set env action for a particular agent
    def _set_action(self, action, agent, action_space, time=None):
        agent.action.u = np.zeros(self.world.dim_p)
        agent.action.c = np.zeros(self.world.dim_c)
        # process action
        if isinstance(action_space, MultiDiscrete):
            act = []
            size = action_space.nvec
            index = 0
            for s in size:
                act.append(action[index:(index+s)])
                index += s
            action = act
        else:
            action = [action]

        if agent.movable:
            # physical action
            if self.discrete_action_input:
                agent.action.u = np.zeros(self.world.dim_p)
                # process discrete action
                if action[0] == 1: agent.action.u[0] = -1.0
                if action[0] == 2: agent.action.u[0] = +1.0
                if action[0] == 3: agent.action.u[1] = -1.0
                if action[0] == 4: agent.action.u[1] = +1.0
            else:
                if self.force_discrete_action:
                    d = np.argmax(action[0])
                    action[0][:] = 0.0
                    action[0][d] = 1.0
                if self.discrete_action_space:
                    agent.action.u[0] += action[0][1] - action[0][2]
                    agent.action.u[1] += action[0][3] - action[0][4]
                else:
                    agent.action.u = action[0]
            sensitivity = 5.0
            if agent.accel is not None:
                sensitivity = agent.accel
            agent.action.u *= sensitivity
            action = action[1:]
        if not agent.silent:
            # communication action
            if self.discrete_action_input:
                agent.action.c = np.zeros(self.world.dim_c)
                agent.action.c[action[0]] = 1.0
            else:
                agent.action.c = action[0]
            action = action[1:]
        # make sure we used all elements of action
        assert len(action) == 0

    # reset rendering assets
    def _reset_render(self):
        self.render_geoms = None
        self.render_geoms_xform = None

    # render environment
    def render(self, mode='human'):
        if mode == 'human':
            alphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
            message = ''
            for agent in self.world.agents:
                comm = []
                for other in self.world.agents:
                    if other is agent: continue
                    if np.all(other.state.c == 0):
                        word = '_'
                    else:
                        word = alphabet[np.argmax(other.state.c)]
                    message += (other.name + ' to ' + agent.name + ': ' + word + '   ')
            print(message)

        for i in range(len(self.viewers)):
            # create viewers (if necessary)
            if self.viewers[i] is None:
                from pymarlzooplus.envs.multiagent_particle_envs.mpe import rendering
                self.viewers[i] = rendering.Viewer(700, 700)

        # create rendering geometry
        if self.render_geoms is None:
            # import rendering only if we need it (and don't import for headless machines)
            from pymarlzooplus.envs.multiagent_particle_envs.mpe import rendering
            self.render_geoms = []
            self.render_geoms_xform = []
            for entity in self.world.entities:
                geom = rendering.make_circle(entity.size)
                xform = rendering.Transform()
                if 'agent' in entity.name:
                    geom.set_color(*entity.color, alpha=0.5)
                else:
                    geom.set_color(*entity.color)
                geom.add_attr(xform)
                self.render_geoms.append(geom)
                self.render_geoms_xform.append(xform)

            # add geoms to viewer
            for viewer in self.viewers:
                viewer.geoms = []
                for geom in self.render_geoms:
                    viewer.add_geom(geom)

        results = []
        for i in range(len(self.viewers)):
            from pymarlzooplus.envs.multiagent_particle_envs.mpe import rendering
            # update bounds to center around agent
            cam_range = 1
            if self.shared_viewer:
                pos = np.zeros(self.world.dim_p)
            else:
                pos = self.agents[i].state.p_pos
            self.viewers[i].set_bounds(pos[0]-cam_range,pos[0]+cam_range,pos[1]-cam_range,pos[1]+cam_range)
            # update geometry positions
            for e, entity in enumerate(self.world.entities):
                self.render_geoms_xform[e].set_translation(*entity.state.p_pos)
            # render to display or array
            results.append(self.viewers[i].render(return_rgb_array = mode=='rgb_array'))

        if self.shared_viewer:
            assert len(results) == 1
            return results[0]

        return results

    def close(self):
        for viewer in self.viewers:
            if viewer:
                viewer.close()
