import numpy as np
np.random.seed(123)
from keras.datasets import imdb
from keras.preprocessing.sequence import pad_sequences
from keras.layers import *
from keras.models import Sequential, Model
from keras.callbacks import EarlyStopping
from keras import optimizers
import matplotlib.pyplot as plot
import time


MAXLEN = 500
BATCHSIZE = 16
EMBSIZE = 50
HIDDENSIZE = 50
KERNELSIZE = 5
VOCABSIZE = 10000

MAXEPOCHS = 20



(x_train, y_train), (x_test, y_test) = imdb.load_data(path = "imdb.npz",
                                                      num_words = VOCABSIZE,
                                                      skip_top = 0,
                                                      maxlen = MAXLEN,
                                                      start_char = 1,
                                                      oov_char = 2,
                                                      index_from = 3)

y_train = np.expand_dims(y_train, -1)
y_test = np.expand_dims(y_test, -1)



x_test = x_test[:1000]
y_test = y_test[:1000]

print("# training samples: {}; # validation samples: {}".format(len(x_train), len(x_test)))

BATCHES_PER_EPOCH = len(x_train) // BATCHSIZE

def generator(x, y, return_positions = False):
    while True:
        for i in range(0, len(x), BATCHSIZE):
            x_batch = x[i:i+BATCHSIZE] # DONE
            print(x_batch)
            y_batch =y[i:i+BATCHSIZE] # DONE
                
            yield(pad_sequences(x_batch), y_batch)

def build_gru_model():
    embedding_L = Embedding(input_dim = VOCABSIZE, output_dim = EMBSIZE, mask_zero = True)
    gru_L = Bidirectional(GRU(units = HIDDENSIZE // 2))
    output_L = Dense(units = 1, activation = "sigmoid")
    return Sequential([embedding_L, dropout_L, gru_L, dropout_L, output_L])

def build_cnn_model():
    #DONE
    embedding_L = Embedding(input_dim = VOCABSIZE, output_dim = EMBSIZE, mask_zero = False)
    dropout_L = Dropout(0.25)
    #Posmotret
    conv_L = Convolution1D(convolution_num_filters, convolution_filter_length, input_length=EMBSIZE,
                         input_dim=EMBSIZE, border_mode='same', activation='tanh')(inp)
    maxpool_L = MaxPooling1D(pool_length=sequence_length)(conv)
    dropout_L = Dropout(0.25)
    output_L = Dense(units = 1, activation = "sigmoid")
    return Sequential([embedding_L, dropout_L, conv_L, maxpool_L, dropout_L, output_L])


def build_emb_model():
    embedding_L = Embedding(input_dim = VOCABSIZE, output_dim = EMBSIZE, mask_zero = False)
    dropout_L = Dropout(0.25)
    avaragepool_L = AveragePooling1D(pool_length=2, stride=None, border_mode='valid')
    dropout_L = Dropout(0.25)
    output_L = Dense(units = 1, activation = "sigmoid")
    return Sequential([embedding_L, dropout_L, avaragepool_L , dropout_L, output_L])
    # Done

def train_model(model, x_train, y_train, x_test, y_test):
    sgd = optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
    model.compile(loss='categorical_crossentropy', optimizer="sgd", metrics=['accuracy']) # Done
    
    earlystop = EarlyStopping(monitor = "val_acc", patience = 7)
    
    history = model.fit_generator(generator(x_train, y_train),
                                  steps_per_epoch = BATCHES_PER_EPOCH,
                                  validation_data = (pad_sequences(x_test), y_test),
                                  epochs = MAXEPOCHS, callbacks = [earlystop])
    
    return history.history

models = {}

#uncomment all models that you have implemented:
#models["gru+attn"] = build_gru_attn_model()
models["gru"] = build_gru_model()
#models["cnn"] = build_cnn_model()
#models["emb"] = build_emb_model()

histories = {}
traintimes = {}

for name in sorted(models.keys()):
    print("Training", name)
    before = time.time()
    histories[name] = train_model(models[name], x_train, y_train, x_test, y_test)
    duration = time.time() - before
    traintimes[name] = duration / len(histories[name]["loss"]) / BATCHES_PER_EPOCH