#sov1

import copy
from locale import ABDAY_1
from torch.distributions import Categorical
from components.episode_buffer import EpisodeBatch
from modules.critics.offpg import OffPGCritic
import torch as th
from utils.offpg_utils import build_target_q
from utils.rl_utils import build_td_lambda_targets
from torch.optim import RMSprop
from modules.mixers.qmix import QMixer
import torch.nn.functional as F
import numpy as np

import torch.distributions as D

import random


class OffPGLearnerSL:
    def __init__(self, mac, scheme, logger, args):
        self.args = args
        self.n_agents = args.n_agents
        self.n_actions = args.n_actions
        self.mac = mac
        self.logger = logger

        self.last_target_update_step = 0
        self.critic_training_steps = 0

        self.log_stats_t = -self.args.learner_log_interval - 1

        self.critic = OffPGCritic(scheme, args)
        self.mixer = QMixer(args)
        self.target_critic = copy.deepcopy(self.critic)
        self.target_mixer = copy.deepcopy(self.mixer)

        '''# dstv1
        self.mac.init_critic(self.critic)'''

        self.agent_params = list(mac.policy_parameters())#list(mac.parameters())
        self.critic_params = list(self.critic.parameters())
        self.mixer_params = list(self.mixer.parameters())
        self.comm_params = list(mac.comm_parameters())
        self.params = self.agent_params + self.critic_params
        self.c_params = self.critic_params + self.mixer_params

        self.agent_optimiser = RMSprop(params=self.agent_params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps)
        self.critic_optimiser = RMSprop(params=self.critic_params, lr=args.critic_lr, alpha=args.optim_alpha, eps=args.optim_eps)
        self.mixer_optimiser = RMSprop(params=self.mixer_params, lr=args.critic_lr, alpha=args.optim_alpha, eps=args.optim_eps)
        self.comm_optimiser = RMSprop(params=self.comm_params, lr=args.lr, alpha=args.optim_alpha, eps=args.optim_eps)

        self.s_mu = th.zeros(1)
        self.s_sigma = th.ones(1)

        self.n_cluster=3
        self.eye_tensor=th.eye(self.n_cluster).cuda()
        self.switched=False
        self.best_cate=th.ones(self.n_cluster).cuda()/self.n_cluster
        self.ret_mean_avg=[0]*self.n_cluster
        self.ret_std_avg=[0]*self.n_cluster
        self.actv_avg=[0.5]*self.n_cluster

    def get_comm_beta(self, t_env):
        comm_beta = self.args.comm_beta
        if self.args.is_comm_beta_decay and t_env > self.args.comm_beta_start_decay:
            comm_beta += 1. * (self.args.comm_beta_target - self.args.comm_beta) / \
                         (self.args.comm_beta_end_decay - self.args.comm_beta_start_decay) * \
                         (t_env - self.args.comm_beta_start_decay)
        return comm_beta

    def get_comm_entropy_beta(self, t_env):
        comm_entropy_beta = self.args.comm_entropy_beta
        if self.args.is_comm_entropy_beta_decay and t_env > self.args.comm_entropy_beta_start_decay:
            comm_entropy_beta += 1. * (self.args.comm_entropy_beta_target - self.args.comm_entropy_beta) / \
                        (self.args.comm_entropy_beta_end_decay - self.args.comm_entropy_beta_start_decay) * \
                        (t_env - self.args.comm_entropy_beta_start_decay)
        return comm_entropy_beta

    def train(self, batch: EpisodeBatch, t_env: int, log):
        # print('training')
        bs = batch.batch_size#batch['batch_size']
        max_t = batch.max_seq_length#batch['max_seq_length']
        actions = batch["actions"][:, :-1]
        terminated = batch["terminated"][:, :-1].float()
        avail_actions = batch["avail_actions"][:, :-1]
        mask = batch["filled"][:, :-1].float() 
        messages=batch['messages'][:,:-1]

        mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1])
        mask_comm=mask.clone().squeeze(2)#mask.clone().unsqueeze(3)#.unsqueeze(4)
        mask_td = mask.repeat(1, 1, self.n_agents).view(bs, -1, self.n_agents)
        mask = mask.repeat(1, 1, self.n_agents).view(-1)
        states = batch['obs'][:,:-1]
        states=states.view(states.shape[0],states.shape[1],-1)#batch["state"][:,:-1]#batch['obs'][:,:].view(actions.shape[0],2,self.n_agents*49)#
        
        inputs = self.critic._build_inputs(batch, bs, max_t)
        q_vals = self.critic.forward(inputs).detach()[:, :-1]

        mac_out = []
        mu_out = []
        mu_trans_out = [[] for _ in range(self.n_cluster)]
        cls_out = []
        sigma_out = []
        logits_out = []
        m_sample_out = []
        g_out = []
        self.mac.init_hidden(batch.batch_size)
        for t in range(batch.max_seq_length - 1):
            if self.args.comm and self.args.use_IB:
                agent_outs, (mu, sigma), logits, m_sample,mu_trans,cls = self.mac.forward(batch, t=t)#,t_env=t_env)
                mu_out.append(mu)
                for i in range(self.n_cluster):
                    mu_trans_out[i].append(mu_trans[i])
                cls_out.append(cls)
                sigma_out.append(sigma)
                logits_out.append(logits)
                m_sample_out.append(m_sample)
            else:
                agent_outs = self.mac.forward(batch, t=t)
            mac_out.append(agent_outs)
        if batch.max_seq_length==2:#todo: make it right
            label_target_max_out = th.stack(mac_out, dim=1)#todo: label may be wrong
            label_target_max_out[avail_actions[:, :] == 0] = -9999999
        else:
            label_target_max_out = th.stack(mac_out[:-1], dim=1)#todo: label may be wrong
            label_target_max_out[avail_actions[:, :-1] == 0] = -9999999
        label_target_actions = label_target_max_out.max(dim=3, keepdim=True)[1]
        
        mac_out = th.stack(mac_out, dim=1)
        if self.args.use_IB:
            if batch.max_seq_length==2:#todo: make it right
                mu_out = th.stack(mu_out, dim=1)[:, :]  # Concat over time
                for i in range(self.n_cluster):
                    mu_trans_out[i] = th.stack(mu_trans_out[i], dim=1)[:, :]
                cls_out = th.stack(cls_out, dim=1)[:, :]  # Concat over time
                sigma_out = th.stack(sigma_out, dim=1)[:, :]  # Concat over time
                logits_out = th.stack(logits_out, dim=1)[:, :]
                m_sample_out = th.stack(m_sample_out, dim=1)[:, :]
            else:
                mu_out = th.stack(mu_out, dim=1)[:, :-1]  # Concat over time
                for i in range(self.n_cluster):
                    mu_trans_out[i] = th.stack(mu_trans_out[i], dim=1)[:, :]
                cls_out = th.stack(cls_out, dim=1)[:, :]    # Concat over time
                sigma_out = th.stack(sigma_out, dim=1)[:, :-1]  # Concat over time
                logits_out = th.stack(logits_out, dim=1)[:, :-1]
                m_sample_out = th.stack(m_sample_out, dim=1)[:, :-1]

        mac_out[avail_actions == 0] = 0
        mac_out = mac_out/mac_out.sum(dim=-1, keepdim=True)
        mac_out[avail_actions == 0] = 0

        q_taken = th.gather(q_vals, dim=3, index=actions).squeeze(3).view(bs, -1, self.n_agents)
        pi = mac_out.view(-1, self.n_actions)
        
        baseline = th.sum(mac_out * q_vals, dim=-1).view(bs, -1, self.n_agents).detach()

        pi_taken = th.gather(pi, dim=1, index=actions.reshape(-1, 1)).squeeze(1)
        pi_taken[mask == 0] = 1.0
        log_pi_taken = th.log(pi_taken)
        log_pi_taken = log_pi_taken.view(bs, -1, self.n_agents)
        
        coe = self.mixer.k(states).view(bs, -1, self.n_agents)

        advantages = (q_taken - baseline)
        beta = 0.1
        advantages = F.softmax(advantages / beta, dim=0)

        coma_loss = - (coe * (len(advantages) * advantages.detach() * log_pi_taken) * mask_td).sum() / mask_td.sum()
        loss=coma_loss

        # Behavior Clone 
        # coma_loss = - (log_pi_taken * mask_td).sum() / mask_td.sum()

        if self.args.only_downstream or not self.args.use_IB:
            expressiveness_loss = th.Tensor([0.])
            compactness_loss = th.Tensor([0.])
            entropy_loss = th.Tensor([0.])
            comm_loss = th.Tensor([0.])
            comm_beta = th.Tensor([0.])
            comm_entropy_beta = th.Tensor([0.])
        else:
            # ### Optimize message
            # Message are controlled only by expressiveness and compactness loss.
            # Compute cross entropy with target q values of the same time step
            expressiveness_loss = 0
            label_prob = th.gather(logits_out, 3, label_target_actions).squeeze(3)
            expressiveness_loss += (-th.log(label_prob + 1e-6)).sum() / mask.sum()

            # Compute KL divergence
            compactness_loss = D.kl_divergence(D.Normal(mu_out, sigma_out), D.Normal(self.s_mu, self.s_sigma)).sum() / \
			                   mask.sum()

            # Entropy loss
            entropy_loss = -D.Normal(self.s_mu, self.s_sigma).log_prob(m_sample_out).sum() / mask.sum()

            # Gate loss
            gate_loss = 0

            # Total loss
            comm_beta = self.get_comm_beta(t_env)
            comm_entropy_beta = self.get_comm_entropy_beta(t_env)
            comm_loss = expressiveness_loss + comm_beta * compactness_loss + comm_entropy_beta * entropy_loss
            comm_loss *= self.args.c_beta
            loss += comm_loss
            comm_beta = th.Tensor([comm_beta])
            comm_entropy_beta = th.Tensor([comm_entropy_beta])

        if t_env>30000:
            self.agent_optimiser.zero_grad()
            loss.backward(retain_graph=True)
            grad_norm = th.nn.utils.clip_grad_norm_(self.agent_params, self.args.grad_norm_clip)
            self.agent_optimiser.step()

        if self.args.imic:
            #print(mu_trans_out[0].shape,messages.shape)
            if True or mu_trans_out[0].shape==messages[:,:-1].shape:
                #(bs,t,n,comm_embed_dim*n) probably

                if False:#mmv2d
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    '''dist=th.exp(dist)
                    softmax=(dist/dist.sum(dim=3,keepdim=True)).detach().unsqueeze(3)#(bs,t,n,1,cluster)'''
                    softmax=self.eye_tensor[dist.argmax(dim=-1)].unsqueeze(3)
                    best=(mu_trans_out*softmax).sum(dim=4)#(bs,t,n,comm_embed_dim*n)
                    loss+=0.1*((best-messages[:,:])**2).sum()

                if False:#mmv4
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3).sum(dim=2)**0.5#(bs,t,cluster)
                    softmax=self.eye_tensor[dist.argmax(dim=-1)].unsqueeze(2).unsqueeze(2)
                    best=(mu_trans_out*softmax).sum(dim=4)#(bs,t,n,comm_embed_dim*n)
                    loss+=0.1*((best-messages[:,:])**2).sum()
                if False:#mmv5
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3).sum(dim=2).sum(dim=1)**0.5#(bs,cluster)
                    softmax=self.eye_tensor[dist.argmin(dim=-1)].unsqueeze(1).unsqueeze(1).unsqueeze(1)
                    best=(mu_trans_out*softmax).sum(dim=4)#(bs,t,n,comm_embed_dim*n)
                    #loss+=1*((best-messages[:,:])**2).sum()
                    comm_loss=0.1*((best-messages[:,:])**2).sum()

                    #update best_cate
                    if t_env<30000:
                        prob=self.eye_tensor[dist.argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2
                        self.mac.comm.sov4db1=int(self.best_cate.argmax())

                    if random.random()<0.1:
                        agmx=dist.argmin(dim=-1)

                        #used in 37908
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt,self.mac.comm.sov4db1)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())
                
                if False:#mle
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    #cls_out=(bs,t,n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    #print(dist.shape,cls_out.shape)
                    comm_loss=0.1*(dist*cls_out).sum()

                    #update best_cate
                    if t_env<30000:
                        prob=self.eye_tensor[dist.sum(dim=1).sum(dim=1).argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2
                        self.mac.comm.sov4db1=int(self.best_cate.argmax())
                        self.mac.comm.prob=self.best_cate

                    if random.random()<0.1:
                        agmx=cls_out[:,1].sum(dim=1).argmax(dim=-1)

                        #used in 37908
                        cate=messages[:,1,:,:].sum(dim=1).view(-1,8,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt,self.mac.comm.sov4db1,self.mac.comm.prob)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())
                
                if False:#mlev3
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    #cls_out=(bs,t,n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    #print(dist.shape,cls_out.shape,mask_comm.shape)
                    #comm_loss=0.1*(dist*cls_out).sum()
                    #comm_loss=0.1*(dist.mean(dim=1).mean(dim=1)*cls_out.mean(dim=1).mean(dim=1)).sum()#mlev3
                    comm_loss=0.1*(dist.sum(dim=1).sum(dim=1)*cls_out.sum(dim=1).sum(dim=1)).sum()#mlev3sum
                    #comm_loss=0.1*((dist*mask_comm).mean(dim=1).mean(dim=1)*(cls_out*mask_comm).mean(dim=1).mean(dim=1)).sum()#mlev4

                    if random.random()<0.1:
                        agmx=cls_out.mean(dim=1).mean(dim=1).argmax(dim=-1)

                        #used in 37908
                        cate=messages[:,1,:,:].sum(dim=1).view(-1,8,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())
                
                if False:#mlev6
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    #cls_out=(bs,t,n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    #print(dist.shape,cls_out.shape)
                    #comm_loss=0.1*(dist*cls_out).sum() if t_env<30000 else 0.1*(dist*self.eye_tensor[cls_out.argmax(dim=-1)]).sum()#mlev7
                    comm_loss=0.1*(dist*cls_out).sum() if t_env<30000 else 0.1*(dist.sum(dim=2)*self.eye_tensor[cls_out.sum(dim=2).argmax(dim=-1)]).sum()#mlev8

                    #update best_cate
                    if t_env<30000:
                        prob=self.eye_tensor[dist.sum(dim=1).sum(dim=1).argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2
                        self.mac.comm.sov4db1=int(self.best_cate.argmax())
                        self.mac.comm.prob=self.best_cate

                    if random.random()<0.1:
                        agmx=cls_out[:,1].sum(dim=1).argmax(dim=-1)

                        #used in 37908
                        #cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        #used in 379082
                        cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)>0
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt,self.mac.comm.sov4db1,self.mac.comm.prob)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())
                
                if False:#mlev11
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    #cls_out=(bs,t,n,cluster)
                    dist_closest=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    dist=(mu_trans_out*cls_out.unsqueeze(3)).sum(dim=4)-messages[:,:]#(bs,t,n,comm_embed_dim*n)
                    dist=(dist**2).sum(dim=3)**0.5#(bs,t,n)
                    #print(dist.shape,cls_out.shape)
                    #comm_loss=0.1*(dist*cls_out).sum() if t_env<30000 else 0.1*(dist*self.eye_tensor[cls_out.argmax(dim=-1)]).sum()#mlev7
                    comm_loss=dist.mean()#0.1*dist.sum()# if t_env<30000 else 0.1*(dist_closest.sum(dim=2)*self.eye_tensor[cls_out.sum(dim=2).argmax(dim=-1)]).sum()#mlev8
                    
                    #mlev14
                    cls_mean=cls_out.mean(dim=0)
                    entropy_loss=(cls_mean*cls_mean.log()).mean()
                    comm_loss+=entropy_loss#10*entropy_loss
                    
                    #deter_entropy_loss=(cls_out*cls_out.log()).sum(dim=3).mean(dim=2).sum()
                    #comm_loss+=0.001*deter_entropy_loss

                    cls_argmax=(cls_out.mean(dim=2).argmax(dim=2)+1)#(bs,t)
                    cls_argmax=cls_argmax*mask_comm.long()
                    ret=batch['reward'].sum(dim=1).repeat(1,cls_argmax.shape[1])
                    ret_mean,ret_std=[],[]
                    for i in range(self.n_cluster):
                        if (cls_argmax==i+1).sum()>cls_argmax.shape[1]:
                            fom=ret[cls_argmax==i+1].sum()/(cls_argmax==i+1).sum().float()
                            som=(ret[cls_argmax==i+1]**2).sum()/(cls_argmax==i+1).sum().float()
                            ret_mean.append(float(fom.cpu()))
                            ret_std.append(float(((som-fom**2)**0.5).cpu()))
                        else:
                            ret_mean.append(0)
                            ret_std.append(0)
                    for i in range(self.n_cluster):
                        self.ret_mean_avg[i]=self.ret_mean_avg[i]*0.8+ret_mean[i]*0.2
                        self.ret_std_avg[i]=self.ret_std_avg[i]*0.8+ret_std[i]*0.2
                    highest=np.argmax(self.ret_mean_avg)
                    actv=[1 if i==highest or (self.ret_mean_avg[i]-self.ret_mean_avg[highest])**2<self.ret_std_avg[i]**2+self.ret_std_avg[highest]**2 else 0 for i in range(self.n_cluster)]
                    self.mac.comm.actv=th.tensor(actv).cuda()

                    #update best_cate
                    if True or t_env<30000:
                        prob=self.eye_tensor[dist_closest.sum(dim=1).sum(dim=1).argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2
                        self.mac.comm.sov4db1=int(self.best_cate.argmax())
                        self.mac.comm.prob=self.best_cate

                    if random.random()<0.1:
                        agmx=cls_out[:,1].sum(dim=1).argmax(dim=-1)

                        #used in 37908
                        #cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        #used in 379082
                        cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)>0

                        #used in 9001014
                        #ret=batch['reward'][:,:-1].sum(dim=2).sum(dim=1)
                        #cate=ret>10
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt,self.mac.comm.sov4db1,self.mac.comm.prob,self.ret_mean_avg,self.ret_std_avg,actv)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())
                
                if False:#mlev19
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    #cls_out=(bs,t,n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    #print(dist.shape,cls_out.shape)
                    #comm_loss=(dist*cls_out).sum(dim=3).mean()
                    comm_loss=(dist*cls_out).sum(dim=3).mean() if t_env<100000 else (dist*cls_out.detach()*self.cls_mask).sum(dim=3).mean()#mlev20
                    #comm_loss=(dist*cls_out).sum(dim=3).mean() if t_env<100000 else (dist.sum(dim=2)*self.eye_tensor[cls_out.sum(dim=2).argmax(dim=-1)]).sum()#mlev21
                    
                    if t_env<100000:
                        #mlev14
                        cls_mean=cls_out.mean(dim=0)
                        entropy_loss=(cls_mean*cls_mean.log()).mean()
                        comm_loss+=entropy_loss#10*entropy_loss
                        
                        #rubbish
                        #deter_entropy_loss=(cls_out*cls_out.log()).sum(dim=3).mean(dim=2).sum()
                        #comm_loss+=0.001*deter_entropy_loss

                        cls_argmax=(cls_out.mean(dim=2).argmax(dim=2)+1)#(bs,t)
                        cls_argmax=cls_argmax*mask_comm.long()
                        ret=batch['reward'].sum(dim=1).repeat(1,cls_argmax.shape[1])
                        ret_mean,ret_std=[],[]
                        for i in range(self.n_cluster):
                            if (cls_argmax==i+1).sum()>cls_argmax.shape[1]:
                                fom=ret[cls_argmax==i+1].sum()/(cls_argmax==i+1).sum().float()
                                som=(ret[cls_argmax==i+1]**2).sum()/(cls_argmax==i+1).sum().float()
                                ret_mean.append(float(fom.cpu()))
                                ret_std.append(float(((som-fom**2)**0.5).cpu()))
                            else:
                                ret_mean.append(0)
                                ret_std.append(0)
                        for i in range(self.n_cluster):
                            self.ret_mean_avg[i]=self.ret_mean_avg[i]*0.8+ret_mean[i]*0.2
                            self.ret_std_avg[i]=self.ret_std_avg[i]*0.8+ret_std[i]*0.2
                        highest=np.argmax(self.ret_mean_avg)
                        actv=[1 if i==highest or (self.ret_mean_avg[i]-self.ret_mean_avg[highest])**2<self.ret_std_avg[i]**2+self.ret_std_avg[highest]**2 else 0 for i in range(self.n_cluster)]
                        self.mac.comm.actv=th.tensor(actv).cuda()

                        self.cls_mask=th.tensor(actv).float().unsqueeze(0).unsqueeze(0).unsqueeze(0).cuda()

                    #update best_cate
                    if True or t_env<30000:
                        prob=self.eye_tensor[dist.sum(dim=1).sum(dim=1).argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2
                        self.mac.comm.sov4db1=int(self.best_cate.argmax())
                        self.mac.comm.prob=self.best_cate

                    if random.random()<0.1:
                        agmx=cls_out[:,1].sum(dim=1).argmax(dim=-1)

                        #used in 37908
                        #cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        #used in 379082
                        #cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)>0

                        #used in 9001014
                        ret=batch['reward'][:,:-1].sum(dim=2).sum(dim=1)
                        cate=ret>10
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt,self.mac.comm.sov4db1,self.mac.comm.prob,self.ret_mean_avg,self.ret_std_avg,self.cls_mask)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())

                if True:#mlev23
                    for i in range(self.n_cluster):
                        mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                    mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                    #cls_out=(bs,t,n,cluster)
                    dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    #print(dist.shape,cls_out.shape)
                    #comm_loss=(dist*cls_out).sum(dim=3).mean()#mlev23
                    #comm_loss=((dist*cls_out).sum(dim=3)*mask_comm.unsqueeze(2)).mean()#mlev25
                    
                    softmax=self.eye_tensor[dist.argmin(dim=-1)]#(bs,t,n,cluster)
                    comm_loss=((dist*softmax).sum(dim=3)*mask_comm.unsqueeze(2)).mean()

                    #mlev26
                    '''dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3)**0.5#(bs,t,n,cluster)
                    dist_loss=(mu_trans_out*cls_out.unsqueeze(3)).sum(dim=4)-messages[:,:]#(bs,t,n,comm_embed_dim*n)
                    dist_loss=(dist_loss**2).sum(dim=3)**0.5#(bs,t,n)
                    comm_loss=(dist_loss*mask_comm.unsqueeze(2)).mean()'''

                    '''#entropy mlev24
                    cls_mean=cls_out.mean(dim=0)
                    entropy_loss=(cls_mean*cls_mean.log()).mean()
                    comm_loss+=entropy_loss#10*entropy_loss'''

                    #update best_cate
                    if True:
                        prob=self.eye_tensor[dist.sum(dim=1).sum(dim=1).argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2
                        self.mac.comm.sov4db1=int(self.best_cate.argmax())
                        self.mac.comm.prob=self.best_cate

                    #check activation
                    cls_argmax=(cls_out.mean(dim=2).argmax(dim=2)+1)#(bs,t)
                    cls_argmax=cls_argmax*mask_comm.long()
                    for i in range(self.n_cluster):
                        if (cls_argmax==i+1).sum()>cls_argmax.shape[1]:#(cls_argmax==i+1).sum().float()>mask_comm.sum()**0.1:#actv1#(cls_argmax==i+1).sum()>cls_argmax.shape[1]:#actv2#
                            self.actv_avg[i]=self.actv_avg[i]*0.8+1*0.2
                        else:
                            self.actv_avg[i]=self.actv_avg[i]*0.8+0*0.2
                    self.mac.comm.actv=(th.tensor(self.actv_avg).cuda())>0.1#0.5
                    #self.mac.comm.actv=(th.tensor(self.actv_avg).cuda())>0.1#actv2
                    #self.mac.comm.actv=(th.tensor(self.best_cate).cuda())>0.1#actv3

                    if random.random()<0.1:
                        agmx=cls_out[:,1].sum(dim=1).argmax(dim=-1)

                        #used in 37908
                        #cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)==0
                        
                        #used in 37208
                        #cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0
                        
                        #used in 379082
                        #cate=messages[:,1,:,:].sum(dim=1).view(-1,self.n_agents,3)[:,:,0].sum(dim=1)>0

                        #used in 9001014
                        ret=batch['reward'][:,:-1].sum(dim=2).sum(dim=1)
                        cate=ret>10
                        
                        prt=[]
                        for j in range(2):
                            new=[]
                            for i in range(self.n_cluster):
                                new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                            prt.append(new)
                        print(prt,self.mac.comm.sov4db1,self.mac.comm.prob,self.actv_avg)
                        #print('type 1',(agmx==0).float().mean(),'type 2',(agmx==1).float().mean(),'type 3',(agmx==2).float().mean())
                
                if False:#sov2
                    if t_env<30000:
                        for i in range(self.n_cluster):
                            mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                        mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                        dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3).sum(dim=2).sum(dim=1)**0.5#(bs,cluster)
                        softmax=self.eye_tensor[dist.argmin(dim=-1)].unsqueeze(1).unsqueeze(1).unsqueeze(1)
                        best=(mu_trans_out*softmax).sum(dim=4)#(bs,t,n,comm_embed_dim*n)
                        #loss+=1*((best-messages[:,:])**2).sum()
                        comm_loss=0.1*((best-messages[:,:])**2).sum()

                        #update best_cate
                        prob=self.eye_tensor[dist.argmin(dim=-1)].mean(dim=0)
                        prob=prob/((prob**2).sum())**0.5
                        self.best_cate=self.best_cate*0.8+prob*0.2

                        if random.random()<0.1:
                            agmx=dist.argmin(dim=-1)
                            cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0#[:,:,0].sum(dim=1)==0
                            prt=[]
                            for j in range(2):
                                new=[]
                                for i in range(self.n_cluster):
                                    new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                                prt.append(new)
                            print(prt)
                    elif not self.switched:
                        self.mac.comm.switch_mode()
                        self.switched=True
                        self.best_cate=int(self.best_cate.argmax())
                        print('best_cate',self.best_cate)
                    else:
                        for i in range(self.n_cluster):
                            mu_trans_out[i]=mu_trans_out[i].unsqueeze(4)
                        mu_trans_out=th.cat(mu_trans_out,dim=4)#(bs,t,n,comm_embed_dim*n,cluster)
                        dist=((mu_trans_out-messages[:,:].unsqueeze(4))**2).sum(dim=3).sum(dim=2).sum(dim=1)**0.5#(bs,cluster)
                        '''softmax=self.eye_tensor[dist.argmin(dim=-1)].unsqueeze(1).unsqueeze(1).unsqueeze(1)
                        best=(mu_trans_out*softmax).sum(dim=4)#(bs,t,n,comm_embed_dim*n)
                        #loss+=1*((best-messages[:,:])**2).sum()
                        comm_loss=0.1*((best-messages[:,:])**2).sum()'''

                        softmax=dist.argmin(dim=-1)==self.best_cate#(bs)
                        comm_loss=0.1*((mu_out-messages[:,:-1])**2)[softmax].sum()

                        if random.random()<0.1:
                            agmx=dist.argmin(dim=-1)
                            cate=messages[:,:,:,:].sum(dim=1).sum(dim=1).view(-1,8,3)[:,:,2].sum(dim=1)==0#[:,:,0].sum(dim=1)==0
                            prt=[]
                            for j in range(2):
                                new=[]
                                for i in range(self.n_cluster):
                                    new.append(float(((agmx==i)*(cate==j)).float().mean().cpu()))
                                prt.append(new)
                            print(prt,self.best_cate)

                #loss+=0.1*((mu_trans_out[0]-messages[:,:])**2).sum()
                self.comm_optimiser.zero_grad()
                comm_loss.backward()
                #grad_norm = th.nn.utils.clip_grad_norm_(self.agent_params, self.args.grad_norm_clip)
                self.comm_optimiser.step()

        p_sum = 0.
        for p in self.agent_params:
            p_sum += p.data.abs().sum().item() / 100.0

        if t_env - self.log_stats_t >= self.args.learner_log_interval:
            self.logger.log_stat("coma_loss", coma_loss.item(), t_env)
            self.logger.log_stat("agent_grad_norm", grad_norm, t_env)
            self.logger.log_stat("pi_max", (pi.max(dim=1)[0] * mask).sum().item() / mask.sum().item(), t_env)
            self.logger.log_stat("loss", loss.item(), t_env)
            self.logger.log_stat("comm_loss", comm_loss.item(), t_env)
            self.logger.log_stat("exp_loss", expressiveness_loss.item(), t_env)
            self.logger.log_stat("comp_loss", compactness_loss.item(), t_env)
            self.logger.log_stat("comm_beta", comm_beta.item(), t_env)
            self.logger.log_stat("entropy_loss", entropy_loss.item(), t_env)
            self.logger.log_stat("comm_beta", comm_beta.item(), t_env)
            self.logger.log_stat("comm_entropy_beta", comm_entropy_beta.item(), t_env)
            self.logger.log_stat("grad_norm", grad_norm, t_env)
            self.log_stats_t = t_env

    def train_critic(self, on_batch, best_batch=None, log=None, t_env=None):
        bs = on_batch.batch_size
        max_t = on_batch.max_seq_length
        rewards = on_batch["reward"][:, :-1]
        actions = on_batch["actions"][:, :]
        terminated = on_batch["terminated"][:, :-1].float()
        mask = on_batch["filled"][:, :-1].float()
        mask[:, 1:] = mask[:, 1:] * (1 - terminated[:, :-1])
        avail_actions = on_batch["avail_actions"][:]
        states = on_batch['obs'][:,:]
        states=states.view(states.shape[0],states.shape[1],-1)#on_batch["state"]#on_batch['obs'][:,:].view(actions.shape[0],2,self.n_agents*49)#

        target_inputs = self.target_critic._build_inputs(on_batch, bs, max_t)
        target_q_vals = self.target_critic.forward(target_inputs).detach()
        
        # -----------------------------Q_lambda-IS-----------------------
        target_q_taken = th.gather(target_q_vals, dim=3, index=actions).squeeze(3)
        target_q_vals_IS = self.target_mixer(target_q_taken, states) 
        beta = 1000
        advantage_Q = F.softmax(target_q_vals_IS / beta, dim=0)
        targets_taken = self.target_mixer(th.gather(target_q_vals, dim=3, index=actions).squeeze(3), states)
        targets_taken = len(advantage_Q) * advantage_Q * targets_taken
        # -----------------------------Q_lambda-IS-----------------------
        
        target_q = build_td_lambda_targets(rewards, terminated, mask, targets_taken, self.n_agents, self.args.gamma, self.args.td_lambda).detach()
        inputs = self.critic._build_inputs(on_batch, bs, max_t)

        mac_out = []
        self.mac.init_hidden(bs)
        for i in range(max_t):
            if self.args.comm and self.args.use_IB:
                agent_outs, (mu, sigma), logits, m_sample,_,_ = self.mac.forward(on_batch, t=i)#,t_env=0)
                #mu_out.append(mu)
                #sigma_out.append(sigma)
                #logits_out.append(logits)
                #m_sample_out.append(m_sample)
            else:
                agent_outs = self.mac.forward(on_batch, t=i)
            mac_out.append(agent_outs)
        mac_out = th.stack(mac_out, dim=1).detach()
        mac_out[avail_actions == 0] = 0
        mac_out = mac_out / mac_out.sum(dim=-1, keepdim=True)
        mac_out[avail_actions == 0] = 0

        mac_out = mac_out.detach()
        for t in range(max_t - 1):
            mask_t = mask[:, t:t+1]
            if mask_t.sum() < 0.5:
                continue
            q_vals = self.critic.forward(inputs[:, t:t+1])
            q_test = q_vals
            q_vals = th.gather(q_vals, 3, index=actions[:, t:t+1]).squeeze(3)
            q_vals = self.mixer.forward(q_vals, states[:, t:t+1])

            #print('s,q',states[0,t],q_vals[0,0,0])

            target_q_t = target_q[:, t:t+1].detach()
            q_err = (q_vals - target_q_t) * mask_t
            critic_loss = (q_err ** 2).sum() / mask_t.sum()

            self.critic_optimiser.zero_grad()
            self.mixer_optimiser.zero_grad()
            critic_loss.backward()
            grad_norm = th.nn.utils.clip_grad_norm_(self.c_params, self.args.grad_norm_clip)
            self.critic_optimiser.step()
            self.mixer_optimiser.step()
            self.critic_training_steps += 1

            log["critic_loss"].append(critic_loss.item())
            log["critic_grad_norm"].append(grad_norm)
            mask_elems = mask_t.sum().item()
            log["td_error_abs"].append((q_err.abs().sum().item() / mask_elems))
            log["target_mean"].append((target_q_t * mask_t).sum().item() / mask_elems)
            log["q_taken_mean"].append((q_vals * mask_t).sum().item() / mask_elems)
            log["q_max_mean"].append((th.mean(q_test.max(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems)
            log["q_min_mean"].append((th.mean(q_test.min(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems)
            log["q_max_var"].append((th.var(q_test.max(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems)
            log["q_min_var"].append((th.var(q_test.min(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems)

            if (t == 0):
                log["q_max_first"] = (th.mean(q_test.max(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems
                log["q_min_first"] = (th.mean(q_test.min(dim=3)[0], dim=2, keepdim=True) * mask_t).sum().item() / mask_elems

        if (self.critic_training_steps - self.last_target_update_step) / self.args.target_update_interval >= 1.0:
            self._update_targets()
            self.last_target_update_step = self.critic_training_steps

        if t_env - self.log_stats_t >= self.args.learner_log_interval:
            self.logger.log_stat("critic_loss", critic_loss.item(), t_env)
            self.logger.log_stat("q_taken_mean", (q_vals * mask_t).sum().item() / mask_elems, t_env)
            self.logger.log_stat("beta_q", beta, t_env)
            self.log_stats_t = t_env

    def _update_targets(self):
        self.target_critic.load_state_dict(self.critic.state_dict())
        self.target_mixer.load_state_dict(self.mixer.state_dict())
        self.logger.console_logger.info("Updated target network")

    def cuda(self):
        self.mac.cuda()
        self.critic.cuda()
        self.mixer.cuda()
        self.target_critic.cuda()
        self.target_mixer.cuda()
        self.s_mu = self.s_mu.cuda()
        self.s_sigma = self.s_sigma.cuda()

    def save_models(self, path):
        self.mac.save_models(path)
        if self.mixer is not None:
            th.save(self.mixer.state_dict(), "{}/mixer.th".format(path))
        #th.save(self.optimiser.state_dict(), "{}/opt.th".format(path))

    def load_models(self, path):
        self.mac.load_models(path)
        # Not quite right but I don't want to save target networks
        self.target_mac.load_models(path)
        if self.mixer is not None:
            self.mixer.load_state_dict(th.load("{}/mixer.th".format(path), map_location=lambda storage, loc: storage))
        #self.optimiser.load_state_dict(th.load("{}/opt.th".format(path), map_location=lambda storage, loc: storage))
