import Basics as midi

from keras.models import Sequential
from keras.layers.core import Dense, Activation, Dropout, RepeatVector, Reshape
from keras.layers.recurrent import LSTM
from keras.preprocessing import sequence
import numpy as np
np.set_printoptions(threshold=np.nan)

midi_data = midi.extract_data()
lahead = 1

#Put music data in 3D numpy array
def gen_note_data(note_data):
    notes = np.zeros((len(note_data[0]), 1, 4), dtype=object)
    for i in range(len(note_data[0])):
        notes[i, 0] = note_data[0][i]
    return notes

n_in = gen_note_data(midi_data)

#Numpy array with expected output. For each note at time t in the sequence
#we get the note at t+1
n_out = np.zeros((len(n_in), 4), dtype=object)
for i in range(len(n_in) - lahead):
    n_out[i] = n_in[i + 1:i + lahead + 1]


model = Sequential()
model.add(LSTM(100, input_dim=4, batch_input_shape=(len(n_in[0]), 1, 4), return_sequences=True))
model.add(LSTM(100, batch_input_shape=(len(n_in[0]), 1, 4)))
model.add(Dense(4))

model.compile(loss='categorical_crossentropy',
              optimizer='rmsprop')
print 'Training...'

model.fit(n_in, n_out, batch_size=len(n_in), nb_epoch=100, show_accuracy=True, verbose=2)

print model.predict(n_in, batch_size=len(n_in))