from modules.agents import REGISTRY as agent_REGISTRY
from components.action_selectors import REGISTRY as action_REGISTRY
import torch as th
import numpy as np
from torch import nn


# This multi-agent controller shares parameters between agents
class MAOSDQNMAC:
    def __init__(self, scheme, groups, args):
        self.n_agents = args.n_agents
        self.n_actions = args.n_actions
        self.args = args
        input_shape = self._get_input_shape(scheme)
        self._build_agents(input_shape)
        self.agent_output_type = args.agent_output_type

        self.action_selector = action_REGISTRY[args.action_selector](args)

        self.hidden_states = None

    def select_actions_by_scores(self, scores, avail_actions):
        # Only select actions for the selected batch elements in bs
        chosen_actions = self.action_selector.select_action(scores, avail_actions, t_env=-1)
        return chosen_actions
    
    def select_actions(self, ep_batch, t_ep, t_env, bs=slice(None), test_mode=False):
        # Only select actions for the selected batch elements in bs
        obs = ep_batch["obs"][:, t_ep]
        avail_actions = ep_batch["avail_actions"][:, t_ep]
        agent_outputs = self.forward(obs)
        chosen_actions = self.action_selector.select_action(agent_outputs[bs], avail_actions[bs], t_env=t_env, test_mode=test_mode)
        return chosen_actions

    def forward(self, obs, id=0, refresh_hidden=False):
        bs = obs.shape[0]
        if refresh_hidden:
            self.init_hidden(bs, id=id)
        agent_inputs = self._build_inputs(obs)
        agent_outs_v, agent_outs_a, self.hidden_states = self.agent[id](agent_inputs, self.hidden_states)
        agent_outs_v = agent_outs_v.view(bs, self.n_agents, 1)
        agent_outs_a = agent_outs_a.view(bs, self.n_agents, self.n_actions)
        agent_outs = agent_outs_v + agent_outs_a
        if self.args.v_type == 'mean':
            agent_outs -= th.mean(agent_outs_a, dim=2, keepdim=True)
        else:
            agent_outs -= th.log(th.sum(th.exp(self.args.alpha * agent_outs_a), dim=2, keepdim=True))/self.args.alpha
        if self.args.random_output and bs==1 and np.random.rand()<0.001:
            print("actor outputs: ",obs,agent_inputs,agent_outs_v,agent_outs)
        return agent_outs

    def init_hidden(self, batch_size, id=0):
        self.hidden_states = self.agent[id].init_hidden().unsqueeze(0).expand(batch_size, self.n_agents, -1)  # bav
        
    def parameters(self):
        return list(self.agent[0].parameters())+list(self.agent[1].parameters())

    def load_state(self, other_mac):
        self.agent[0].load_state_dict(other_mac.agent[0].state_dict())
        self.agent[1].load_state_dict(other_mac.agent[1].state_dict())

    def cuda(self):
        self.agent[0].cuda()
        self.agent[1].cuda()

    def save_models(self, path):
        th.save(self.agent[0].state_dict(), "{}/agent_0.th".format(path))
        th.save(self.agent[1].state_dict(), "{}/agent_1.th".format(path))

    def load_models(self, path):
        self.agent[0].load_state_dict(th.load("{}/agent_0.th".format(path), map_location=lambda storage, loc: storage))
        self.agent[1].load_state_dict(th.load("{}/agent_1.th".format(path), map_location=lambda storage, loc: storage))

    def _build_agents(self, input_shape):
        self.agent = nn.ModuleList([
            agent_REGISTRY[self.args.agent](input_shape, self.args),
            agent_REGISTRY[self.args.agent](input_shape, self.args)
        ])

    def _build_inputs(self, obs):
        # Assumes homogenous agents with flat observations.
        # Other MACs might want to e.g. delegate building inputs to each agent
        bs = obs.shape[0]
        inputs = [obs]
        if self.args.obs_agent_id:
            inputs.append(th.eye(self.n_agents, device=obs.device).unsqueeze(0).expand(bs, -1, -1))

        inputs = th.cat([x.reshape(bs*self.n_agents, -1) for x in inputs], dim=1)
        return inputs

    def _get_input_shape(self, scheme):
        input_shape = scheme["obs"]["vshape"]
        if self.args.obs_agent_id:
            input_shape += self.n_agents

        return input_shape
