import os
import torch
import torch.nn.functional as F
from torch.optim import Adam
from sac.utils import soft_update, hard_update
from sac.model import GaussianPolicy, QNetwork, DeterministicPolicy
from sac.get_kernel import get_rf, get_nystrom, get_omegas_bs, get_rd_state, get_gram, get_r
import numpy as np
import time

np.random.seed(42)
class SAC(object):
    def __init__(self, num_inputs, action_space, env, args):

        self.gamma = args.gamma
        self.tau = args.tau
        self.alpha = args.alpha
        self.policy_type = args.policy
        self.target_update_interval = args.target_update_interval
        self.automatic_entropy_tuning = args.automatic_entropy_tuning

        self.device = torch.device("cuda" if args.cuda else "cpu")

        self.critic = QNetwork(num_inputs, action_space.shape[0], args.hidden_size).to(device=self.device)
        self.critic_optim = Adam(self.critic.parameters(), lr=args.lr)

        self.critic_target = QNetwork(num_inputs, action_space.shape[0], args.hidden_size).to(device=self.device)
        hard_update(self.critic_target, self.critic)
        
        
        self.env = env
        self.kernel = args.kernel
        self.sigma = args.sigma
        if args.kernel:
            self.m = args.m
            self.theta = torch.randn(args.m+1).to(device=self.device)
            self.theta = self.theta.requires_grad_(True)
            self.theta_optim = Adam([self.theta], lr=args.lr)
            if self.kernel == "rf":
                omegas, bs = get_omegas_bs(self.env, self.sigma, self.m)
                self.omegas = np.array(omegas)
                self.bs = np.array(bs)
            elif self.kernel == "nystrom":
                m_states = []
                for i in range(args.m):
                    if self.env.metadata["name"] =="simple_push_v3":
                        rd = get_rd_state(env.world)
                    elif self.env.metadata["name"] == "simple_tag_v3":
                        rd = get_tag_rd_state(env.world)
                    m_states.append(rd)
                gram = get_gram(m_states)
                self.m_states = m_states
                gram = np.array(gram)
                eigenvalues, eigenvectors = np.linalg.eig(gram)
                sorted_indices = np.argsort(eigenvalues)[::-1]
                sorted_eigenvalues = eigenvalues[sorted_indices]
                sorted_eigenvectors = eigenvectors[:, sorted_indices]
                you = np.zeros_like(sorted_eigenvectors)
                for i in range(you.shape[0]):
                    you[i][i] = sorted_eigenvalues[i]
                self.gram = np.dot(sorted_eigenvectors, you)

        if self.policy_type == "Gaussian":
            # Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
            if self.automatic_entropy_tuning is True:
                self.target_entropy = -torch.prod(torch.Tensor(action_space.shape).to(self.device)).item()
                self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
                self.alpha_optim = Adam([self.log_alpha], lr=args.lr)

            self.policy = GaussianPolicy(num_inputs, action_space.shape[0], args.hidden_size, action_space).to(self.device)
            self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)

        else:
            self.alpha = 0
            self.automatic_entropy_tuning = False
            self.policy = DeterministicPolicy(num_inputs, action_space.shape[0], args.hidden_size, action_space).to(self.device)
            self.policy_optim = Adam(self.policy.parameters(), lr=args.lr)

    def select_action(self, state, evaluate=False):
        state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
        if evaluate is False:
            action, _, _ = self.policy.sample(state)
        else:
            _, _, action = self.policy.sample(state)
        return action.detach().cpu().numpy()[0]

    def update_parameters1(self, memory, batch_size, updates, adv):
        # torch.autograd.set_detect_anomaly(True)
        state_batch, action1_batch, action2_batch, reward1_batch, reward2_batch, next_state_batch, mask_batch = memory.sample(batch_size=batch_size)

        state_batch = torch.FloatTensor(state_batch).to(self.device)
        next_state_batch = torch.FloatTensor(next_state_batch).to(self.device)
        action1_batch = torch.FloatTensor(action1_batch).to(self.device)
        action2_batch = torch.FloatTensor(action2_batch).to(self.device)
        reward1_batch = torch.FloatTensor(reward1_batch).to(self.device).unsqueeze(1)
        reward2_batch = torch.FloatTensor(reward2_batch).to(self.device).unsqueeze(1)
        mask_batch = torch.FloatTensor(mask_batch).to(self.device).unsqueeze(1)
        
        if self.kernel:
            if self.kernel == "rf":
                phi = get_rf(self.env, self.m, state_batch, action1_batch, action2_batch, self.omegas, self.bs,reward1_batch.detach().cpu().numpy(), self.device)
            elif self.kernel == "nystrom":
                phi = get_nystrom(self.env, self.m, state_batch, action1_batch, action2_batch, self.m_states, self.gram,reward1_batch.detach().cpu().numpy(), self.device)
            phi = torch.FloatTensor(phi).to(self.device)
            Q = torch.matmul(phi, self.theta)
            Q = Q.unsqueeze(1)
            with torch.no_grad():
                next_state_action, next_state_log_pi, _ = self.policy.sample(next_state_batch)
                next_state_action_adv, next_state_log_pi_adv, _ = adv.policy.sample(next_state_batch)
                next_r = get_r(self.env, next_state_batch.detach().cpu().numpy())
                next_r = np.expand_dims(next_r, axis=1)
                if self.kernel == "rf":
                    
                    phi_next = get_rf(self.env, self.m, next_state_batch, next_state_action, next_state_action_adv, self.omegas, self.bs, next_r, self.device)
                elif self.kernel =="nystrom":
                    phi_next = get_nystrom(self.env, self.m, next_state_batch, next_state_action, next_state_action_adv, self.m_states, self.gram,next_r, self.device)
                    
                phi_next = torch.FloatTensor(phi_next).to(self.device)
                next_Q = torch.matmul(phi_next, self.theta)
                next_Q = next_Q.unsqueeze(1)
                # next_target = next_Q - self.alpha * next_state_log_pi
                next_target = next_Q
                next_q_value = reward1_batch + mask_batch * self.gamma * (next_target)
            
            qf_loss = F.mse_loss(Q, next_q_value)
            self.theta_optim.zero_grad()
            qf_loss.backward()
            self.theta_optim.step()
        else:
            with torch.no_grad():
                next_state_action, next_state_log_pi, _ = self.policy.sample(next_state_batch)
                next_state_action_adv, next_state_log_pi_adv, _ = adv.policy.sample(next_state_batch)
                qf1_next_target, qf2_next_target = self.critic_target(next_state_batch, next_state_action, next_state_action_adv)
                min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - self.alpha * next_state_log_pi
                # 用phi的话，min_qf_next_target换成phi*theta就好了
                
                next_q_value = reward1_batch + mask_batch * self.gamma * (min_qf_next_target)
            qf1, qf2 = self.critic(state_batch, action1_batch, action2_batch)  # Two Q-functions to mitigate positive bias in the policy improvement step
            qf1_loss = F.mse_loss(qf1, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
            qf2_loss = F.mse_loss(qf2, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
            qf_loss = qf1_loss + qf2_loss

            self.critic_optim.zero_grad()
            qf_loss.backward()
            self.critic_optim.step()

        pi, log_pi, _ = self.policy.sample(state_batch)
        pi_adv, log_pi_adv, _ = adv.policy.sample(state_batch)

        if self.kernel:
            
            if self.kernel == "rf":
                phi = get_rf(self.env, self.m, state_batch, pi, pi_adv, self.omegas, self.bs,reward1_batch.detach().cpu().numpy(), self.device)
                
            elif self.kernel == "nystrom":
                phi = get_nystrom(self.env, self.m, state_batch, pi, pi_adv, self.m_states, self.gram,reward1_batch.detach().cpu().numpy(), self.device)
            phi = torch.FloatTensor(phi).to(self.device)
            min_qf_pi = torch.matmul(phi, self.theta)
            min_qf_pi = min_qf_pi.unsqueeze(1)
        else:
            qf1_pi, qf2_pi = self.critic(state_batch, pi,pi_adv)
            min_qf_pi = torch.min(qf1_pi, qf2_pi)

        policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean() # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]

        self.policy_optim.zero_grad()
        policy_loss.backward()
        self.policy_optim.step()

        if self.automatic_entropy_tuning:
            alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()

            self.alpha_optim.zero_grad()
            alpha_loss.backward()
            self.alpha_optim.step()

            self.alpha = self.log_alpha.exp()
            alpha_tlogs = self.alpha.clone() # For TensorboardX logs
        else:
            alpha_loss = torch.tensor(0.).to(self.device)
            alpha_tlogs = torch.tensor(self.alpha) # For TensorboardX logs

        if self.kernel == None and updates % self.target_update_interval == 0:
            soft_update(self.critic_target, self.critic, self.tau)

        return qf_loss.item(), policy_loss.item(), alpha_loss.item(), alpha_tlogs.item()

    def update_parameters2(self, memory, batch_size, updates, adv):
        # Sample a batch from memory
        state_batch, action1_batch, action2_batch, reward1_batch, reward2_batch, next_state_batch, mask_batch = memory.sample(batch_size=batch_size)

        state_batch = torch.FloatTensor(state_batch).to(self.device)
        next_state_batch = torch.FloatTensor(next_state_batch).to(self.device)
        action1_batch = torch.FloatTensor(action1_batch).to(self.device)
        action2_batch = torch.FloatTensor(action2_batch).to(self.device)
        reward1_batch = torch.FloatTensor(reward1_batch).to(self.device).unsqueeze(1)
        reward2_batch = torch.FloatTensor(reward2_batch).to(self.device).unsqueeze(1)
        mask_batch = torch.FloatTensor(mask_batch).to(self.device).unsqueeze(1)

        if self.kernel:
            if self.kernel == "rf":
                phi = get_rf(self.env, self.m, state_batch, action1_batch, action2_batch, self.omegas, self.bs,reward2_batch.detach().cpu().numpy(), self.device)
            elif self.kernel == "nystrom":
                phi = get_nystrom(self.env, self.m, state_batch, action1_batch, action2_batch, self.m_states, self.gram,reward2_batch.detach().cpu().numpy(), self.device)
            phi = torch.FloatTensor(phi).to(self.device)
            Q = torch.matmul(phi, self.theta)
            Q = Q.unsqueeze(1)
            with torch.no_grad():
                next_state_action, next_state_log_pi, _ = self.policy.sample(next_state_batch)
                next_state_action_adv, next_state_log_pi_adv, _ = adv.policy.sample(next_state_batch)
                next_r = - get_r(self.env, next_state_batch.detach().cpu().numpy())
                next_r = np.expand_dims(next_r, axis=1)
                if self.kernel =="rf":
                    phi_next = get_rf(self.env, self.m, next_state_batch, next_state_action_adv, next_state_action, self.omegas, self.bs,next_r, self.device)
                elif self.kernel =="nystrom":
                    phi_next = get_nystrom(self.env, self.m, next_state_batch, next_state_action_adv, next_state_action, self.m_states, self.gram,next_r, self.device)
                phi_next = torch.FloatTensor(phi_next).to(self.device)
                next_Q = torch.matmul(phi_next, self.theta)
                next_Q = next_Q.unsqueeze(1)
                # next_target = next_Q - self.alpha * next_state_log_pi
                next_target = next_Q
                next_q_value = reward2_batch + mask_batch * self.gamma * (next_target)
            
            
            qf_loss = F.mse_loss(Q, next_q_value)
            self.theta_optim.zero_grad()
            qf_loss.backward()
            self.theta_optim.step()

        else:
            with torch.no_grad():
                next_state_action, next_state_log_pi, _ = self.policy.sample(next_state_batch)
                next_state_action_adv, next_state_log_pi_adv, _ = adv.policy.sample(next_state_batch)
                qf1_next_target, qf2_next_target = self.critic_target(next_state_batch, next_state_action_adv, next_state_action)
                min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - self.alpha * next_state_log_pi
                next_q_value = reward2_batch + mask_batch * self.gamma * (min_qf_next_target)
            qf1, qf2 = self.critic(state_batch, action1_batch, action2_batch)  # Two Q-functions to mitigate positive bias in the policy improvement step
            qf1_loss = F.mse_loss(qf1, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
            qf2_loss = F.mse_loss(qf2, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
            qf_loss = qf1_loss + qf2_loss

            self.critic_optim.zero_grad()
            qf_loss.backward()
            self.critic_optim.step()

        pi, log_pi, _ = self.policy.sample(state_batch)
        pi_adv, log_pi_adv, _ = adv.policy.sample(state_batch)

        if self.kernel:
            if self.kernel == "rf":
                phi = get_rf(self.env, self.m, state_batch, pi_adv, pi, self.omegas, self.bs,reward2_batch.detach().cpu().numpy(), self.device)
            elif self.kernel == "nystrom":
                phi = get_nystrom(self.env, self.m, state_batch, pi_adv, pi, self.m_states, self.gram,reward2_batch.detach().cpu().numpy(), self.device)
            phi = torch.FloatTensor(phi).to(self.device)
            min_qf_pi = torch.matmul(phi, self.theta)
            min_qf_pi = min_qf_pi.unsqueeze(1)
        else:
            qf1_pi, qf2_pi = self.critic(state_batch, pi_adv,pi)
            min_qf_pi = torch.min(qf1_pi, qf2_pi)

        policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean() # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]

        self.policy_optim.zero_grad()
        policy_loss.backward()
        self.policy_optim.step()

        if self.automatic_entropy_tuning:
            alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()

            self.alpha_optim.zero_grad()
            alpha_loss.backward()
            self.alpha_optim.step()

            self.alpha = self.log_alpha.exp()
            alpha_tlogs = self.alpha.clone() # For TensorboardX logs
        else:
            alpha_loss = torch.tensor(0.).to(self.device)
            alpha_tlogs = torch.tensor(self.alpha) # For TensorboardX logs

        if updates % self.target_update_interval == 0:
            soft_update(self.critic_target, self.critic, self.tau)

        return qf_loss.item(), policy_loss.item(), alpha_loss.item(), alpha_tlogs.item()

    # Save model parameters
    def save_checkpoint(self,agent_name, suffix="", ckpt_path=None):
        if not os.path.exists('Checkpoints/'):
            os.makedirs('Checkpoints/')
        if ckpt_path is None:
            if self.kernel is None:
                ckpt_path = "Checkpoints/{}_{}_{}_{}".format(self.env.metadata["name"],agent_name , self.sigma, suffix)
            else:
                ckpt_path = "Checkpoints/{}_{}_{}_{}_{}_{}".format(self.env.metadata["name"],agent_name,  self.kernel, self.m,self.sigma, suffix)
        print('Saving models to {}'.format(ckpt_path))
        if self.kernel:
            torch.save({'policy_state_dict': self.policy.state_dict(),
                        'theta': self.theta,
                        'theta_optimizer':self.theta_optim.state_dict(),
                        'policy_optimizer_state_dict': self.policy_optim.state_dict()}, ckpt_path)
        else:
            torch.save({'policy_state_dict': self.policy.state_dict(),
                        'critic_state_dict': self.critic.state_dict(),
                        'critic_target_state_dict': self.critic_target.state_dict(),
                        'critic_optimizer_state_dict': self.critic_optim.state_dict(),
                        'policy_optimizer_state_dict': self.policy_optim.state_dict()}, ckpt_path)

    # Load model parameters
    def load_checkpoint(self, ckpt_path, evaluate=False):
        print('Loading models from {}'.format(ckpt_path))
        if ckpt_path is not None:
            checkpoint = torch.load(ckpt_path)
            self.policy.load_state_dict(checkpoint['policy_state_dict'])
            self.policy_optim.load_state_dict(checkpoint['policy_optimizer_state_dict'])
            if self.kernel:
                if 'theta' in checkpoint:
                    self.theta = checkpoint['theta']
                else:
                    raise KeyError("Key 'theta' not found in the checkpoint.")

                self.theta = checkpoint['theta']
                self.theta_optim.load_state_dict(checkpoint['theta_optimizer'])
            else:
                self.critic.load_state_dict(checkpoint['critic_state_dict'])
                self.critic_target.load_state_dict(checkpoint['critic_target_state_dict'])
                self.critic_optim.load_state_dict(checkpoint['critic_optimizer_state_dict'])

            if evaluate:
                self.policy.eval()
                self.critic.eval()
                self.critic_target.eval()
            else:
                self.policy.train()
                self.critic.train()
                self.critic_target.train()

