#!/usr/bin/env python3

import tensorflow as tf
import settings
from tensorflow.keras import layers
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
import numpy as np
import sys
import pickle

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 = np.unique(train_X)
    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')

# model training
print('Training...')
model.fit(train_X, train_y, epochs=epochs, batch_size=64)

# save trained model
pickle.dump(model, open(save_model_path,'wb'))
pickle.dump(int_to_note, open('{}_dict'.format(save_model_path),'wb'))
print('Done!')
print("Model saved to: {}".format(save_model_path))