# AutoBach
# Autoencoder generation inspired from DeepHear (Felix Sun)
# With some import (dataset and some code) from DeepBach (Gaetan Hadjeres)
# Adaptation by Jean-Pierre Briot
# 10/09/2018

import pickle
import numpy as np
# import math
import random

from music21 import corpus, converter, stream, note, duration, analysis, interval

from data_utils import BACH_DATASET, START_SYMBOL, END_SYMBOL, to_onehot, indexed_chorale_to_score, standard_note
from keras.engine import Input, Model
from keras.layers import Dense, Activation, Dropout
from keras.models import Sequential, load_model
from keras import regularizers

# Dataset -- Analog to MiniBach

SOP_INDEX = 0

# Construct an return input (X) and output (y) examples
def features_and_target(chorale, time_index,
                        timesteps, num_pitches,
                        num_voices):
    """                       
    :param chorale: time major 
    :param time_index: 
    :param timesteps: 
    :param num_pitches: 
    :param num_voices: 
    :return: 
    """
    
    features = np.array([to_onehot(time_slice, num_pitches[SOP_INDEX])
    # create np.array of elements to_onehot(time_slice, num_pitches[SOP_INDEX]) with time_slice varying within chorale[time_index: time_index + timesteps, SOP_INDEX]
                         for time_slice in chorale[time_index: time_index + timesteps, SOP_INDEX]
                         ])
    # (time, num_pitches)
    # flatten the array
    features = np.reshape(features, (timesteps * num_pitches[SOP_INDEX],))
    # (time * num_pitches)

    target = [to_onehot(time_slice, num_pitches[voice_index])
    # create np.array of elements to_onehot(time_slice, num_pitches[voice_index]) with voice_index varying within [1 num_voices] and time_slice varying within chorale[time_index: time_index + timesteps, voice_index]]
              for voice_index in range(1, num_voices)
              for time_slice in chorale[time_index: time_index + timesteps, voice_index]]
    # flatten the array
    target = np.concatenate(target)

    # target[(i * timesteps * num_pitches +  j * num_pitches + k)] = voice i, timestep j, pitch k
    return features, target

def generator_pianoroll(batch_size, timesteps,
                        phase='train', percentage_train=0.8, pickled_dataset=BACH_DATASET):
    """
     Returns a generator of
            (left_features,
            central_features,
            right_features,
            beats,
            metas,
            labels,
            fermatas) tuples
            where fermatas = (fermatas_left, central_fermatas, fermatas_right)
    """

    X, X_metadatas, voice_ids, index2notes, note2indexes, metadatas = pickle.load(open(pickled_dataset, 'rb'))
    num_pitches = list(map(lambda x: len(x), index2notes))
    num_voices = len(voice_ids)
    # Set chorale_indices
    if phase == 'train':
        chorale_indices = np.arange(int(len(X) * percentage_train))
    if phase == 'test':
        chorale_indices = np.arange(int(len(X) * percentage_train), len(X))
    if phase == 'all':
        chorale_indices = np.arange(int(len(X)))

    features = []
    targets = []
    batch = 0

    while True:
        chorale_index = np.random.choice(chorale_indices)
        extended_chorale = np.transpose(X[chorale_index])

        padding_dimensions = (timesteps,) + extended_chorale.shape[1:]

        start_symbols = np.array(list(map(lambda note2index: note2index[START_SYMBOL], note2indexes)))
        end_symbols = np.array(list(map(lambda note2index: note2index[END_SYMBOL], note2indexes)))

        extended_chorale = np.concatenate((np.full(padding_dimensions, start_symbols),
                                           extended_chorale,
                                           np.full(padding_dimensions, end_symbols)),
                                          axis=0)

        chorale_length = len(extended_chorale)

        time_index = np.random.randint(0, chorale_length - timesteps)

        feature, target = features_and_target(chorale=extended_chorale, time_index=time_index,
                                              timesteps=timesteps, num_pitches=num_pitches,
                                              num_voices=num_voices)

        features.append(feature)
#      targets.append(target)
        targets.append(feature)

        batch += 1

        # if there is a full batch
        if batch == batch_size:
            next_element = (
                np.array(features, dtype=np.float32),
                np.array(features, dtype=np.float32))

            yield next_element

            batch = 0

            features = []
            targets = []

_, _, _, index2notes, note2indexes, _ = pickle.load(open(BACH_DATASET, 'rb'))

num_pitches = list(map(lambda x: len(x), index2notes))

num_pitches_melody = num_pitches[0]

def from_melody_output_to_note_index_melody(melody):
	note_index_melody = []
	for i in range(0, timesteps * num_pitches_melody, num_pitches_melody):
		index = np.argmax(melody[i:i + num_pitches_melody])
		note_index_melody.append(index)
	return(note_index_melody)

SUBDIVISION = 4

# Rest : 'rest' : 9
rest_index = note2indexes[0]['rest']
# Hold : '__' : 1
hold_index = note2indexes[0]['__']

def from_note_index_melody_to_score(note_index_melody):
	# simplified version (hold : repeated notes; no rest/silence)
	score = stream.Score()
	previous_note_index = 0
	current_duration = 1
	for i, note_index in enumerate(note_index_melody):
		print('i: ' + str(i) + ' note_index: ' + str(note_index) + ' note: ' + str(index2notes[0][note_index]))
		if note_index == hold_index:
			# hold
			new_note = standard_note(index2notes[0][previous_note_index])
			new_note.duration = duration.Duration(current_duration / SUBDIVISION)
			score.append(new_note)
			previous_note_index = note_index
		elif note_index == rest_index:
			# rest
			'rien'
		else:
			# note
			new_note = standard_note(index2notes[0][note_index])
			new_note.duration = duration.Duration(current_duration / SUBDIVISION)		
			score.append(new_note)
			previous_note_index = note_index
	return score

# batch_size = 128
batch_size = 32

timesteps = 12

# Generate example
# load dataset
gen_train = generator_pianoroll(batch_size,
	timesteps = timesteps,
	phase = 'train')
gen_test = generator_pianoroll(batch_size,
  	timesteps = timesteps,
	phase = 'test')

features, target = next(gen_train)
input_size = features.shape[-1]
output_size = target.shape[-1]

# Setup the parameters
input_layer_size  = input_size
hidden_layer_size =  200	# 200 hidden units

# Create the autoencoder model

autoencoder = Sequential()

autoencoder.add(Dense(hidden_layer_size,
	input_dim = input_layer_size,
	activation = 'sigmoid',
	activity_regularizer = regularizers.l1(10e-5)))
autoencoder.add(Dense(input_layer_size,
	activation = 'sigmoid'))

# Organize decoder input

# Encoder_layer = autoencoder.layers[0]
Decoder_layer = autoencoder.layers[1]

Decoder_Input_layer = Input(shape = (hidden_layer_size, ))

decoder = Model(input = Decoder_Input_layer,
	output = Decoder_layer(Decoder_Input_layer))

autoencoder.compile(optimizer = 'adadelta',
	loss = 'binary_crossentropy',
	metrics = ['accuracy'])

autoencoder.fit_generator(generator = gen_train,
    	steps_per_epoch = 10,
	epochs = 10,
	validation_data = gen_test,
	validation_steps = 10)

label1 = []
for i in range(hidden_layer_size):
	label1.append(random.uniform(-1, 1))

label2 = []
for i in range(hidden_layer_size):
	label2.append(random.uniform(-10, 10))

label3 = []
for i in range(hidden_layer_size):
	label3.append(random.uniform(-1, 1))

melody1 = decoder.predict(np.array([label1]))[0]
melody2 = decoder.predict(np.array([label2]))[0]
melody3 = decoder.predict(np.array([label3]))[0]

score1 = from_note_index_melody_to_score(from_melody_output_to_note_index_melody(melody1))
score1.write("midi", "auto1.mid")

score2 = from_note_index_melody_to_score(from_melody_output_to_note_index_melody(melody2))
score2.write("midi", "auto2.mid")

score3 = from_note_index_melody_to_score(from_melody_output_to_note_index_melody(melody3))
score3.write("midi", "auto3.mid")