import torch
import torchvision
import torch.nn as nn
import numpy as np
from torch.distributions import Categorical
import torch.nn.utils as utils
from torch.autograd import Variable
class AriaAC_ER(nn.Module):
    def __init__(self,opt_params, with_memory=False, aidi=None):
        super(AriaAC_ER, self).__init__()
        self.aidi = aidi
        self.batch_size = opt_params["batch_size"]
        self.gamma = opt_params["gamma"]
        self.vocabulary_size = opt_params["vocab_size"]
        self.epsilon = opt_params["eps"]
        self.replay_size = opt_params["replay_size"]
        self.training_loops = opt_params["training_loops"]
        self.eps = np.finfo(np.float32).eps.item()
        self.hiddenDim = 8
        self.memory_size = 8
        self.with_memory = with_memory
        self.obs_Mod = lin_Mod([2, self.hiddenDim//2])
        self.msg_Enc = lin_Mod([4, self.hiddenDim//2, self.hiddenDim//2], sftmx = False)
        self.rep_Mod = lin_Mod([self.hiddenDim, self.hiddenDim])
        if self.with_memory:
            self.memory = nn.LSTMCell(self.hiddenDim, self.memory_size)
            self.memories = [torch.zeros([1, 2*self.memory_size], dtype=torch.float32) for _ in range(self.replay_size+1)]
            self.action_Mod = nn.Sequential(lin_Mod([self.memory_size, 1], sftmx = False), nn.Sigmoid())
            self.msg_Dec = lin_Mod([self.memory_size, self.memory_size, self.vocabulary_size], sftmx=True)
            self.value_head = lin_Mod([self.memory_size, self.memory_size, 1])
            self.hidden_statesTable = np.zeros(self.replay_size)
        else : 
            self.memory = None
            self.memories = None
            self.action_Mod = lin_Mod([self.hiddenDim, 1], sftmx = True)
            self.value_head = lin_Mod([self.hiddenDim, 1])

        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt_params["lr"])

        self.saved_act_Logp = []
        self.saved_values = []
        self.saved_entropies = []
        self.saved_msg_Logp = []
        self.saved_hid_states = []
        self.saved_rewards = []
        self.saved_obs = []
        self.saved_downlink_msgs = []
        self.batch_counter = 0
        
    def forward(self, obs, msg, memory):
        o = self.obs_Mod(obs)
        m = self.msg_Enc(msg)
        z = self.rep_Mod(torch.cat([o, m], -1))
        if self.with_memory:
            hz, cz = self.memory(z, (memory[:, :self.memory_size], memory[:, self.memory_size:]))
            
            out_memory = torch.cat([hz, cz], dim=1)
        else:
            hz = z
            out_memory = None
        action = self.action_Mod(hz)
        message = self.msg_Dec(hz)
        value = self.value_head(hz)
        return action, message, out_memory, value

    def select_action(self, obs, msg):

        obs_t = torch.tensor([obs], dtype=torch.float32)
        msg_t = torch.tensor([msg], dtype=torch.float32)
        
        if self.with_memory:
            action, message, hid_state, value = self.forward(obs_t.float(), msg_t.float(), self.memories[self.hidden_statesTable[-1]])
            self.memories[self.replay_counter+1] = hid_state
        else:
            action, message, _, value = self.forward(obs_t.float(), msg_t.float(), None)

        a_distrib = Categorical(torch.cat([action, 1-action], -1))
        m_distrib = Categorical(message)
        a = a_distrib.sample()
        m = m_distrib.sample()
        a_entropy = a_distrib.entropy() 
        m_entropy = m_distrib.entropy()
        self.pushBufferSelect(a_distrib.log_prob(a), m_distrib.log_prob(m), value, a_entropy + m_entropy)
        return a, m


    def train_on_batch(self, state, reward):
        #self.popBuffer()
        self.pushBufferEpisode(state[0], state[1], reward)
        self.batch_counter += 1
        if self.batch_counter >= self.batch_size:
            
            self.optimizer.zero_grad()
            returns = self.getReturns(normalize=True)
            returns = torch.tensor(returns)
            rewards = torch.tensor(self.saved_rewards)
            policy_loss = 0
            value_loss = 0
            for i in range(self.batch_size-1):
                advantage = rewards[i]+self.saved_values[i+1].item()-self.saved_values[i].item()
                policy_loss += -(self.saved_act_Logp[i] + self.saved_msg_Logp[i])*advantage.detach()        
                value_loss += advantage.pow(2)
            """for a_logp, m_logp, ret, val, entro in zip(self.saved_act_Logp, self.saved_msg_Logp, returns, self.saved_values, self.saved_entropies):
                advantage = ret - val.item()
                policy_loss += -(a_logp + m_logp)*advantage.detach()        
                value_loss += advantage.pow(2)"""
            policy_loss /= self.batch_size
            value_loss /= self.batch_size
            entropy_loss =  -self.epsilon*torch.cat(self.saved_entropies).sum()
            loss = policy_loss+ value_loss + entropy_loss
            loss.backward(retain_graph=True)
            self.optimizer.step()
            mean_policy = torch.cat(self.saved_act_Logp, 0).exp().mean(dim=0)
            rewards = np.copy(self.saved_rewards)
            self.reset_Buffer()
            return np.round([policy_loss.item(), value_loss.item(), entropy_loss.item()], 2), rewards, mean_policy
        return None, None, None

    def getReturns(self, normalize=False):
        _R = 0
        Gs = []
        for r in self.saved_rewards[:self.batch_size]:
            _R = r + self.gamma * _R
            Gs.insert(0, _R)
        if normalize==True:
            return (Gs-np.mean(Gs))/np.std(Gs)
        return Gs
    
    
    def popBuffer(self):
        self.saved_act_Logp[:-1] = self.saved_act_Logp[1:]
        self.saved_values[:-1] = self.saved_values[1:]
        self.saved_entropies[:-1] = self.saved_entropies[1:]
        self.saved_msg_Logp[:-1] = self.saved_msg_Logp[1:]
        self.saved_hid_states[:-1] = self.saved_hid_states[1:]
        self.saved_rewards[:-1] = self.saved_rewards[1:]
        self.saved_obs[:-1] = self.saved_obs[1:]
        self.saved_downlink_msgs[:-1] = self.saved_downlink_msgs[1:]
        self.memories[self.hidden_statesTable[0]] = torch.zeros([1, 2*self.memory_size], dtype=torch.float32)
        self.memories[self.hidden_statesTable[1]] = self.hidden_statesTable[1].detach() 
        self.hidden_statesTable[:-1] = self.hidden_statesTable[1:]
    def sampleBatch(self):
        indices = np.random.randint(0, self.replay_size, self.batch_size)
        return indices
    def pushBufferEpisode(self, obs, msg, reward):
        self.saved_obs.append(obs)
        self.saved_downlink_msgs.append(msg)
        self.saved_rewards.append(reward)

    def pushBufferSelect(self, a_lp, m_lp, val, ent):
        self.saved_act_Logp.append(a_lp)
        self.saved_msg_Logp.append(m_lp)
        self.saved_values.append(val)
        self.saved_entropies.append(ent)
    def reset_Buffer(self):
        self.memories[0] = self.memories[-1].detach()
        self.memories[1:] = [torch.zeros([1, 2*self.memory_size], dtype=torch.float32) for _ in range(self.batch_size)]
        self.saved_act_Logp = []
        self.saved_values = []
        self.saved_entropies = []
        self.saved_msg_Logp = []
        self.saved_hid_states = []
        self.saved_rewards = []
        self.saved_obs = []
        self.saved_downlink_msgs = []
        self.batch_counter = 0
class lin_Mod(nn.Module):
    def __init__(self, sizes, sftmx = False):
        super(lin_Mod, self).__init__()
        self.sizes = sizes
        L = []
        
        for i, s in enumerate(sizes[:-1]):
            L.append(nn.Linear(sizes[i], sizes[i+1]))
            if i==len(sizes)-2 and sftmx==True:
                L.append(nn.Softmax(dim=-1))
            else :
                L.append(nn.ReLU())
        self.mod = nn.ModuleList(L)
    
    def forward(self, x):
        x_ = x
        for m in self.mod:
            x_ = m(x_)
        return x_
class ariaActor(nn.Module):
    def __init__(self, hidden_dim=10):
        super(ariaActor, self).__init__()
        self.hiddenDim = hidden_dim
        self.obs_Mod = lin_Mod([2+self.hiddenDim, 5])
        self.action_Mod = lin_Mod([self.hiddenDim, 2], sftmx = True)
        self.msg_Enc = lin_Mod([4, 5], sftmx = False)
        self.msg_Dec = lin_Mod([self.hiddenDim, 4], sftmx=True)
        self.rep_Mod = lin_Mod([self.hiddenDim, self.hiddenDim])
       
    def forward(self, obs, msg, last_state):
        inO = torch.cat([obs, last_state[self.batch_counter]], -1)
        o = self.obs_Mod(inO)
        m = self.msg_Enc(msg)
        new_state = self.rep_Mod(torch.cat([o, m], -1))
        action = self.action_Mod(new_state)
        message = self.msg_Dec(new_state)

        return action, message, new_state
