from data.load import x_manifest, x_manifest_mal, x_code, x_code_mal
from data.load import train_gen, x_test, y_test, train_good, train_mal
from model.model import model

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

import os

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

for mwr in [0.3, 0.4, 0.5]:

    for temp in [10, 33, 66, 100, 150, 300]:

        m = model(distil_temp=temp)

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

        epochs = 10

        model_path = model_dir + "regnonegsoftmax3.l1-0.l2-0.mwr-{}-epoch-{}-distil-train-{}.hdf5".format(mwr, epochs, temp)

        samples_per_epoch = int(train_mal / mwr)
        
        print ("Looking for {}".format(model_path))

        if os.path.isfile(model_path):
            m.load_weights(model_path)
        else:
            history = m.fit_generator(train_gen(mwr=mwr), samples_per_epoch=samples_per_epoch,
                            nb_epoch=epochs, pickle_safe=False, #nb_worker=4,
                            #validation_data=(x_test, y_test),
                            )
            m.save_weights(model_path)

        y_good_distill = m.predict([x_manifest, x_code], batch_size=2000).squeeze()
        y_mal_distill = m.predict([x_manifest_mal, x_code_mal], batch_size=2000).squeeze()

        import cPickle as pickle

        #Save distillation output
        with open( "../data/processed/distil/y_good_distill.mwr-{}.temp-{}.pickle".format(mwr, temp), "wb" ) as f:
            pickle.dump( y_good_distill, f )

        with open( "../data/processed/distil/y_mal_distill.mwr-{}.temp-{}.pickle".format(mwr, temp), "wb" ) as f:
            pickle.dump( y_mal_distill, f )

        K.clear_session()