#!python3
#!/usr/bin/env python3
import sys
import pickle
import settings

import numpy as np
from keras.layers import Input, Dense, Conv2D, Flatten, LSTM, Dropout, TimeDistributed, RepeatVector, Activation, Bidirectional, Reshape
from keras.models import Model, Sequential
from keras.utils.np_utils import to_categorical


def load_data(samples_path):
    print('Loading... {}'.format(train_data_path))
    train_X = np.load(train_data_path, allow_pickle=True)['arr_0']
    train_y = np.load(train_data_path, allow_pickle=True)['arr_1']
    return train_X, train_y

# TODO: make transformer class with fit, transform and reverse definitions
def preprocess_samples(train_X, train_y):
    vocab_X = np.unique(train_X)
    vocab_y = np.unique(train_y)
    vocab = np.concatenate([vocab_X, vocab_y])
    n_vocab = vocab.shape[0]
    note_to_int = dict((note, number) for number, note in enumerate(vocab))
    int_to_note = dict((number, note) for number, note in enumerate(vocab))
    _train_X = []
    _train_y = []
    for sample in train_X:
        # TODO: add normalizasion
        _train_X.append([note_to_int[note] for note in sample])

    train_X = np.array(_train_X).reshape(train_X.shape[0], train_X.shape[1], 1)
    train_y = np.array([note_to_int[note] for note in train_y]).reshape(-1,1)
    train_y = to_categorical(train_y)

    return train_X, train_y, n_vocab, int_to_note

train_data_path = sys.argv[1]

train_X, train_y = load_data(train_data_path)
train_X, train_y, n_vocab, int_to_note = preprocess_samples(train_X, train_y)

save_model_path = sys.argv[2]
epochs = int(sys.argv[3])

model = Sequential()
model.add(LSTM(512, input_shape=(train_X.shape[1], train_X.shape[2]), return_sequences=True))
model.add(Dropout(0.3))
model.add(LSTM(512, return_sequences=True))
model.add(Dropout(0.3))
model.add(LSTM(512))
model.add(Dense(256))
model.add(Dropout(0.3))
model.add(Dense(n_vocab))
model.add(Activation('softmax'))
model.compile(loss='categorical_crossentropy', optimizer='rmsprop')

# This code will train our model, with given by parameter number of epochs
print('Training...')
model.fit(train_X, train_y, epochs=epochs, batch_size=64)

# it saves model, and additional informations of model 
# that is needed to generate music from it 
pickle.dump(model, open(save_model_path,'wb'))
pickle.dump((int_to_note, n_vocab, train_X.shape[1]), open('{}_dict'.format(save_model_path),'wb'))
print('Done!')
print("Model saved to: {}".format(save_model_path))