#!/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
import numpy as np
import sys
import pickle

train_data_path_X = sys.argv[1]
train_data_path_y = sys.argv[2]
save_model_path = sys.argv[3]
epochs = int(sys.argv[4])

# model architecture
# model = Sequential()
# model.add(LSTM(128, activation='relu', input_shape=(96,128)))
# model.add(RepeatVector(96))
# model.add(LSTM(128, activation='softmax', return_sequences=True))
# model.add(TimeDistributed(Dense(128)))
#
# model.compile(optimizer='adam',
#               loss='categorical_crossentropy',
#               metrics=['accuracy'])

model = Sequential()
model.add(LSTM(128,input_shape=(96, 128),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(512))
# model.add(Dropout(0.3))
# model.add(Dense(128))
model.add(Activation('softmax'))
model.add(Reshape((4, 128)))
model.compile(loss='categorical_crossentropy', optimizer='rmsprop')

# load training data
print('Reading samples from: {}'.format(train_data_path_X))
train_X = np.load(train_data_path_X)['arr_0']
train_y = np.load(train_data_path_y)['arr_0']

# model training
model.fit(train_X, train_y, epochs=epochs, batch_size=32)

# save trained model
pickle_path = '{}.pickle'.format(save_model_path)
pickle.dump(model, open(pickle_path,'wb'))
print("Model save to {}".format(pickle_path))