import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os

class Linear_QNet(nn.Module):
	def __init__(self, input_size, hidden_size, hidden_size2, output_size):
		# feed forward neural network
		super().__init__()
		self.linear1 = nn.Linear(input_size, hidden_size)
		self.linear2 = nn.Linear(input_size, hidden_size2)
		self.linear3 = nn.Linear(hidden_size2, output_size)

	def forward(self, x):
		x = F.relu(self.linear1(x))
		x = self.linear2(x)
		x = self.linear3(x)
		return x

	def save(self, state_dict, file_name="model.pth"):
		model_folder_path = "./models"
		if not os.path.exists(model_folder_path):
			os.makedirs(model_folder_path)

		file_name = os.path.join(model_folder_path,file_name)
		torch.save(state_dict, file_name)

	def load(self, file_name="model.pth"):
		model_folder_path = "./models"
		file_name = os.path.join(model_folder_path,file_name)
		checkpoint = torch.load(file_name)
		self.load_state_dict(checkpoint['state_dict'])

class QTrainer():

	def __init__(self,model, lr, gamma):
		self.lr = lr
		self.gamma = gamma
		self.model = model
		self.optimizer = optim.Adam(model.parameters(), lr=self.lr)
		self.criterion = nn.MSELoss()

	def train_step(self, state, action, reward, next_state, done):
		state = torch.tensor(state,dtype=torch.float)
		next_state = torch.tensor(next_state,dtype=torch.float)
		action = torch.tensor(action,dtype=torch.long)
		reward = torch.tensor(reward,dtype=torch.float)

		if len(state.shape) == 1:
			state = torch.unsqueeze(state, 0)
			next_state = torch.unsqueeze(next_state, 0)
			action = torch.unsqueeze(action, 0)
			reward = torch.unsqueeze(reward, 0)
			done = (done, )

		# 1. predicted Q values with current state
		pred = self.model(state)

		# 2. Q_new = r + gamma * max(next_predicted_q_value) -> only if not done
		# pred.clone()
		# preds[argmax(action)] = Q_new
		target = pred.clone()
		for idx in range(len(done)):
			Q_new = reward[idx]
			if not done[idx]:
				Q_new = reward[idx] + self.gamma * torch.max(self.model(next_state[idx]))
			target[idx][torch.argmax(action[idx]).item()] = Q_new
	
		self.optimizer.zero_grad()
		loss = self.criterion(target,pred)
		loss.backward()

		self.optimizer.step()



torch.nn.Linear torch.nn.MSELoss torch.save torch.unsqueeze torch.tensor torch.load torch.argmax