from torch.utils.data import DataLoader
from torch.optim import Adam
from dataset.e_piano import create_epiano_datasets, compute_epiano_accuracy
from utilities.argument_funcs import parse_train_args, print_train_args, write_model_params, parse_eval_args, \
	parse_generate_args
import torch.nn as nn
from model.Generator import Generator
from model.Discriminator import Discriminator
from utilities.constants import *
from utilities.device import get_device
import torch.nn.functional as F
import random


class GANInstructor:
	def __init__(self, dis_path = None):
		self.train_args = parse_train_args()
		self.eval_args = parse_eval_args()
		self.parse_args = parse_generate_args()

		self.dis_path = dis_path
		self.gen = Generator(n_layers=self.train_args.n_layers, num_heads=self.train_args.num_heads,
							 d_model=self.train_args.d_model, dim_feedforward=self.train_args.dim_feedforward,
							 max_sequence=self.train_args.max_sequence, rpr=self.eval_args.rpr).to(get_device())
		self.dis = Discriminator().to(get_device())
		self.dis_pretrain_criterion = nn.BCELoss().to(get_device())

		self.loadGenerator()
		self.loadDiscriminator()
		self.gen_opt = Adam(self.gen.parameters(), lr=0.00001, betas=(ADAM_BETA_1, ADAM_BETA_2), eps=ADAM_EPSILON)
		self.dis_opt = Adam(self.dis.parameters(), lr=0.00001, betas=(ADAM_BETA_1, ADAM_BETA_2), eps=ADAM_EPSILON)

		self.train_dataset, self.val_dataset, self.test_dataset = create_epiano_datasets(INPUT_DIR, max_seq=256)
		self.train_loader = DataLoader(self.train_dataset, batch_size=self.train_args.batch_size, num_workers=self.train_args.n_workers, shuffle=True, drop_last=True)
		self.val_loader = DataLoader(self.val_dataset, batch_size=self.train_args.batch_size, num_workers=self.train_args.n_workers)
		self.test_loader = DataLoader(self.test_dataset, batch_size=self.train_args.batch_size, num_workers=self.train_args.n_workers)

		#TODO define the results path

	def loadGenerator(self):
		self.gen.load_state_dict(torch.load(self.eval_args.model_weights))

	def loadDiscriminator(self):
		pass
		#self.dis.load_state_dict(torch.load("./new_discriminator.pth"))

	def run(self):
		if (self.dis_path == None): #Model not pre-trained
			print("Pre-Training Discriminator")


	def preTrainDiscriminator(self):
		self.dis.train()
		self.gen.eval()

		for epoch in range (0,self.train_args.pretrain_epochs):
			true_loss = []
			false_loss = []
			for batch_num, batch in enumerate(self.train_loader):
				if (batch_num%100 == 0) : print("Iterating", batch_num)
				self.dis_opt.zero_grad()
				x = batch[0].to(get_device())
				tgt = batch[1].to(get_device())
				ran = random.randint(0,1)
				if (389 in x or 389 in tgt) : print(389)
				else:
					if (ran == 0) :
						x_one = F.one_hot(x, 388)
						generated = self.gen.sample(x_one, self.parse_args.target_seq_length, one_hot=True)
						generated = torch.cat((x_one, generated), 1)
						generated = generated.argmax(dim=2)
						prediction_dis_fake = self.dis(generated)
						prediction_dis_fake = torch.squeeze(prediction_dis_fake, 0)
						label_0 = torch.zeros([self.train_args.batch_size,1], dtype=torch.float32).to(get_device())
						loss_dis_fake = self.dis_pretrain_criterion(prediction_dis_fake, label_0)
						loss_dis_fake.backward()
						self.dis_opt.step()
						false_loss.append(loss_dis_fake.detach().cpu().numpy().item())

					else:
						tgt = torch.cat((x, tgt), 1)
						prediction_dis_true = self.dis(tgt)
						prediction_dis_true = torch.squeeze(prediction_dis_true, 0)
						label_1 = torch.ones([self.train_args.batch_size,1], dtype=torch.float32).to(get_device())
						loss_dis_true = self.dis_pretrain_criterion(prediction_dis_true, label_1)
						loss_dis_true.backward()
						self.dis_opt.step()
						true_loss.append(loss_dis_true.detach().cpu().numpy().item())

				try:
					if (batch_num % 100 == 0): print("false loss:", sum(false_loss) / len(false_loss))
					if (batch_num % 100 == 0): print("true loss:", sum(true_loss) / len(true_loss))
				except: pass


			torch.save(self.dis.state_dict(), "./new_discriminator.pth")

			with torch.no_grad():

				true_loss = []
				false_loss = []
				total_loss = []
				acc = []

				for batch_num, batch in enumerate(self.test_loader):

					try:
						x = batch[0].to(get_device())
						tgt = batch[1].to(get_device())

						if (x.size(dim=0) == self.train_args.batch_size):
							generated = self.gen.sample(x[0], self.parse_args.target_seq_length)
							#generated = torch.randint(0, 388, (2, 256)).to(get_device())
							generated = torch.cat((x,generated), 1)
							prediction_dis_fake = self.dis(generated)
							prediction_dis_fake = torch.squeeze(prediction_dis_fake, 0)
							label_0 = torch.zeros([1], dtype=torch.float32).to(get_device())
							loss_dis_fake = self.dis_pretrain_criterion(prediction_dis_fake, label_0)
							false_loss.append(loss_dis_fake.detach().cpu().numpy().item())

							tgt = torch.cat((x,tgt), 1)
							prediction_dis_true = self.dis(tgt)
							prediction_dis_true = torch.squeeze(prediction_dis_true, 0)
							label_1 = torch.ones([1], dtype=torch.float32).to(get_device())
							loss_dis_true = self.dis_pretrain_criterion(prediction_dis_true, label_1)
							true_loss.append(loss_dis_true.detach().cpu().numpy().item())

							total_loss.append((sum(true_loss)+sum(false_loss)) / (2 * len(true_loss)))
							print("True Loss : ", true_loss[-1])
							print("False Loss : ", false_loss[-1])
							print("Overall Loss : ", total_loss[-1])
					except:
						pass

	def iterateOnGenerator(self):
		self.dis.eval()
		self.gen.train()

		for epoch in range(0, self.train_args.gen_epochs):
			loss = []
			for batch_num, batch in enumerate(self.train_loader):
				if (batch_num % 100 == 0): print("Iterating", batch_num)
				self.gen_opt.zero_grad()
				x = batch[0].to(get_device())
				generated = self.gen.sample(x, self.parse_args.target_seq_length)
				generated = torch.cat((x, generated), 1)
				prediction_dis_fake = self.dis(generated)
				prediction_dis_fake = torch.squeeze(prediction_dis_fake, 0)
				label_1 = torch.ones([self.train_args.batch_size, 1], dtype=torch.float32).to(get_device())
				loss_dis_fake = self.dis_pretrain_criterion(prediction_dis_fake, label_1)
				loss_dis_fake.backward()
				self.dis_opt.step()
				loss.append(loss_dis_fake.detach().cpu().numpy().item())
				if (batch_num % 100 == 0): print("loss:", sum(loss) / len(loss))
