# Simplistic melody generation LSTM example
# Jean-Pierre Briot
# 11/09/2018

from keras.models import Sequential
from keras.layers import LSTM, Dense
import numpy as np
import math
import random
import copy
from music21 import stream, note, duration

note_name_list = ['DO', 'DO#', 'RE', 'RE#', 'MI', 'FA', 'FA#', 'SOL', 'SOL#', 'LA', 'LA#', 'SI']
# corresponding values (index of the name in the list) : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

# From a symbolic note name to the note value
# Ex: note_name_to_value('A') -> 9
def note_name_to_value(note_name):
	return(note_name_list.index(note_name))

# From a note value to the symbolic note name
def note_value_to_name(note_value):
	return(note_name_list[round(math.fmod(round(note_value), 12))])

# From a melody note name to the note value melody
def note_name_melody_to_value(note_name_melody):
	note_value_melody = []
	for i in range(len(note_name_melody)):
		note_value_melody.append(note_name_to_value(note_name_melody[i]))
	return(note_value_melody)

# From a melody note value to the symbolic note name melody
def note_value_melody_to_name(note_value_melody):
	note_name_melody = []
	for i in range(len(note_value_melody)):
		note_name_melody.append(note_value_to_name(note_value_melody[i]))
	return(note_name_melody)

# One hot encoding of a note value (into a list, not an array)
# Ex: note_value_to_one_hot(2) ('RE') -> [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]
def note_value_to_one_hot(note_value):
	note_one_hot_list = []
	for i in range(number_notes):
		if i == note_value:
			note_one_hot_list.append(1)
		else:
			note_one_hot_list.append(0)
	return(note_one_hot_list)

# One hot encoding of a note value melody
def note_value_melody_to_one_hot(note_value_melody):
	note_one_hot_melody = []
	for i in range(len(note_value_melody)):
		note_one_hot_melody.append(note_value_to_one_hot(note_value_melody[i]))
	return(note_one_hot_melody)

# Generate a random note value (integer between 0 and number_notes - 1)
def random_note_value():
	return(random.randint(0, number_notes - 1))

# Generate a random melody (of notes values) with time steps number of notes
def random_note_value_melody():
	note_value_melody = []
	for i in range(time_steps):
		note_value_melody.append(random_note_value())
	return(note_value_melody)

# Note name of a one hot encoded note (in a numpy array)
# Ex: note_one_hot_to_name(np.array([0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0])) -> 'RE'
# Deterministic interpretation : the note with the highest probability
# argmax returns the index of the array with the highest value
def note_one_hot_to_name(note_one_hot_encoding):
	return(note_value_to_name(np.argmax(note_one_hot_encoding)))

# Sample a note from a multinoulli probability distribution
# Stochastic interpretation : sample a note along their respective probabilities
# Return a one-hot encoding
def sample_note_one_hot(note_one_hot_encoding):
	return(np.random.multinomial(1, note_one_hot_encoding))

# Note name melody of a one hot encoded note melody (in a numpy array)
# Deterministic interpretation : the note with the highest probability
def note_one_hot_melody_to_name(note_one_hot_melody):
	note_name_melody = []
	for i in range(len(note_one_hot_melody)):
		note_name_melody.append(note_one_hot_to_name(note_one_hot_melody[i]))
	return(note_name_melody)

def note_name_melody_to_score(note_name_melody):
	# simplified version (hold : repeated notes; no rest/silence)
	score = stream.Score()
	for i in range(len(note_name_melody)):
		new_note = note.Note(note_name_to_value(note_name_melody[i]) + 60)		# DO : 0 -> DO : 60 (C4 : middle C)
		new_note.duration = duration.Duration(1)		
		score.append(new_note)
	return score

time_steps = 1
# time_steps = 8
number_notes = len(note_name_list)

# is_deterministic = True
is_deterministic = False

input_size = number_notes
lstm_layer_size = 32
output_size = number_notes

# expected input data shape: (batch_size, time_steps, number_notes)
model = Sequential()
model.add(LSTM(lstm_layer_size,
               input_shape = (time_steps, input_size)))
model.add(Dense(output_size,
	activation = 'softmax'))

model.compile(loss = 'categorical_crossentropy',
	optimizer = 'rmsprop',
	metrics =	['accuracy'])

# Generate dummy training data
# X_train = np.random.random((100, time_steps, input_size))
# y_train = np.random.random((100, output_size))

full_moon_melody = ['DO', 'DO', 'RE', 'MI', 'RE', 'DO', 'MI', 'RE', 'RE', 'DO']
moon_melody = full_moon_melody[0:time_steps]

def hot_m(note_name_melody):
	return(note_value_melody_to_one_hot(note_name_melody_to_value(note_name_melody)))

def hot_n(note_name):
	return(note_value_to_one_hot(note_name_to_value(note_name)))

# Generate DO major chord arpeggio serie (with only ONE example !)
# if time steps = 1
X_train = np.array([hot_m(['DO']),
			hot_m(['MI']),
			hot_m(['SOL']),
			hot_m(['SI']),
			hot_m(['DO'])])
y_train = np.array([hot_n('MI'),
			hot_n('SOL'),
			hot_n('SI'),
			hot_n('DO'),
			hot_n('MI')])

# if time steps = 8
# X_train = np.array([hot_m(['DO', 'MI', 'SOL', 'SI', 'DO', 'MI', 'SOL', 'SI'])])
# y_train = np.array([hot_n('DO')])

model.fit(X_train,
	y_train,
	epochs = 10)

print('Prediction from single note melodies')

for i in range(number_notes):
	current_note_name_melody = [note_name_list[i]] * time_steps
	current_note_value_melody = note_name_melody_to_value(current_note_name_melody)
	current_note_melody_one_hot = note_value_melody_to_one_hot(current_note_value_melody)
	current_input = np.array([current_note_melody_one_hot])
	next_note_one_hot = model.predict(current_input)[0]
	if not is_deterministic:
		next_note_one_hot = sample_note_one_hot(next_note_one_hot)
	next_note_name = note_one_hot_to_name(next_note_one_hot)
	print(str(current_note_name_melody) + ' -> ' + str(next_note_name))

print('Prediction from random notes melody')

current_note_value_melody = random_note_value_melody()
current_note_name_melody = note_value_melody_to_name(current_note_value_melody)
current_note_melody_one_hot = note_value_melody_to_one_hot(current_note_value_melody)
current_input = np.array([current_note_melody_one_hot])
next_note_one_hot = model.predict(current_input)[0]
if not is_deterministic:
	next_note_one_hot = sample_note_one_hot(next_note_one_hot)
next_note_name = note_one_hot_to_name(next_note_one_hot)
print(str(current_note_name_melody) + ' -> ' + str(next_note_name))

def generate_melody(note_name_melody_beginning, generated_melody_additional_length):
# Warning : note_name_melody_beginning must be time_steps long
	note_name_melody_generated = copy.copy(note_name_melody_beginning)
	note_one_hot_melody_beginning = note_value_melody_to_one_hot(note_name_melody_to_value(note_name_melody_beginning))
	current_input = np.array([note_one_hot_melody_beginning])
	for i in range(generated_melody_additional_length):
		next_note_one_hot = model.predict(current_input)[0]
		if not is_deterministic:
			next_note_one_hot = sample_note_one_hot(next_note_one_hot)
		next_note_name = note_one_hot_to_name(next_note_one_hot)
		note_name_melody_generated.append(next_note_name)
		current_input = np.delete(current_input, 0, 1)
		current_input = np.append(current_input, [[next_note_one_hot]], 1)
	return(note_name_melody_generated)

print('Prediction from fixed notes melody: ' + str(moon_melody))

melody_beginning = moon_melody
note_name_melody_generated = generate_melody(melody_beginning, 24)
print(str(melody_beginning) + ' -> ' + str(note_name_melody_generated))

score = note_name_melody_to_score(note_name_melody_generated)
score.write("midi", "lstm1.mid")

melody_beginning = moon_melody
note_name_melody_generated = generate_melody(melody_beginning, 24)
print(str(melody_beginning) + ' -> ' + str(note_name_melody_generated))

score = note_name_melody_to_score(note_name_melody_generated)
score.write("midi", "lstm2.mid")