# Do the import here for serialization purposes
from data.load import train_gen, x_test, y_test, train_good, train_mal
from model.model import model, as_dense
from adversarial import adversarial_eval

import keras
import keras.backend as K
import numpy as np

import tensorflow as tf
from sklearn.metrics import classification_report, roc_curve, auc, confusion_matrix
from hyperopt import STATUS_OK

import os

def do_run(space):

    model_dir = '/home/alex/research/malware/models/'

    l1 = space['l1']
    l2 = space['l2']
    mwr = space['mwr']
    epochs = space['epochs']
    hard = space['hard']
    distil_temp = space['distil'] if 'distil' in space else 1.0
    restriction = space['restriction'] if 'restriction' in space else 'weights'
    two_outs = space['two_outs'] if 'two_outs' in space else False

    m = model(l1=l1, l2=l2, hard_constraint=hard, distil_temp=distil_temp, restriction=restriction, two_outs=two_outs)

    samples_per_epoch = int(train_mal / mwr)

    optimizer = keras.optimizers.Adam()
    m.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy', 'matthews_correlation'])

    name = "regnonegsoftmax3.l1-{}.l2-{}.mwr-{}-epoch-{}".format(l1, l2, mwr, epochs)

    if restriction and restriction != 'weights':
      name = name + "-" + restriction

    if distil_temp != 1.0:
        name = name + "-distil-{}".format(distil_temp)

    if hard:
        name = name + "-hard"

    print ("Begining trial "+name) 
    model_path = model_dir+name+".hdf5"
    #saver = keras.callbacks.ModelCheckpoint(model_dir+name+'.weights.{epoch:02d}-{val_loss:.2f}.hdf5', monitor='val_loss', verbose=0, save_best_only=False, save_weights_only=False, mode='auto')
    callbacks = [] #[saver]

    if os.path.isfile(model_path):
        history = None
        m.load_weights(model_path)
    else:
        split = True
        if split:
            history = m.fit_generator(train_gen(mwr=mwr, distil_temp=distil_temp),
                            samples_per_epoch=samples_per_epoch,
                            nb_epoch=epochs, pickle_safe=False, #nb_worker=4,
                            #validation_data=(x_test, y_test),
                            callbacks=callbacks)
        else:
            history = m.fit(x_train, y_train,  batch_size=1000,
                            #train_gen(), samples_per_epoch=samples_per_epoch,
                            nb_epoch=epochs,
                            #validation_data=(x_test, y_test),
                            callbacks=callbacks)

        history = history.history
        m.save_weights(model_path)

    y_pred = m.predict(x_test)
    y_truth = y_test

    y_round = np.copy(y_pred)
    y_round[y_round < 0.5] = 0
    y_round[y_round >= 0.5] = 1
    
    confusion = confusion_matrix(y_truth, y_round)

    print (confusion)

    fpr = float(confusion[0][1]) / (confusion[0][0] + confusion[0][1]) * 100
    fnr =float(confusion[1][0]) / (confusion[1][0] + confusion[1][1]) * 100

    print ("FPR: {}% FNR: {}%".format(fpr, fnr))

    try_on = model(l1=0.0, l2=0.0, hard_constraint=False, distil_temp=100)
    try_on.load_weights("models/regnonegsoftmax3.l1-0.0.l2-0.0.mwr-0.3-epoch-10-distil-100.hdf5")

    max_samples = 45
    adv_res = adversarial_eval(as_dense(m), max_samples=max_samples, try_on=as_dense(try_on))

    print (adv_res)

    K.clear_session()

    return {
        'loss': 1,
        'confusion': confusion,
        'adversarial': adv_res,
        'history': history,
        'status': STATUS_OK
    }

