# Celtic GAN Generation
# Inspired from Demystifying Generative Adversarial Networks, Towards Data Science Blog, Stefan Hosein, 10 May 2018
# https://towardsdatascience.com/demystifying-generative-adversarial-networks-c076d8db8f44
# Jean-Pierre Briot, 05/11/2019

import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import glob
import random

from keras.layers import Input
from keras.models import Model, Sequential
from keras.layers.core import Dense, Dropout, Activation
from keras.layers.advanced_activations import LeakyReLU
from keras.datasets import mnist
from keras.optimizers import Adam
from keras import initializers
from keras.utils import np_utils

from representation import *
from metrics import *

config.deep_music_analysis_verbose = False	# analysis verbose flag - default value = False
config.deep_music_training_verbose = False	# training verbose flag - default value = False
config.deep_music_generate_verbose = True	# generation verbose flag - default value = False

corpus_names_list = []

corpus_files_names_list = glob.glob('Data/ABC_Session/*.abc')

print(str(len(corpus_files_names_list)) + ' songs considered')

print('Song names: ' + str(corpus_files_names_list))

def load_corpus_from_files(corpus_files_names_list):
	corpus_list = []
	for n, file_name in enumerate(corpus_files_names_list):
		corpus_list.append(converter.parse(file_name))
	return(corpus_list)

is_transposed = False

print('Loading and parsing the corpus')

corpus_list = load_corpus_from_files(corpus_files_names_list)

if is_transposed:
	print('Transpose/Extend the corpus into all keys')
	corpus_list = transpose_corpus_into_all_keys(corpus_list)

# corpus_list = corpus_list[:20]

print('Analyze the corpus')

analyze_corpus(corpus_list)

# bug !!
# Quarter length smallest note duration: 0.0 !!
# config.time_step_quarter_length == 0 !!

config.time_step_quarter_length = 0.25

print('Durations (in quarter length): ' + str(music_duration_quarter_length_v))

max_number_quarters = min(music_duration_quarter_length_v)

# max_number_quarters = 4

config.max_number_time_steps = int(max_number_quarters / config.time_step_quarter_length)

print('Max number quarters (in quarter length): ' + str(max_number_quarters))
print('Max number time steps: ' + str(config.max_number_time_steps))

print('Encode the corpus (Irish music)')

corpus_data = encode_corpus(corpus_list)
# structure: hierarchy: music/part/encoded_data

print('Construction of the training and validation (test) datasets')

X_train = []

for i in range(len(corpus_data)):
	X_train.append(corpus_data[i][0])

# X_train : [[0, 0, 0 ..  0, (:) ... (:) 0, 0, 0 .. 0], ... [0, 0, 0 ..  0, (:) ... (:) 0, 0, 0 .. 0]]
#                  chorale 1				       chorale P
#		   time step 1  		 time step N     time step 1		      time step N

# numpy representations

X_train = np.array(X_train)

print('X_train shape: ' + str(X_train.shape))		# (64, 3520)

# Define the deep learning model/architecture

output_size = config.max_number_time_steps * config.one_hot_size_v[0]

print('Output size: ' + str(output_size))

# To make sure that we can reproduce the experiment and get the same results
np.random.seed(10)

# Random noise vector dimension
random_dim=120

epochs = 100
batch_size = 20

optimizer=Adam(lr=0.0002, beta_1=0.5)

# Generator

Generator = Sequential()
Generator.add(Dense(256,
	input_dim=random_dim,
	kernel_initializer=initializers.RandomNormal(stddev=0.02)))
Generator.add(LeakyReLU(0.2))
Generator.add(Dense(512))
Generator.add(LeakyReLU(0.2))
Generator.add(Dense(1024))
# Generator.add(LeakyReLU(0.2))
Generator.add(Dense(output_size,
	activation='sigmoid'))

Generator.compile(loss='binary_crossentropy',
	optimizer=optimizer)

# Discriminator

Discriminator = Sequential()
Discriminator.add(Dense(1024,
	input_dim=output_size,
	kernel_initializer=initializers.RandomNormal(stddev=0.02)))
Discriminator.add(LeakyReLU(0.2))
Discriminator.add(Dropout(0.3))
Discriminator.add(Dense(512))
Discriminator.add(LeakyReLU(0.2))
Discriminator.add(Dropout(0.3))
Discriminator.add(Dense(256))
Discriminator.add(LeakyReLU(0.2))
Discriminator.add(Dropout(0.3))
Discriminator.add(Dense(1,
	activation='sigmoid'))

Discriminator.compile(loss='binary_crossentropy',
	optimizer=optimizer)

# GAN

# def get_gan_network(discriminator, random_dim, generator, optimizer):

# We initially set trainable to False since we only want to train either the
# generator or discriminator at a time
Discriminator.trainable = False

GAN_input = Input(shape=(random_dim,))

GAN_sample = Generator(GAN_input)

GAN_output = Discriminator(GAN_sample)

GAN = Model(inputs=GAN_input,
		outputs=GAN_output)

GAN.compile(loss='binary_crossentropy',
	optimizer=optimizer)

def train(epochs, batch_size):
# Split the training data into batches of size batch_size
# X_train.shape[0] = number of examples
	batch_count = X_train.shape[0] / batch_size
	for e in range(1, epochs+1):
		print('-'*15, 'Epoch %d' % e, '-'*15)
		for _ in range(int(batch_count)):
			# Get a random set of input noise and music
			noises = np.random.normal(0, 1, size=[batch_size, random_dim])
			music_batch = X_train[np.random.randint(0, X_train.shape[0], size=batch_size)]
			# Generate fake music
			generated_musics = Generator.predict(noises)
			# Training dataset = set of real music samples followed by set of faked samples
			X = np.concatenate([music_batch, generated_musics])
			# Labels for generated and real data
			y_dis = np.zeros(2*batch_size)
			# One-sided label smoothing
			# For real samples, label = 0.9 and for faked samples, label = 0
			y_dis[:batch_size] = 0.9
			# Train Discriminator
			Discriminator.trainable = True
			Discriminator.train_on_batch(X, y_dis)
			# Train Generator
			noises = np.random.normal(0, 1, size=[batch_size, random_dim])
			y_gen = np.ones(batch_size)
			Discriminator.trainable = False
			# Train Generator to be able to fool Discriminator (GAN/Discriminator(noises) labels=1)
			GAN.train_on_batch(noises, y_gen)

def create_labels(n):
	return(np.random.normal(size=[n, random_dim]))

def create_melodies_from_labels(labels_list):
	data_list = Generator.predict(np.array(labels_list))
	data_list2 = []
	for i in range(len(data_list)):			# transforms	[[0, 0 .. 0 (:) ... (:) 0, 0 .. 0], ... [0, 0 .. 0 (:) ... (:) 0, 0 .. 0]]
		data_list2.append([data_list[i]])		# into		[[[0, 0 .. 0 (:) ... (:) 0, 0 .. 0]], ... [[0, 0 .. 0 (:) ... (:) 0, 0 .. 0]]]
	return(data_list2)

def create_melodies(n):
	return(create_melodies_from_labels(create_labels(n)))

generator = train(epochs, batch_size)
		
melodies_list = create_melodies(5)
scores = create_scores(1, melodies_list)

for i in range(len(scores)):
		scores[i].write('midi', 'mid/gan_mel_irish_' + str(i) + '.mid')

# Checking the mode collapse

def check_samples(n):
	argmax_list_list = []
	original_sample = True
	number_samples = 0
	for i in range(n):
		labels_list = np.random.normal(size=[1, random_dim])
		melody_data = create_melodies_from_labels(labels_list)[0]
		argmax_list = extract_argmax_list(melody_data[0])
		for j in range(len(argmax_list_list)):
			original_sample = True
			if argmax_list_list[j] == argmax_list:
				original_sample = False
				break
		if original_sample:
			argmax_list_list.append(argmax_list)
			number_samples = number_samples + 1
	return(number_samples)

one_hot_size = config.one_hot_size_v[0]

def extract_argmax_list(data):
	argmax_list = []
	for i in range(int(len(data) / one_hot_size)):
		argmax_item = np.argmax(data[(i * one_hot_size):(i * one_hot_size) + one_hot_size])
		argmax_list.append(argmax_item)
	return(argmax_list)

print(check_samples(100))

# Corpus[:20][:4 quarters] -> 6
# Corpus[:20] -> 28

# Full Corpus[:4 quarters] -> 1
# Full Corpus -> 16

# Full Corpus Transposed[:4 quarters] -> 3
# Full Corpus Transposed -> 93