import torch
import random
import numpy as np
from utils.network import PTQNetwork, PTACAgent, LEARN_RATE, EPS_MIN, REPLAY_BATCH_SIZE

class DDQNetwork(PTQNetwork):
	def __init__(self, state_size, action_size, lr=LEARN_RATE, gpu=True, load=None): 
		super().__init__(state_size, action_size, lr=lr, gpu=gpu, load=load, name="ddqn")

	def get_action(self, state, use_target=False, numpy=True, sample=True):
		with torch.no_grad():
			q_values = self.critic_local(state) if not use_target else self.critic_target(state)
			return q_values.softmax(-1).cpu().numpy() if numpy else q_values.softmax(-1)

	def get_q_value(self, state, action, use_target=False, numpy=True):
		with torch.no_grad():
			q_values = self.critic_local(state) if not use_target else self.critic_target(state)
			out_dims = q_values.size()[:-1]
			q_values = q_values.reshape(-1, q_values.size(-1))
			q_indices = action.argmax(-1).reshape(-1)
			q_selected = q_values[np.arange(q_indices.size(0)), q_indices].reshape(*out_dims, 1)
			return q_selected.cpu().numpy() if numpy else q_selected
	
	def optimize(self, states, actions, q_targets, importances=1):
		q_values = self.critic_local(states)[np.arange(actions.size(0)), actions.argmax(-1)].unsqueeze(-1)
		critic_error = q_values - q_targets.detach()
		critic_loss = importances.to(self.device) * critic_error.pow(2)
		self.step(self.critic_optimizer, critic_loss.mean())
		self.soft_copy(self.critic_local, self.critic_target)
		return critic_error.cpu().detach().numpy().squeeze(-1)

class DDQNAgent(PTACAgent):
	def __init__(self, state_size, action_size, lr=LEARN_RATE, gpu=True, load=None):
		super().__init__(state_size, action_size, DDQNetwork, lr=lr, gpu=gpu, load=load)

	def get_action(self, state, eps=None, sample=True, e_greedy=True):
		eps = self.eps if eps is None else eps
		action_random = super().get_action(state, eps)
		if e_greedy and random.random() < eps: return action_random
		action_greedy = self.network.get_action(self.to_tensor(state), sample=sample)
		return action_greedy
		
	def train(self, state, action, next_state, reward, done):
		self.buffer.append((state, action, reward, done))
		if np.any(done[0]) or len(self.buffer) >= self.update_freq:
			states, actions, rewards, dones = map(self.to_tensor, zip(*self.buffer))
			self.buffer.clear()	
			states = torch.cat([states, self.to_tensor(next_state).unsqueeze(0)], dim=0)
			actions = torch.cat([actions, self.network.get_action(states[-1], use_target=True, numpy=False).unsqueeze(0)], dim=0)
			values = self.network.get_q_value(states, actions, use_target=True, numpy=False)
			targets = self.compute_gae(values[-1], rewards.unsqueeze(-1), dones.unsqueeze(-1), values[:-1])[0]
			states, actions, targets = [x.view(x.size(0)*x.size(1), *x.size()[2:]).cpu().numpy() for x in (states[:-1], actions[:-1], targets)]
			self.replay_buffer.extend(list(zip(states, actions, targets)), shuffle=True)	
		if len(self.replay_buffer) > REPLAY_BATCH_SIZE:
			(states, actions, targets), indices, importances = self.replay_buffer.sample(REPLAY_BATCH_SIZE, dtype=self.to_tensor)
			errors = self.network.optimize(states, actions, targets, importances**(1-self.eps))
			self.replay_buffer.update_priorities(indices, errors)
			if np.any(done[0]): self.eps = max(self.eps * self.decay, EPS_MIN)
