from hyperopt import Trials, STATUS_OK, tpe
from hyperas import optim
from hyperas.distributions import choice, uniform
import numpy as np
from keras.models import Sequential
from keras.layers.core import Dense, Activation, Dropout
from keras.layers import LSTM
from keras.utils import np_utils
import pandas as pd
import os
from gensim.models import Doc2Vec as d2v
from keras.callbacks import TensorBoard


def data():
    '''
    Read pandas dataframe from pickle and prepare for keras input.
    :param datafile: Pickle containing pandas dataframe with training data
    :return: numpy array containing document vectors and list containing labels
             Label 0 - benign; 1 - Malware
    '''
    TRAINFILE = 'train_ds_1500_clean.pickle'
    TRAINVECS = d2v.load(os.path.join('DATA') + '/d2v_200_train_ds_1500_clean.model')
    TESTFILE = 'test_ds_1900_clean.pickle'
    test_data = pd.read_pickle(os.path.join('DATA') +
                               '/test_ds_1900_clean_inf_vectors_200.pickle',
                               compression='gzip')
    train_data = []
    train_labels = []
    test_labels = []
    train = pd.read_pickle(os.path.join('DATA') + '/' + TRAINFILE,
                           compression='gzip')
    for i in list(train.index):
        train_data.append(TRAINVECS[i])
        # train_labels.append(i)
        if train.loc[i]['label'] == 'MALWARE':
            train_labels.append(1)
        else:
            train_labels.append(0)
    train_data = np.array(train_data)
    test = pd.read_pickle(os.path.join('DATA') + '/' + TESTFILE,
                          compression='gzip')
    # test_labels = list(test.loc[:]['label'])
    for i in list(test.index):
        if test.loc[i]['label'] == 'MALWARE':
            test_labels.append(1)
        else:
            test_labels.append(0)

    train_labels = np_utils.to_categorical(train_labels, 2)
    print(len(train_labels))
    test_labels = np_utils.to_categorical(test_labels, 2)

    return train_data, train_labels, test_data, test_labels


def lstm_modl(train_data, train_labels, test_data, test_labels):
    '''
    Define the LSM model.
    '''
    np.random.seed(123)
    # INPUT_SHAPE = (1, 1, 200)
    model = Sequential()
    # model.add(Dropout({{uniform(0, 1)}}, input_shape=INPUT_SHAPE))

    model.add(LSTM(200))
    model.add(Dense(2))
    model.add(Activation({{choice(["softmax", "sigmoid", "hard_sigmoid"])}}))

    model.compile(loss={{choice(['categorical_crossentropy',
                         'binary_crossentropy'])}}, metrics=['accuracy'],
                  optimizer={{choice(['rmsprop', 'adam', 'sgd', 'nadam',
                                     'adamax'])}})

    history = model.fit(train_data, train_labels,
                        batch_size={{choice([32, 64, 128, 256, 512])}},
                        epochs={{choice([10, 20, 50, 100])}}, verbose=1,
                        validation_split={{uniform(0, 1)}},
                        callbacks=[TensorBoard(log_dir=os.path.join('LOGS/lstm'),
                                               histogram_freq=0,
                                               write_graph=True,
                                               write_grads=False,
                                               write_images=False,
                                               embeddings_freq=0,
                                               embeddings_layer_names=None,
                                               embeddings_metadata=None)])

    score, acc = model.evaluate(test_data, test_labels, verbose=1)

    print('History')
    print(history.history.keys())

    print('Test Score: {}'.format(score))
    print('Test Accuracy: {}'.format(acc))

    out = {'loss': -acc, 'status': STATUS_OK, 'model': model}

    return out


def main():
    train_data, train_labels, test_data, test_labels = data()
    best_run, best_model = optim.minimize(model=lstm_modl, data=data,
                                          algo=tpe.suggest, max_evals=20,
                                          trials=Trials())
    train_data, train_labels, test_data, test_labels = data()
    print("Evalutation of best performing model:")
    print(best_model.evaluate(test_data, test_labels))
    print("Best performing model chosen hyper-parameters:")
    print(best_run)

    best_model.save(os.path.join('MODELS') + '/lstm_d2v200.h5')


if __name__ == '__main__':
    main()
