import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import os
from network.wavenet import WaveNet
from network.disnet import Discriminator
import time
import datasets
import config
import wave_gen


############arguments assighment: TODO use config file#########################
in_out_size = 256
res_size=512
layer_size = 10
stack_size = 2#args.stack_size
rfield_size = int(np.sum([2 ** i for i in range(0, layer_size)] * stack_size))

data_dir_g1 = '/home/jn1664/DGM/DiscoGAN/data_piano'
data_dir_g2 = '/home/jn1664/DGM/DiscoGAN/data_glock'
seed_dir_g1 = data_dir_g1 + '/A_Grand.wav'
seed_dir_g2 = data_dir_g2 + '/A_Glockenspiel.wav'
model_file_path = 'model/'
source_name = 'piano_grand'
target_name = 'glockenspiel'

sample_size = 10000
train_steps = 1000000
##################################################################################


train_loader = datasets.DataLoader(data_dir_g1, data_dir_g2, rfield_size, sample_size)

g1 = WaveNet(in_out_size, res_size, stack_size, layer_size)#P->G
#g2 = WaveNet(in_out_size, res_size, stack_size, layer_size)#G->P
if torch.cuda.is_available():
    g1.cuda()
    #g2.cuda()
CEloss = nn.CrossEntropyLoss()


d1 = Discriminator(in_out_size, 1)#G 
d2 = Discriminator(in_out_size, 1)#P
if torch.cuda.is_available():
    d1.cuda()
    d2.cuda()

criterion = nn.BCELoss()
lr_g = 0.002
lr_d = 0.01
g_optimizer = torch.optim.Adam(g1.parameters(), lr=lr_g)
d1_optimizer = torch.optim.SGD(d1.parameters(), lr=lr_d)
d2_optimizer = torch.optim.SGD(d2.parameters(), lr=lr_d)


def infinite_batch(train_loader):
    while True:
        for dataset in train_loader:
            for inputs1, targets1, inputs2, targets2 in dataset:
                yield inputs1, targets1, inputs2, targets2


def train(epoch):
    #if os.path.exists(model_file_path+'model_1100_G.pth'):
	#g.load_state_dict(torch.load(model_file_path+'model_1100_G.pth')) 
    #if os.path.exists(model_file_path+'model_1100_D.pth'):
        #d.load_state_dict(torch.load(model_file_path+'model_1100_D.pth'))


    g1.train()
    #g2.train()
    d1.train()
    d2.train()


    total_steps = 0
    for data_g1, target_g1, data_g2, target_g2 in infinite_batch(train_loader):
	total_steps += 1

	num_target = target_g1.shape[1]


	#prepare-1. seed data for generators
	seed_g1, l1 = wave_gen.get_seed_from_audio(seed_dir_g1, rfield_size, sample_size)
	seed_g2, l2 = wave_gen.get_seed_from_audio(seed_dir_g2, rfield_size, sample_size)
	#prepare-2. generate fake audio
	fake_audio1 = g1(seed_g2)#(1,7954,256)
	fake_audio2 = g1(seed_g1)#(1,7954,256)
	#prepare-3. real label
	real_label = Variable(torch.ones(num_target).cuda())#(7954i,)
	#prepare-4. fake label
	fake_label = Variable(torch.zeros(num_target).cuda())#(7954,)


        #D1-1. discriminate for real audio
        d1_real_output = d1(target_g1).view(-1)#3D(1,7954,256)
        #D1-2. compute real loss
        d1_real_loss = criterion(d1_real_output, real_label)
        #D1-3. discriminate for fake audio generated by G
        d1_fake_output = d1(fake_audio1).view(-1)[:fake_label.shape[0]]#3D(1,7954,256)->(1,1,7954)
        #D1-4. compute fake loss
        d1_fake_loss = criterion(d1_fake_output, fake_label)
	#D1-5. compute total loss and optimize
	d1_loss = d1_real_loss + d1_fake_loss
	d1.zero_grad()
        d1_loss.backward(retain_graph=True)
        d1_optimizer.step()


	#D2-1. discriminate for real audio
        d2_real_output = d2(target_g2).view(-1)
        #D2-2. compute real loss
        d2_real_loss = criterion(d2_real_output, real_label)
        #D2-3. discriminate for fake audio generated by G
        d2_fake_output = d2(fake_audio2).view(-1)[:fake_label.shape[0]]
        #D2-4. compute fake loss
        d2_fake_loss = criterion(d2_fake_output, fake_label)
        #D2-5. compute total loss and optimize
        d2_loss = d2_real_loss + d2_fake_loss
        d2.zero_grad()
        d2_loss.backward(retain_graph=True)
        d2_optimizer.step()


	#G-1. compute GAN loss
        g1_gan_loss = criterion(d1_fake_output, real_label)
	g2_gan_loss = criterion(d2_fake_output, real_label)
	#G-2. compute reconstruction loss
        recons_audio1 = g1(fake_audio2)#(1,5908, 256)
        recons_audio1 = recons_audio1.view(-1, in_out_size)#(5908,256)
        target_g1 = datasets.one_hot_decode(target_g1.data, 2)[:,rfield_size:]
        target_g1 = Variable(target_g1.cuda()).view(-1)
	g1_reconst_loss = CEloss(recons_audio1[:target_g1.shape[0]], target_g1)

        recons_audio2 = g1(fake_audio1)#(1,5908, 256)
        recons_audio2 = recons_audio2.view(-1, in_out_size)#(5908,256)
        target_g2 = datasets.one_hot_decode(target_g2.data, 2)[:,rfield_size:]
        target_g2 = Variable(target_g2.cuda()).view(-1)
	g2_reconst_loss = CEloss(recons_audio2[:target_g2.shape[0]], target_g2)
	#G-3. compute total loss and optimize
	g_loss = g1_gan_loss + g2_gan_loss + g1_reconst_loss + g2_reconst_loss
	g1.zero_grad()
        g_loss.backward()
        g_optimizer.step()


        if total_steps % 10 == 0:
            print('Train steps: ['+source_name+'] [{0}]\tG Loss: {1}\tD1 Loss: {2}\tD2 Loss{3}'.format(total_steps, g_loss.data[0], d1_loss.data[0], d2_loss.data[0]))

	if (total_steps % 100 == 0) or (total_steps == train_steps):
	    model_file = model_file_path+'model_' + str(total_steps)
	    torch.save(g1.state_dict(), model_file + '_G.pth') 
	    torch.save(d1.state_dict(), model_file + '_D1.pth')
	    torch.save(d2.state_dict(), model_file + '_D2.pth')

	if total_steps> train_steps:
	    break



num_epochs = 1#args.epochs
for epoch in range(num_epochs):
    start_t = time.time()
    train(epoch)
    end_D = time.time()
    print('''[Total time(D)]: %.2fsec'''% (end_D - start_t))
