from keras.models import Sequential
from keras.layers import Dense, Activation, Dropout
from keras.layers.recurrent import GRU, LSTM
from keras.preprocessing import sequence
import numpy as np
from BuySessionData import *


def generator(batch_size, test_type='train'):
    data = BuySessionData('/home/eric/data/formatted/buy-sessions-train.dat',
                          '/home/eric/data/formatted/buy-sessions-test.dat')
    while True:
        x, y = data.next_train_batch(batch_size) if test_type == 'train'\
            else data.next_test_batch(batch_size, test_type)
        x = sequence.pad_sequences(x)
        yield (x, y)

# Parameters
HIDDEN_SIZE = 64
BATCH_SIZE = 50
LAYERS = 2
FEATURES = 4 + 339
CLASSES = 2
EPOCH_SAMPLES = 100000
EPOCHS = 1
TEST_SAMPLES = 25000

model = Sequential()
model.add(LSTM(HIDDEN_SIZE, input_shape=(None, FEATURES), return_sequences=True))
model.add(Activation('relu'))

model.add(LSTM(HIDDEN_SIZE, input_shape=(None, FEATURES)))
model.add(Activation('relu'))

model.add(Dense(CLASSES))
model.add(Activation('softmax'))

model.compile(loss='binary_crossentropy',
              optimizer='rmsprop',
              metrics=['binary_accuracy'])

model.fit_generator(generator(BATCH_SIZE),
                    samples_per_epoch=EPOCH_SAMPLES,
                    nb_epoch=EPOCHS,
                    class_weight={0: 5.5, 1: 94.5})

score, acc = model.evaluate_generator(generator(BATCH_SIZE, 'buys'), val_samples=TEST_SAMPLES)
print('Test accuracy buys:', acc)

score, acc = model.evaluate_generator(generator(BATCH_SIZE, 'non_buys'), val_samples=TEST_SAMPLES)
print('Test accuracy non buys:', acc)