from discriminators.base import Discriminator
from networks.discriminator_network import Q_phi, Empowerment, Reward

import torch
import torch.nn as nn

    
class EAIRL(Discriminator):
    def __init__(self,writer, device, state_dim, action_dim, args):
        super(EAIRL, self).__init__()
        self.writer = writer
        self.device = device
        self.args = args
        self.q_phi = Q_phi(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function,self.args.last_activation,self.args.trainable_std)
        self.empowerment = Empowerment(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function ,self.args.last_activation)
        self.empowerment_t = Empowerment(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function ,self.args.last_activation)
        self.reward = Reward(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function , self.args.last_activation)
        self.empowerment_t.load_state_dict(self.empowerment.state_dict())
        self.criterion = nn.BCELoss()
        self.q_phi_optimizer = torch.optim.Adam(self.q_phi.parameters(), lr=self.args.lr)
        self.empowerment_optimizer = torch.optim.Adam(self.empowerment.parameters(), lr=self.args.lr)
        self.reward_optimizer = torch.optim.Adam(self.reward.parameters(), lr=self.args.lr)
        self.network_init()
        self.mse = nn.MSELoss()
        
        self.iter = 0
    def get_d(self,state,next_state,action,done_mask,log_prob):
        exp_f = torch.exp(self.get_f(state,next_state,action,done_mask))
        return exp_f / (exp_f+torch.exp(log_prob)) 
    
    def get_f(self,state,next_state,action,done_mask):
        return self.reward(state,action) +\
    done_mask * (self.args.gamma * self.empowerment_t(next_state).detach() - self.empowerment(state).detach())

    def get_reward(self,log_prob,state,action,next_state,done):
        done_mask = 1 - done.float()
        return (self.get_f(state,next_state,action,done_mask) - log_prob - self.args.i_lambda * self.get_loss_i(state,next_state,action,log_prob)).detach() 
    def get_loss_q(self,state,next_state,action):
        mu,sigma = self.q_phi(state,next_state)
        loss = self.mse(mu,action)
        return loss
    
    def get_loss_i(self,state,next_state,action,log_prob):
        mu,sigma = self.q_phi(state,next_state)
        dist = torch.distributions.Normal(mu,sigma)
        q_log_prob = dist.log_prob(action).sum(-1,keepdim=True).detach()
        approx_1 = self.args.beta * q_log_prob
        approx_2 = log_prob + self.empowerment(state)
        loss = self.mse(approx_1,approx_2)
        return loss
    
    def forward(self,log_prob,state,action,next_state,done_mask):
        return self.get_d(state,next_state,action,done_mask,log_prob)
        
    def train_network(self,writer,n_epi,agent_s,agent_a,agent_next_s,\
                      agent_log_prob,agent_done_mask,expert_s,expert_a,expert_next_s,\
                      expert_log_prob,expert_done_mask):
        
        loss_q = self.get_loss_q(agent_s,agent_next_s,agent_a)
        self.q_phi_optimizer.zero_grad()
        loss_q.backward()
        self.q_phi_optimizer.step()
        
        loss_i = self.get_loss_i(agent_s,agent_next_s,agent_a,agent_log_prob)
        self.empowerment_optimizer.zero_grad()
        loss_i.backward()
        self.empowerment_optimizer.step()
        
        expert_preds = self.forward(expert_log_prob,expert_s,expert_a,expert_next_s,expert_done_mask)
        expert_loss = self.criterion(expert_preds,torch.ones(expert_preds.shape[0],1).to(self.device)) 
        
        agent_preds = self.forward(agent_log_prob,agent_s,agent_a,agent_next_s,agent_done_mask)
        agent_loss = self.criterion(agent_preds,torch.zeros(agent_preds.shape[0],1).to(self.device))
        
        reward_loss = expert_loss+agent_loss
        self.reward_optimizer.zero_grad()
        reward_loss.backward()
        self.reward_optimizer.step()
        
        if writer != None:
            writer.add_scalar("loss/discriminator_loss_q", loss_q.item(), n_epi)
            writer.add_scalar("loss/discriminator_loss_i", loss_i.item(), n_epi)
            writer.add_scalar("loss/discriminator_reward_loss", reward_loss.item(), n_epi)
        if self.iter % self.args.update_cycle == 0:
            self.empowerment_t.load_state_dict(self.empowerment.state_dict())
        self.iter += 1
        
from discriminators.base import Discriminator
from networks.discriminator_network import Q_phi, Empowerment, Reward

import torch
import torch.nn as nn

    
class EAIRL(Discriminator):
    def __init__(self,writer, device, state_dim, action_dim, args):
        super(EAIRL, self).__init__()
        self.writer = writer
        self.device = device
        self.args = args
        self.q_phi = Q_phi(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function,self.args.last_activation,self.args.trainable_std)
        self.empowerment = Empowerment(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function ,self.args.last_activation)
        self.empowerment_t = Empowerment(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function ,self.args.last_activation)
        self.reward = Reward(self.args.layer_num, state_dim, action_dim, self.args.hidden_dim, self.args.activation_function , self.args.last_activation)
        self.empowerment_t.load_state_dict(self.empowerment.state_dict())
        self.criterion = nn.BCELoss()
        self.q_phi_optimizer = torch.optim.Adam(self.q_phi.parameters(), lr=self.args.lr)
        self.empowerment_optimizer = torch.optim.Adam(self.empowerment.parameters(), lr=self.args.lr)
        self.reward_optimizer = torch.optim.Adam(self.reward.parameters(), lr=self.args.lr)
        self.network_init()
        self.mse = nn.MSELoss()
        
        self.iter = 0
    def get_d(self,state,next_state,action,done_mask,log_prob):
        exp_f = torch.exp(self.get_f(state,next_state,action,done_mask))
        return exp_f / (exp_f+torch.exp(log_prob)) 
    
    def get_f(self,state,next_state,action,done_mask):
        return self.reward(state,action) +\
    done_mask * (self.args.gamma * self.empowerment_t(next_state).detach() - self.empowerment(state).detach())

    def get_reward(self,log_prob,state,action,next_state,done):
        done_mask = 1 - done.float()
        return (self.get_f(state,next_state,action,done_mask) - log_prob - self.args.i_lambda * self.get_loss_i(state,next_state,action,log_prob)).detach() 
    def get_loss_q(self,state,next_state,action):
        mu,sigma = self.q_phi(state,next_state)
        loss = self.mse(mu,action)
        return loss
    
    def get_loss_i(self,state,next_state,action,log_prob):
        mu,sigma = self.q_phi(state,next_state)
        dist = torch.distributions.Normal(mu,sigma)
        q_log_prob = dist.log_prob(action).sum(-1,keepdim=True).detach()
        approx_1 = self.args.beta * q_log_prob
        approx_2 = log_prob + self.empowerment(state)
        loss = self.mse(approx_1,approx_2)
        return loss
    
    def forward(self,log_prob,state,action,next_state,done_mask):
        return self.get_d(state,next_state,action,done_mask,log_prob)
        
    def train_network(self,writer,n_epi,agent_s,agent_a,agent_next_s,\
                      agent_log_prob,agent_done_mask,expert_s,expert_a,expert_next_s,\
                      expert_log_prob,expert_done_mask):
        
        loss_q = self.get_loss_q(agent_s,agent_next_s,agent_a)
        self.q_phi_optimizer.zero_grad()
        loss_q.backward()
        self.q_phi_optimizer.step()
        
        loss_i = self.get_loss_i(agent_s,agent_next_s,agent_a,agent_log_prob)
        self.empowerment_optimizer.zero_grad()
        loss_i.backward()
        self.empowerment_optimizer.step()
        
        expert_preds = self.forward(expert_log_prob,expert_s,expert_a,expert_next_s,expert_done_mask)
        expert_loss = self.criterion(expert_preds,torch.ones(expert_preds.shape[0],1).to(self.device)) 
        
        agent_preds = self.forward(agent_log_prob,agent_s,agent_a,agent_next_s,agent_done_mask)
        agent_loss = self.criterion(agent_preds,torch.zeros(agent_preds.shape[0],1).to(self.device))
        
        reward_loss = expert_loss+agent_loss
        self.reward_optimizer.zero_grad()
        reward_loss.backward()
        self.reward_optimizer.step()
        
        if writer != None:
            writer.add_scalar("loss/discriminator_loss_q", loss_q.item(), n_epi)
            writer.add_scalar("loss/discriminator_loss_i", loss_i.item(), n_epi)
            writer.add_scalar("loss/discriminator_reward_loss", reward_loss.item(), n_epi)
        if self.iter % self.args.update_cycle == 0:
            self.empowerment_t.load_state_dict(self.empowerment.state_dict())
        self.iter += 1
        