from modules.agents.multi_task import REGISTRY as agent_REGISTRY
from modules.decomposers import REGISTRY as decomposer_REGISTRY

from components.action_selectors import REGISTRY as action_REGISTRY
import torch as th
import torch.distributions as D
import numpy as np
import torch.nn.functional as F
from copy import deepcopy

# This multi-agent controller shares parameters between agents
class BiKTMAC:
    def __init__(self, train_tasks, task2scheme, task2args, main_args):
        # set some task-specific attributes
        self.train_tasks = train_tasks
        self.task2scheme = task2scheme
        self.task2args = task2args
        self.task2n_agents = {task: self.task2args[task].n_agents for task in train_tasks}
        self.main_args = main_args

        # set some common attributes
        self.agent_output_type = main_args.agent_output_type
        self.action_selector = action_REGISTRY[main_args.action_selector](main_args)

        # get decomposer for each task
        self.task2decomposer, self.task2dynamic_decoder = {}, {}
        self.surrogate_decomposer = None
        for task in train_tasks:
            task_args = self.task2args[task]
            task_decomposer = decomposer_REGISTRY[task_args.env](task_args)
            self.task2decomposer[task] = task_decomposer
            if not self.surrogate_decomposer:
                self.surrogate_decomposer = task_decomposer
            # set obs_shape
            task_args.obs_shape = task_decomposer.obs_dim

        # build agents
        task2input_shape_info = self._get_input_shape()
        self._build_agents(task2input_shape_info)
        

        if self.main_args.pretrain_vqvae:
            self._build_vqvae_agents(task2input_shape_info)
        
        elif self.main_args.train_DT_w_glsk:
            self._build_dt_agents(task2input_shape_info)
        
        self.hidden_states_enc = None
        self.hidden_states_dec = None
        self.skill = None
        self.skill_dim = main_args.skill_dim
        self.c_step = main_args.c_step # doesn't use it

    def select_actions(self, ep_batch, t_ep, t_env, task, bs=slice(None), test_mode=False):
        if self.main_args.pretrain_vae:
            # Only select actions for the selected batch elements in bs
            avail_actions = ep_batch["avail_actions"][:, t_ep]
            agent_outputs = self.forward(ep_batch, t_ep, task, test_mode=test_mode)
            chosen_actions = self.action_selector.select_action(agent_outputs[bs], avail_actions[bs], t_env, test_mode=test_mode)
            
        elif self.main_args.train_DT:
            # Only select actions for the selected batch elements in bs
            avail_actions = ep_batch["avail_actions"][:, t_ep]
            agent_outputs = self.forward(ep_batch, t_ep, task, test_mode=test_mode)
            chosen_actions = self.action_selector.select_action(agent_outputs[bs], avail_actions[bs], t_env, test_mode=test_mode)
        
        elif self.main_args.train_DT_w_glsk:
            # Only select actions for the selected batch elements in bs
            avail_actions = ep_batch["avail_actions"][:, t_ep]
            agent_outputs = self.forward(ep_batch, t_ep, task, test_mode=test_mode)
            chosen_actions = self.action_selector.select_action(agent_outputs[bs], avail_actions[bs], t_env, test_mode=test_mode)
            
        return chosen_actions

    def forward_qvalue(self, ep_batch, t, task, test_mode=False):
        agent_inputs = self._build_inputs(ep_batch, t, task)
        agent_outs, self.hidden_states_enc = self.agent.forward_qvalue(agent_inputs, self.hidden_states_enc, task)

        return agent_outs.view(ep_batch.batch_size, self.task2n_agents[task], -1)

    def forward_skill(self, ep_batch, t, task, test_mode=False, actions=None):
        agent_inputs = ep_batch["state"][:, t]
        agent_outs, self.hidden_states_enc = self.agent.forward_skill(agent_inputs, self.hidden_states_enc, task, actions=actions)

        return agent_outs.view(ep_batch.batch_size, self.task2n_agents[task], -1)

    def forward_both(self, ep_batch, t, task, test_mode=False):
        agent_inputs = self._build_inputs(ep_batch, t, task)
        agent_outs, pri_outs, self.hidden_states_enc = self.agent.forward_both(agent_inputs, self.hidden_states_enc, task)

        return agent_outs.view(ep_batch.batch_size, self.task2n_agents[task], -1), \
               pri_outs.view(ep_batch.batch_size, self.task2n_agents[task], -1)

    def forward_obs_skill(self, ep_batch, t, task, test_mode=False):
        agent_inputs = self._build_inputs(ep_batch, t, task)
        agent_outs, self.hidden_states_enc = self.agent.forward_qvalue(agent_inputs, self.hidden_states_enc, task)

        return agent_outs.view(ep_batch.batch_size, self.task2n_agents[task], -1)

    def forward_seq_action(self, ep_batch, skill_action, t, task, test_model=False):
        agent_seq_inputs = []
        skill_action = skill_action.reshape(-1, self.skill_dim)
        for i in range(self.c_step):
            agent_inputs = self._build_inputs(ep_batch, t+i, task)
            agent_seq_inputs.append(agent_inputs)
        agent_seq_inputs = th.stack(agent_seq_inputs, dim=1)

        agent_seq_outs, self.hidden_states_dec = self.agent.forward_seq_action(agent_seq_inputs, self.hidden_states_dec, task, skill_action)

        return agent_seq_outs.view(ep_batch.batch_size, self.c_step, self.task2n_agents[task], -1)
    
    
    def forward(self, ep_batch, t, task, test_mode=False):
        if self.main_args.pretrain_vae:
            agent_inputs = self._build_inputs(ep_batch, t, task)
            avail_actions = ep_batch["avail_actions"][:, t]

            bs = agent_inputs.shape[0] // self.task2n_agents[task]

            if t % self.c_step == 0:
                agent_outs, self.hidden_states_enc, self.hidden_states_dec, self.skill = self.agent(agent_inputs, self.hidden_states_enc, self.hidden_states_dec, task, None)
            else:
                agent_outs, self.hidden_states_enc, self.hidden_states_dec, _= self.agent(agent_inputs, self.hidden_states_enc, self.hidden_states_dec, task, self.skill)

            if self.agent_output_type == "pi_logits":

                if getattr(self.main_args, "mask_before_softmax", True):
                    reshaped_avail_actions = avail_actions.reshape(ep_batch.batch_size * self.task2n_agents[task], -1)
                    agent_outs[reshaped_avail_actions == 0] = -1e10

                agent_outs = th.nn.functional.softmax(agent_outs, dim=-1)
                if not test_mode:
                    epsilon_action_num = agent_outs.size(-1)
                    if getattr(self.main_args, "mask_before_softmax", True):
                        epsilon_action_num = reshaped_avail_actions.sum(dim=1, keepdim=True).float()

                    agent_outs = ((1 - self.action_selector.epsilon) * agent_outs
                                + th.ones_like(agent_outs) * self.action_selector.epsilon / epsilon_action_num)

                    if getattr(self.main_args, "mask_before_softmax", True):
                        agent_outs[reshaped_avail_actions == 0] = 0.0

            return agent_outs.view(ep_batch.batch_size, self.task2n_agents[task], -1)
        
        
        elif self.main_args.train_DT_w_glsk:
            #* basic info
            step_end = t
            step_start = max(0, t+1 - self.main_args.step_block_size)
            if t <= self.main_args.step_block_size-1:
                padding_step_num = self.main_args.step_block_size - (t+1)
            else:
                padding_step_num  = 0
            agent_num = self.task2n_agents[task]
            
            #* obs_input 
            # [batch_size*(t+1-step_start)*agent_num, obs_dim] --> [batch_size, context_length, agent_num, obs_dim] --> [batch_size*context_length*agent_num, obs_dim]
            obs_input = self._build_inputs_4dt(ep_batch, step_start, t+1, task).reshape(1, (t+1-step_start),agent_num,-1)
            obs_input = th.cat([obs_input, th.zeros(1, padding_step_num, *obs_input.shape[2:]).to('cuda')], dim=1)
            obs_input = obs_input.reshape(-1, obs_input.shape[-1])
            
            
            #* tactic
            gl_sk_id = ep_batch['global_skill_id'][:, step_start:t+1,:,:].reshape(1, (t+1-step_start),agent_num,-1)
            gl_sk_id = th.cat([gl_sk_id, th.zeros(1, self.main_args.step_block_size - (t+1-step_start), *gl_sk_id.shape[2:]).to('cuda')], dim=1)
            assert self.main_args.step_block_size - (t+1-step_start) >=0  
            
            #* skill
            sk = ep_batch['skill'][:, step_start:t+1,:,:].reshape(1, (t+1-step_start),agent_num,-1)
            sk = th.cat([sk, th.zeros(1, self.main_args.step_block_size - (t+1-step_start), *sk.shape[2:]).to('cuda')], dim=1)
            assert self.main_args.step_block_size - (t+1-step_start) >=0  
   
            
            #* rtg
            if t == 0:
                ep_batch['rtg'][:,0] = 20.0
                ep_batch.update({"rtg": deepcopy([ep_batch['rtg'][0,0].cpu()])}, ts=0)
            else:
                ep_batch['rtg'][:,t] = ep_batch['rtg'][:,t-1] - ep_batch['reward'][:,t-1]
            # [batch_size, (t+1-step_start), agent_num, 1] --> [batch_size, context_length, agent_num, 1]
            rtg = ep_batch['rtg'][:,step_start:t+1].reshape(1, (t+1-step_start),agent_num,-1)
            rtg = th.cat([rtg, th.zeros(1, self.main_args.step_block_size - ((t+1-step_start)), *rtg.shape[2:]).to('cuda')], dim=1)
            
            #* output skills
            input = (obs_input, sk, gl_sk_id, rtg, step_start, t+1+padding_step_num)
            with th.no_grad():
                # [batch_size*timesteps*agent_num, 8, n_embd]
                gl_skill_id_probs_Ts, _ = self.dt_agent(input, task)
            gl_skill_id_probs_t = gl_skill_id_probs_Ts.reshape(1, self.main_args.step_block_size, agent_num, -1)[0,t-step_start]
            gl_skill_id_t = th.argmax(gl_skill_id_probs_t, dim=1).unsqueeze(-1)
            gl_sk_id[:,t-step_start] = gl_skill_id_t
            input = (obs_input, sk, gl_sk_id, rtg, step_start, t+1+padding_step_num)
            
            with th.no_grad():
                # [batch_size*timesteps*agent_num, 8, n_embd]
                gl_skill_id_probs_Ts, logits_skill = self.dt_agent(input, task)
                gl_skill_id_probs_t = gl_skill_id_probs_Ts.reshape(1, self.main_args.step_block_size, agent_num, -1)[0,t-step_start]
                gl_skill_id_t = th.argmax(gl_skill_id_probs_t, dim=1).unsqueeze(-1)
            

            #* update skill & rtg   
            ep_batch.update({ "skill": [logits_skill.reshape(1, self.main_args.step_block_size, agent_num, -1)[0,t-step_start].cpu()]}, ts=t)
            ep_batch.update({ "global_skill_id": [gl_skill_id_t.cpu()]}, ts=t) 
            ep_batch.update({ "rtg": [rtg[0,t-step_start].cpu()]}, ts=t)
            
            #* skill to actions
            agent_inputs4dt = self._build_inputs(ep_batch, t, task)
            if self.main_args.skill_softmax:
                logits_skill_softmax = th.nn.functional.softmax(logits_skill.reshape(1, self.main_args.step_block_size, agent_num, -1)[0,t-step_start], dim=-1)
                agent_outs, self.hidden_states_dec = self.agent.decoder(agent_inputs4dt, self.hidden_states_dec, task, logits_skill_softmax)
            
            elif self.main_args.skill_max:
                logits_skill_one_hot = th.nn.functional.one_hot(th.argmax(logits_skill, dim=-1), num_classes=logits_skill.size(-1)).float()
                agent_outs, self.hidden_states_dec = self.agent.decoder(agent_inputs4dt, self.hidden_states_dec, task, logits_skill_one_hot)
                
            return agent_outs.view(ep_batch.batch_size, agent_num, -1) 
            
            
            

    def init_hidden(self, batch_size, task):
        if self.main_args.pretrain_vae:
            n_agents = self.task2n_agents[task]
            hidden_states_enc, hidden_states_dec = self.agent.init_hidden()
            self.hidden_states_enc = hidden_states_enc.unsqueeze(0).expand(batch_size, n_agents, -1)
            self.hidden_states_dec = hidden_states_dec.unsqueeze(0).expand(batch_size, n_agents, -1)
        
        elif  self.main_args.pretrain_vqvae:
            n_agents = self.task2n_agents[task]
            hidden_states_enc, hidden_states_dec = self.vqvae_model.init_hidden()
            self.hidden_states_enc = hidden_states_enc.unsqueeze(0).expand(batch_size, n_agents, -1)
            self.hidden_states_dec = hidden_states_dec.unsqueeze(0).expand(batch_size, n_agents, -1)
        
        elif self.main_args.train_DT:
            n_agents = self.task2n_agents[task]
            hidden_states_enc, hidden_states_dec = self.agent.init_hidden()
            self.hidden_states_enc = hidden_states_enc.unsqueeze(0).expand(batch_size, n_agents, -1)
            self.hidden_states_dec = hidden_states_dec.unsqueeze(0).expand(batch_size, n_agents, -1)
        
        elif self.main_args.train_DT_w_glsk:
            n_agents = self.task2n_agents[task]
            hidden_states_enc, hidden_states_dec = self.agent.init_hidden()
            self.hidden_states_enc = hidden_states_enc.unsqueeze(0).expand(batch_size, n_agents, -1)
            self.hidden_states_dec = hidden_states_dec.unsqueeze(0).expand(batch_size, n_agents, -1)

    def parameters(self):
        return self.agent.parameters()

    def load_state(self, other_mac):
        """ we don't load the state of task dynamic decoder """
        self.agent.load_state_dict(other_mac.agent.state_dict())

    def cuda(self):
        self.agent.cuda()

    def save_models(self, path):
        """ we don't save the state of task dynamic decoder """
        if self.main_args.pretrain_vae:
            th.save(self.agent.state_dict(), "{}/agent.th".format(path))
        
        elif self.main_args.train_DT:
            th.save(self.dt_agent.state_dict(), "{}/dt_agent.th".format(path))
            
        elif self.main_args.pretrain_vqvae:
            th.save(self.vqvae_model.state_dict(), "{}/vqvae_model.th".format(path))

    def load_models(self, path):
        """ we don't load the state of task_encoder """
        self.agent.load_state_dict(th.load("{}/agent.th".format(path), map_location=lambda storage, loc: storage))

    def _build_agents(self, task2input_shape_info):
        self.agent = agent_REGISTRY[self.main_args.agent](task2input_shape_info,
                                                          self.task2decomposer, self.task2n_agents,
                                                          self.surrogate_decomposer, self.main_args)
        
        
    def _build_vqvae_agents(self, task2input_shape_info):
        self.vqvae_model = agent_REGISTRY[self.main_args.vqvae_agent](task2input_shape_info,
                                                                self.main_args, 
                                                                self.task2decomposer,
                                                                self.task2n_agents,
                                                                self.surrogate_decomposer,
                                                                self.task2args)
        
        print("initial vqvae model.")
    
    def _build_dt_agents(self, task2input_shape_info):
        self.dt_agent = agent_REGISTRY[self.main_args.dt_agent](task2input_shape_info,
                                                                self.main_args, 
                                                                self.task2decomposer,
                                                                self.task2n_agents,
                                                                self.surrogate_decomposer)
        
        print("initial decision transformer agent.")
        

    def _build_actions(self, actions):
        actions = actions.reshape(-1) - 5
        zeros = th.zeros_like(actions).to(self.main_args.device)
        actions = th.where(actions >= 0, actions, zeros)
        return actions

    def _build_inputs(self, batch, t, task):
        # Assumes homogenous agents with flat observations.
        # Other MACs might want to e.g. delegate building inputs to each agent
        bs = batch.batch_size
        inputs = []
        inputs.append(batch["obs"][:, t])
        # get args, n_agents for this specific task
        task_args, n_agents = self.task2args[task], self.task2n_agents[task]
        if task_args.obs_last_action:
            if t == 0:
                inputs.append(th.zeros_like(batch["actions_onehot"][:, t]))
            else:
                inputs.append(batch["actions_onehot"][:, t - 1])
        if task_args.obs_agent_id:
            inputs.append(th.eye(n_agents, device=batch.device).unsqueeze(0).expand(bs, -1, -1))

        inputs = th.cat([x.reshape(bs * n_agents, -1) for x in inputs], dim=1)
        return inputs
    
    def _build_inputs_4dt(self, batch, step_start, step_end, task):
        # Assumes homogenous agents with flat observations.
        # Other MACs might want to e.g. delegate building inputs to each agent
        bs = batch.batch_size
        inputs = []
        inputs.append(batch["obs"][:, step_start: step_end])
        # get args, n_agents for this specific task
        task_args, n_agents = self.task2args[task], self.task2n_agents[task]
        if self.main_args.obs_last_skill:
            if step_start == 0:
                inputs.append(th.zeros_like(batch["skill"][:, step_start: step_end]))
                # inputs.append(th.zeros_like(batch["actions_onehot"][:, t]))
            else:
                inputs.append(batch["skill"][:, step_start - 1: step_end-1])
        if task_args.obs_agent_id:
            inputs.append(th.eye(n_agents, device=batch.device).unsqueeze(0).unsqueeze(0).expand(bs,step_end-step_start, -1, -1))

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

    def _get_input_shape(self):
        task2input_shape_info = {}
        for task in self.train_tasks:
            task_scheme = self.task2scheme[task]
            input_shape = task_scheme["obs"]["vshape"]
            last_action_shape, agent_id_shape = 0, 0
            if self.task2args[task].obs_last_action:
                input_shape += task_scheme["actions_onehot"]["vshape"][0]
                last_action_shape = task_scheme["actions_onehot"]["vshape"][0]
            if self.task2args[task].obs_agent_id:
                input_shape += self.task2n_agents[task]
                agent_id_shape = self.task2n_agents[task]
            task2input_shape_info[task] = {
                "input_shape": input_shape,
                "last_action_shape": last_action_shape,
                "agent_id_shape": agent_id_shape,
            }
        return task2input_shape_info
