#!/usr/bin/env python3

import tensorflow as tf
import settings
from tensorflow.keras import layers
from keras.layers import Input, Dense, Conv2D, Flatten
from keras.models import Model, Sequential
import numpy as np
from sys import exit
import pickle

print('Reading samples from: {}'.format(settings.samples_path))

train_X = np.load(settings.samples_path)['arr_0']

n_samples = train_X.shape[0]
input_shape = settings.midi_resolution*128
train_X = train_X.reshape(n_samples, input_shape)

# encoder model
input_img = tf.keras.layers.Input(shape=(input_shape,))
encoded = tf.keras.layers.Dense(160, activation='relu')(input_img)
decoded = tf.keras.layers.Dense(input_shape, activation='sigmoid')(encoded)
autoencoder = tf.keras.models.Model(input_img, decoded)

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

autoencoder.fit(train_X, train_X, epochs=settings.epochs, batch_size=32)

autoencoder.save_weights(settings.model_path)
print("Model save to {}".format(settings.model_path))