import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from ES import Searcher
from torch.distributions import Normal
import math
import matplotlib.pyplot as plt
import datetime
import os
# Implementation of Self-Guided and Self-Regularized Actor-Critic Algorithm.


epsilon = 1e-6


class Actor(nn.Module):
	def __init__(self, state_dim, action_dim, max_action):
		super(Actor, self).__init__()
		self.l1 = nn.Linear(state_dim, 256)
		self.l2 = nn.Linear(256, 256)
		self.l3_sigma = nn.Linear(256, action_dim)
		self.l3_mean = nn.Linear(256, action_dim)
		
		self.max_action = max_action
		self.action_dim = action_dim
		self.state_dim = state_dim

	def forward(self, state, *args):
		a = F.relu(self.l1(state))
		a = F.relu(self.l2(a))
		mean = self.max_action * torch.tanh(self.l3_mean(a))
		sigma = F.softplus(self.l3_sigma(a)) + 0.001
		normal = Normal(mean, sigma)
		action = normal.rsample().clamp(-self.max_action, self.max_action)
		return action

	def forward_all(self, state, *args):
		a = F.relu(self.l1(state))
		a = F.relu(self.l2(a))
		mean = self.max_action * torch.tanh(self.l3_mean(a))
		sigma = F.softplus(self.l3_sigma(a)) + 0.001
		normal = Normal(mean, sigma)
		action = normal.rsample()
		log_prob = normal.log_prob(action).sum(1,keepdim=True)
		action = action.clamp(-self.max_action, self.max_action)
		return action, log_prob, mean, sigma

	def forward_all_sample(self, state, *args):
		a = F.relu(self.l1(state))
		a = F.relu(self.l2(a))
		mean = self.max_action * torch.tanh(self.l3_mean(a))
		sigma = F.softplus(self.l3_sigma(a)) + 0.001
		normal = Normal(mean, sigma)
		action1 = normal.rsample().clamp(-self.max_action, self.max_action)
		action2 = normal.rsample().clamp(-self.max_action, self.max_action)
		prob1 = normal.log_prob(action1).sum(1,keepdim=True)
		prob2 = normal.log_prob(action2).sum(1,keepdim=True)
		probm = normal.log_prob(mean).sum(1,keepdim=True)
		prob1 = torch.exp(prob1)
		prob2 = torch.exp(prob2)
		probm = torch.exp(probm)
		return action1, action2, mean, sigma, prob1, prob2, probm

class Critic(nn.Module):
	def __init__(self, state_dim, action_dim):
		super(Critic, self).__init__()

		# Q1 architecture
		self.l1 = nn.Linear(state_dim + action_dim, 256)
		self.l2 = nn.Linear(256, 256)
		self.l3 = nn.Linear(256, 1)

	def forward(self, state, action):

		sa = torch.cat([state, action], len(action.shape)-1)

		q1 = F.relu(self.l1(sa))
		q1 = F.relu(self.l2(q1))
		q1 = self.l3(q1)
		return q1


	def Q1(self, state, action):
		sa = torch.cat([state, action], len(action.shape)-1)

		q1 = F.relu(self.l1(sa))
		q1 = F.relu(self.l2(q1))
		q1 = self.l3(q1)
		return q1

class GRAC():
	def __init__(
		self,
		state_dim,
		action_dim,
		max_action,
		batch_size=256,
		discount=0.99,
		max_timesteps=3e6,
                actor_lr = 3e-4,
                critic_lr = 3e-4,
                loss_decay = 0.95,
                log_freq = 200,
		cem_loss_coef =1.0,
		device=torch.device('cuda'),
	):
		self.action_dim = action_dim
		self.state_dim = state_dim
		self.max_action = max_action
		self.discount = discount
		self.total_it = 0

		self.device = device
		self.actor_lr = actor_lr # here is actor lr is not the real actor learning rate
		self.critic_lr = critic_lr
		self.loss_decay = loss_decay

		self.actor = Actor(state_dim, action_dim, max_action).to(device)
		self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.actor_lr)

		self.critic = Critic(state_dim, action_dim).to(device)
		self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.critic_lr)

		cem_sigma = 1e-2
		cem_clip = 0.5 * max_action
		self.cem_clip_init = cem_clip
		self.cem_clip = cem_clip
		self.searcher = Searcher(action_dim, max_action, device=device, sigma_init=cem_sigma, clip=cem_clip, batch_size=batch_size)
		self.action_dim = float(action_dim)
		self.log_freq = log_freq
		self.max_timesteps = max_timesteps
		self.cem_loss_coef = cem_loss_coef/float(self.action_dim)

	def select_action(self, state, writer=None, test=False):
		state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
		if test is False:
			with torch.no_grad():
				action = self.actor(state)
				better_action = action
			return better_action.cpu().data.numpy().flatten()
		else:
			_, _, action, _ = self.actor.forward_all(state)
			return action.cpu().data.numpy().flatten()

	def lr_scheduler(self, optimizer,lr):
		for param_group in optimizer.param_groups:
			param_group['lr'] = lr
		return optimizer

	def update_critic(self, critic_loss):
		self.critic_optimizer.zero_grad()
		critic_loss.backward()
		self.critic_optimizer.step()

	def save(self, filename):
		torch.save(self.critic.state_dict(), filename + "_critic")
		torch.save(self.critic_optimizer.state_dict(), filename + "_critic_optimizer")
		torch.save(self.actor.state_dict(), filename + "_actor")
		torch.save(self.actor_optimizer.state_dict(), filename + "_actor_optimizer")

	def load(self, filename):
		self.critic.load_state_dict(torch.load(filename + "_critic"))
		self.critic_optimizer.load_state_dict(torch.load(filename + "_critic_optimizer"))
		self.actor.load_state_dict(torch.load(filename + "_actor"))
		self.actor_optimizer.load_state_dict(torch.load(filename + "_actor_optimizer"))

	def train(self, replay_buffer, batch_size=100, writer=None, reward_range=20.0, reward_max=0, episode_step_max=100, reward_min=0, episode_step_min=1):
		self.total_it += 1
		log_it = (self.total_it % self.log_freq == 0)
		ratio_it = max(1.0 - self.total_it/float(self.max_timesteps), 0.1)
		if log_it:
			writer.add_scalar('train_critic/ratio_it', ratio_it, self.total_it)
		self.cem_clip = self.cem_clip_init * ratio_it
		# Sample replay buffer 
		state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)
		with torch.no_grad():
			Q_max = reward_max / (1 - self.discount) * (1 - self.discount ** int(episode_step_max))
			if reward_min >= 0:
				Q_min = reward_min / (1 - self.discount) * (1 - self.discount ** int(episode_step_min))
			else:
				Q_min = reward_min / (1 - self.discount) * (1 - self.discount ** int(episode_step_max))

			# Select action according to policy and add clipped noise
			#better_next_action = next_action1
			next_action1, next_action2, next_action_mean, next_action_sigma, prob1, prob2, probm= self.actor.forward_all_sample(next_state)
			coef1 = prob1 / (prob1 + prob2 + probm)
			coef2 = prob2 / (prob1 + prob2 + probm)
			coefm = probm / (prob1 + prob2 + probm)
			#print("next_action1",next_action1)
			#print("next_action2",next_action2)
			#print("prob1",prob1)
			#print("prob2",prob2)
			#next_action1, next_action2 = self.actor.forward_sample(next_state)
			#better_next_action,_ = self.searcher.search(next_state, next_action, self.critic.Q2, clip=self.cem_clip)
			#next_action2 = next_action_mean	
			#next_action1 = next_action
			#better_next_action = next_action1

			target_Q1 = self.critic(next_state, next_action1)
			target_Q2 = self.critic(next_state, next_action2)
			target_mean = self.critic(next_state, next_action_mean)
			target_Q1[target_Q1 > Q_max] = Q_max
			target_Q1[target_Q1 < Q_min] = Q_min
			target_Q2[target_Q2 > Q_max] = Q_max
			target_Q2[target_Q2 < Q_min] = Q_min
			target_mean[target_mean > Q_max] = Q_max
			target_mean[target_mean < Q_min] = Q_min

			#target_Q_w = target_Q1 * coef1 + target_Q2 * coef2 + target_mean * coefm
			#target_Q = torch.min(target_Q_w, target_mean)
			target_Q = torch.max(torch.min(target_Q1, target_mean), torch.min(target_Q2, target_mean))
			target_Q_mean_diff = target_mean - target_Q
			target_Q1_diff = target_Q1 - target_Q
			target_Q2_diff = target_Q2 - target_Q	
			#action_index = (target_Q1 > target_Q2).squeeze()
			#better_next_action[action_index] = next_action2[action_index]
			#better_target_Q = self.critic(next_state, better_next_action)
			#target_Q1 = target_Q

			target_Q[target_Q > Q_max] = Q_max
			target_Q[target_Q < Q_min] = Q_min
			
			target_Q_final = reward + not_done * self.discount * target_Q
			target_Q_final[target_Q_final > Q_max] = Q_max
			target_Q_final[target_Q_final < Q_min] = Q_min

			target_Q1 = target_Q
			next_action = next_action_mean
		# Get current Q estimates
		current_Q1 = self.critic(state, action)

		# Compute critic loss
		critic_loss = F.mse_loss(current_Q1, target_Q_final)
		self.update_critic(critic_loss)

		current_Q1_ = self.critic(state, action)
		target_Q1_ = self.critic(next_state, next_action)
		critic_loss2_1 = F.mse_loss(current_Q1_, target_Q_final)
		critic_loss2_2 = F.mse_loss(target_Q1_, target_Q1)
		weight1 = critic_loss2_1.item()
		weight2 = critic_loss2_2.item()	
		weight_loss = (math.sqrt(weight1) + 1.)/( math.sqrt(weight2) + 1.)
		critic_loss3 = critic_loss2_1 + critic_loss2_2 * weight_loss
		self.update_critic(critic_loss3)
		init_critic_loss3 = critic_loss3.clone()
		ratio = 0.0
		max_step = 0

		idi = 0
		while True:
			idi = idi + 1
			current_Q1_ = self.critic(state, action)
			target_Q1_ = self.critic(next_state, next_action)
			critic_loss3_1 = F.mse_loss(current_Q1_, target_Q_final)
			critic_loss3_2 = F.mse_loss(target_Q1_, target_Q1)
			critic_loss3 = critic_loss3_1 + critic_loss3_2 * weight_loss
			self.update_critic(critic_loss3)
			if critic_loss3_1 < critic_loss * self.loss_decay and critic_loss3_1 < critic_loss2_1 * self.loss_decay:# and torch.sqrt(critic_loss3_2) < torch.max(torch.mean(torch.abs(target_Q)) * 0.01, torch.mean(torch.abs(reward))):
                                        break
			if idi > 50:
				break
		#critic_loss = F.mse_loss(current_Q1, target_Q_final) + F.mse_loss(current_Q2, target_Q_final)
		weights_actor_lr = critic_loss2_1.item()
		if log_it:
			writer.add_scalar('train_critic/weight_loss', weight_loss, self.total_it)
			writer.add_scalar('train_critic/third_loss_num', idi, self.total_it)
			writer.add_scalar('train_critic/Q_max', Q_max, self.total_it)
			writer.add_scalar('train_critic/episode_step_max',episode_step_max, self.total_it)
			writer.add_scalar('train_critic/Q_min', Q_min, self.total_it)
			writer.add_scalar('train_critic/cem_clip', self.cem_clip, self.total_it)
			#writer.add_scalar('train_critic/Q_min_mean', torch.mean(Q_min), self.total_it)
			#writer.add_scalar('train_critic/Q_min_min', torch.min(Q_min), self.total_it)
			writer.add_scalar('train_critic/episode_step_min',episode_step_min, self.total_it)
		if log_it:
			writer.add_scalar('train_loss/loss2_1',critic_loss2_1,self.total_it)
			writer.add_scalar('train_loss/loss2_2',critic_loss2_2,self.total_it)
			writer.add_scalar('train_loss/loss3_1_r',critic_loss3_1/critic_loss2_1,self.total_it)
			writer.add_scalar('train_loss/loss3_2_r',critic_loss3_2/critic_loss2_2,self.total_it)
			writer.add_scalar('train_loss/loss3_1_r_loss',critic_loss3_1/critic_loss,self.total_it)
			writer.add_scalar('train_loss/sqrt_critic_loss3_2',torch.sqrt(critic_loss3_2),self.total_it)
			writer.add_scalar('train_loss/max_reward',reward_max,self.total_it)	
			#writer.add_scalar('train_loss/min_reward',reward_min,self.total_it)
		if self.total_it % 1 == 0:
			#lr_tmp = self.actor_lr / (float(weights_actor_lr)+1.0) * self.actor_lr_ratio
			#self.actor_optimizer = self.lr_scheduler(self.actor_optimizer, lr_tmp)

			# Compute actor loss
			actor_action, log_prob, action_mean, action_sigma = self.actor.forward_all(state)
			q_actor_action = self.critic.Q1(state, actor_action)
			m = Normal(action_mean, action_sigma)

			better_action,_ = self.searcher.search(state, actor_action, self.critic.Q1, batch_size=batch_size, clip=self.cem_clip)#####
			q_better_action = self.critic.Q1(state, better_action)
			
			q_better_action[q_better_action > Q_max] = Q_max
			q_better_action[q_better_action < Q_min] = Q_min
			q_actor_action[q_actor_action > Q_max] = Q_max
			q_actor_action[q_actor_action < Q_min] = Q_min

			log_prob_better_action = m.log_prob(better_action).sum(1,keepdim=True)

			adv = (q_better_action - q_actor_action).detach()
			adv = torch.min(adv,torch.ones_like(adv) * Q_max * 0.05)
			cem_loss = log_prob_better_action * adv#torch.min(reward_range * torch.ones_like(adv) * ratio_it, adv)
			actor_loss = -(cem_loss * self.cem_loss_coef + q_actor_action).mean()

			# Optimize the actor 
			Q_before_update = self.critic.Q1(state, actor_action)

			self.actor_optimizer.zero_grad()
			actor_loss.backward()
			self.actor_optimizer.step()
			#self.actor_optimizer = self.lr_scheduler(self.actor_optimizer, self.actor_lr)

		if log_it:
			with torch.no_grad():
				writer.add_scalar('train_critic/third_loss_num', idi, self.total_it)
				writer.add_scalar('train_critic/critic_loss', critic_loss, self.total_it)
				writer.add_scalar('train_critic/critic_loss3', critic_loss3, self.total_it)
	
				#target_Q1
				writer.add_scalar('train_critic/target_Q1/mean', torch.mean(target_Q1), self.total_it)
				writer.add_scalar('train_critic/target_Q1/max', target_Q1.max(), self.total_it)
				writer.add_scalar('train_critic/target_Q1/min', target_Q1.min(), self.total_it)
				writer.add_scalar('train_critic/target_Q1/std', torch.std(target_Q1), self.total_it)
	
				#current_Q1
				writer.add_scalar('train_critic/current_Q1/mean', current_Q1.mean(), self.total_it)
				writer.add_scalar('train_critic/current_Q1/std', torch.std(current_Q1), self.total_it)
				writer.add_scalar('train_critic/current_Q1/max', current_Q1.max(), self.total_it)
				writer.add_scalar('train_critic/current_Q1/min', current_Q1.min(), self.total_it)

				# advantage 
				writer.add_scalar('train_critic/adv/mean', adv.mean(), self.total_it)
				writer.add_scalar('train_critic/adv/std', torch.std(adv), self.total_it)
				writer.add_scalar('train_critic/adv/max', adv.max(), self.total_it)
				writer.add_scalar('train_critic/adv/min', adv.min(), self.total_it)

				# targetQ1_diff
				writer.add_scalar('train_critic/target_Q1_diff/mean', target_Q1_diff.mean(), self.total_it)
				writer.add_scalar('train_critic/target_Q1_diff/std', torch.std(target_Q1_diff), self.total_it)
				writer.add_scalar('train_critic/target_Q1_diff/max', target_Q1_diff.max(), self.total_it)
				writer.add_scalar('train_critic/target_Q1_diff/min', target_Q1_diff.min(), self.total_it)

				# targetQ2_diff
				writer.add_scalar('train_critic/target_Q2_diff/mean', target_Q2_diff.mean(), self.total_it)
				writer.add_scalar('train_critic/target_Q2_diff/std', torch.std(target_Q2_diff), self.total_it)
				writer.add_scalar('train_critic/target_Q2_diff/max', target_Q2_diff.max(), self.total_it)
				writer.add_scalar('train_critic/target_Q2_diff/min', target_Q2_diff.min(), self.total_it)
	
				
				# target_Q_mean_diff
				writer.add_scalar('train_critic/target_Q_mean_diff/mean', target_Q_mean_diff.mean(), self.total_it)
				writer.add_scalar('train_critic/target_Q_mean_diff/std', torch.std(target_Q_mean_diff), self.total_it)
				writer.add_scalar('train_critic/target_Q_mean_diff/max', target_Q_mean_diff.max(), self.total_it)
				writer.add_scalar('train_critic/target_Q_mean_diff/min', target_Q_mean_diff.min(), self.total_it)


	
	def save(self, filename):
		torch.save(self.critic.state_dict(), filename + "_critic")
		torch.save(self.critic_optimizer.state_dict(), filename + "_critic_optimizer")
		torch.save(self.actor.state_dict(), filename + "_actor")
		torch.save(self.actor_optimizer.state_dict(), filename + "_actor_optimizer")

	def load(self, filename):
		self.critic.load_state_dict(torch.load(filename + "_critic"))
		self.critic_optimizer.load_state_dict(torch.load(filename + "_critic_optimizer"))
		self.actor.load_state_dict(torch.load(filename + "_actor"))
		self.actor_optimizer.load_state_dict(torch.load(filename + "_actor_optimizer"))
