# Irish Music Disentangling Variational Autoencoder generation example
# Implementation inspired from EC2-VAE Architecture
# DEEP MUSIC ANALOGY VIA LATENT REPRESENTATION DISENTANGLEMENT
# by Ruihan Yang, Dingsu Wang, ZiyuWang, Tianyao Chen, Junyan Jiang and Gus Xia
# arXiv:1906.03626v4 [cs.SD] 20 Oct 2019
# Pytorch

# Jean-Pierre Briot
# 26/11/2019

import glob
import random

import torch
import torch.nn as nn
import torch.optim as optim
import torch.random

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 = False	# generation verbose flag - default value = False

number_epochs = 100

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)

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 = 8

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 melodies)')

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

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

X_train = []
y_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

# melody soprano (unique) voice representation one hot size
one_hot_size = config.one_hot_size_v[0]

# rhythm representation one hot size: only 3 items for the one hot: hold, rest (the two special elements), note (independent of its actual pitch)
r_one_hot_size = number_special_elements_one_hot + 1

# construct associated rhythm representation examples - only 3 items one hot size: hold, rest, note
X_r_train = []

for i in range(len(corpus_data)):
	soprano_data_i = corpus_data[i][0]
	X_r_train_i = []
	for j in range(int(len(soprano_data_i) / one_hot_size)):
		X_r_train_i.extend(soprano_data_i[j * one_hot_size:(j * one_hot_size) + r_one_hot_size - 1])
		if ((soprano_data_i[j * one_hot_size] == 0) and (soprano_data_i[(j * one_hot_size) + 1] == 0)):
			# if no hold and no rest, note
			note = 1.
		else:
			# otherwise, no note
			note = 0.
		X_r_train_i.append(note)
	X_r_train.append(X_r_train_i)

print('X_train shape: ' + str(np.array(X_train).shape))
print('X_r_train shape: ' + str(np.array(X_r_train).shape))

# torch tensors representations
X_train = torch.tensor(X_train)
X_r_train = torch.tensor(X_r_train)

# Define the deep learning model/architecture
in_dim = config.max_number_time_steps * one_hot_size
r_out_dim = config.max_number_time_steps * r_one_hot_size

print('Input size: ' + str(in_dim))					# 3968
print('Rhythm Output size: ' + str(r_out_dim))		# 384

# Define the deep learning model/architecture
i_dim = 1000		# intermediate layer dimension
z_dim = 500		# latent variables (hidden layer) dimension (including rhythm latent variables)
z_r_dim = 300		# rhythm latent variables (hidden layer) dimension

# Extract the rhythm latent variables component
def z_r(z):
	return(z[:,:z_r_dim])

# Extract the rest latent variables component
def z_rest(z):
	return(z[:,z_r_dim:])

print('Hidden layer dimension: ' + str(z_dim))
print('Hidden layer rhythm label dimension: ' + str(z_r_dim))

print('Define and create the autoencoder model/network')

# Model of the Variational Autoencoder (VAE)
class VAE(nn.Module):
	def __init__(self, in_dim, i_dim, z_dim, z_r_dim, r_out_dim):
		super(VAE, self).__init__()

		self.encoder_pre = nn.Linear(in_dim, i_dim)
	
		self.encoder_mean = nn.Linear(i_dim, z_dim)
		self.encoder_var = nn.Linear(i_dim, z_dim)
	
		self.r_decoder = nn.Linear(z_r_dim, i_dim)
		self.r_decoder_post = nn.Linear(i_dim, r_out_dim)
	
		self.decoder = nn.Linear(((z_dim - z_r_dim) + r_out_dim), i_dim)
		self.decoder_post = nn.Linear(i_dim, in_dim)
								
	def encode(self, x):
		pre_x = torch.sigmoid(self.encoder_pre(x))
		return([self.encoder_mean(pre_x), self.encoder_var(pre_x)])

	def sample(self, mean_and_var):
		mean, var = mean_and_var
		std = var.mul(0.5).exp_()
		eps = torch.FloatTensor(std.size()).normal_()
		return(eps.mul(std).add_(mean))
	
	def decode(self, z):
		r_out = self.r_decode(z_r(z))
		return(self.rest_decode(z_rest(z), r_out))
		
	def r_decode(self, z_r):
		post_r_x = torch.sigmoid(self.r_decoder(z_r))
		return(torch.sigmoid(self.r_decoder_post(post_r_x)))

	def rest_decode(self, z_rest, r_out):
		post_x = torch.sigmoid(self.decoder(torch.cat((z_rest, r_out), dim = 1)))
		return(torch.sigmoid(self.decoder_post(post_x)))

	def forward(self, x):
		mean_and_var = self.encode(x)
		mean, var = mean_and_var
		z = self.sample(mean_and_var)
		r_out = self.r_decode(z_r(z))
		return([self.rest_decode(z_rest(z), r_out), r_out, mean, var])

# Construct our model by instantiating the class defined above
model = VAE(in_dim, i_dim, z_dim, z_r_dim, r_out_dim)

# Loss and Optimizer
reconstr_criterion = nn.BCELoss()

# balance hyperparameters
plain_reconstr_w = 1
rhythm_reconstr_w = 1
prior_normal_criterion_w = 1

def criterion(reconstr_x, x, reconstr_x_r, x_r, mean, var):
	plain_reconstr_criterion = reconstr_criterion(reconstr_x, x)
	rhythm_reconstr_criterion = reconstr_criterion(reconstr_x_r, x_r)
	# loss = 0.5 * sum(1 + log(var^2) - mean^2 - var^2)
	prior_normal_criterion = torch.sum(mean.pow(2).add_(var.exp()).mul_(-1).add_(1).add_(var)).mul_(-0.5)
	return((plain_reconstr_w * plain_reconstr_criterion)
		+ (rhythm_reconstr_w * rhythm_reconstr_criterion)
		+ (prior_normal_criterion_w * prior_normal_criterion))

optimizer = optim.Adam(model.parameters())

# Training
print('Training the model/network')

# Simple training, no mini batch, all examples are used during each epoch
for epoch in range(number_epochs):
	# ===================forward=====================
	X_pred, X_r_pred, mean, var = model(X_train)
	loss = criterion(X_pred, X_train, X_r_pred, X_r_train, mean, var)
	print('epoch: ' + str(epoch) + ' loss: ' + str(loss.item))
	# ===================backward====================
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	# ===================log========================

print('Model trained')

def create_scores_melodies_from_labels(labels_list):
	labels_list = torch.tensor(labels_list)
	data_list = model.decode(labels_list)
	data_list = data_list.detach()
	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(create_scores(1, data_list2))

print('Create specific melodies')

# Extreme labels
label_rh_min_rest_min = ([-10.] * z_r_dim) + ([-10.] * (z_dim - z_r_dim))
label_rh_max_rest_min = ([10.] * z_r_dim) + ([-10.] * (z_dim - z_r_dim))
label_rh_min_rest_max = ([-10.] * z_r_dim) + ([10.] * (z_dim - z_r_dim))
label_rh_max_rest_max = ([10.] * z_r_dim) + ([10.] * (z_dim - z_r_dim))

scores = create_scores_melodies_from_labels([label_rh_min_rest_min, label_rh_max_rest_min, label_rh_min_rest_max, label_rh_max_rest_max])

scores[0].write('midi', 'mid/auto_var_dis_irish_mel_pt_rh_min_rest_min.mid')
scores[1].write('midi', 'mid/auto_var_dis_irish_mel_pt_rh_max_rest_min.mid')
scores[2].write('midi', 'mid/auto_var_dis_irish_mel_pt_rh_min_rest_max.mid')
scores[3].write('midi', 'mid/auto_var_dis_irish_mel_pt_rh_max_rest_max.mid')

# Mixing up melody with rhythm melody profile

# Computing the label of one of the training example (first one)
first_training_example_label = model.sample(model.encode(X_train[0]))

# Create representation of a rhythm melody with 16th notes
rhythm_melody_16ths_data = [0.] * one_hot_size * config.max_number_time_steps

for i in range(config.max_number_time_steps):
	rhythm_melody_16ths_data[(i * one_hot_size) + number_special_elements_one_hot + int(one_hot_size / 2)] = 1.
	
rhythm_melody_16ths_data = torch.tensor(rhythm_melody_16ths_data)
rhythm_melody_16ths_label = model.sample(model.encode(rhythm_melody_16ths_data))

# Create representation of a rhythm melody with 4th notes
rhythm_melody_4ths_data = [0.] * one_hot_size * config.max_number_time_steps

for i in range(int(config.max_number_time_steps / 4)):
	index = i * one_hot_size
	rhythm_melody_4ths_data[index + number_special_elements_one_hot + int(one_hot_size / 2)] = 1.
	for j in range(1, 4):
		rhythm_melody_4ths_data[index + (j * one_hot_size)] = 1.

rhythm_melody_4ths_data = torch.tensor(rhythm_melody_4ths_data)
rhythm_melody_4ths_label = model.sample(model.encode(rhythm_melody_4ths_data))

# Constructing the hybrid labels - rhythm part (from 16ths or 4ths melody) concatenated with rest part (from 1st example)
# I reconstruct item per item the labels parts, otherwise I cannot concatenate the labels (which are tensors) because of an error...
# There should be a wiser way...
label_rhythm_16ths_latent_rhythm_part = []
label_rhythm_4ths_latent_rhythm_part = []

for i in range(z_r_dim):
	label_rhythm_16ths_latent_rhythm_part.append(rhythm_melody_16ths_label[i])
	label_rhythm_4ths_latent_rhythm_part.append(rhythm_melody_4ths_label[i])

label_first_example_latent_rest_part = []

for i in range(z_r_dim, z_dim):
	label_first_example_latent_rest_part.append(first_training_example_label[i])

label_first_training_example_rhythm_16ths  = label_rhythm_16ths_latent_rhythm_part + label_first_example_latent_rest_part
label_first_training_example_rhythm_4ths = label_rhythm_4ths_latent_rhythm_part + label_first_example_latent_rest_part

# Creating the hybrid melodies
score_first_training_example_rhythm_16ths = create_scores_melodies_from_labels([label_first_training_example_rhythm_16ths])[0]
score_first_training_example_rhythm_16ths.write('midi', 'mid/auto_var_dis_irish_mel_pt_ex1_16ths.mid')

score_label_first_training_example_rhythm_4ths = create_scores_melodies_from_labels([label_first_training_example_rhythm_4ths])[0]
score_label_first_training_example_rhythm_4ths.write('midi', 'mid/auto_var_dis_irish_mel_ex1_4ths.mid')