import cPickle as pickle
import os

dir_path = os.path.dirname(os.path.realpath(__file__))

processed_dir = dir_path+'/../data/processed/'

with open(processed_dir+"x_manifest_mal.pickle", "rb") as f:
    x_manifest_mal = pickle.load( f )

with open(processed_dir+"x_code_mal.pickle", "rb") as f:
    x_code_mal = pickle.load( f )

valid_mal = int(x_manifest_mal.shape[0]/10)

mal_train = [x_manifest_mal[:-valid_mal], x_code_mal[:-valid_mal]]
mal_val = [x_manifest_mal[-valid_mal:], x_code_mal[-valid_mal:]]

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



code_csr, manifest_csr = mal_val

def adversarial_eval(model, max_attempts=20, batch=15, max_samples=None, try_on=None):
    sess = K.get_session()
    
    feed_dict = {'keras_learning_phase:0': 0}

    done = 0

    flipped = 0
    total = 0
    val_samples = code_csr.shape[0]
    if max_samples:
        val_samples = min(max_samples, val_samples)

    flipped_sample_batches = []
    
    while done < val_samples:
        
        end = min(done+batch, val_samples-1)

        curr_code = code_csr[done:end].todense()
        curr_manifest = manifest_csr[done:end].todense()

        for j in xrange(max_attempts):
            feed_dict[model.inputs[0]] = curr_code
            feed_dict[model.inputs[1]] = curr_manifest

            ret = sess.run([tf.gradients(-model.outputs[0], model.inputs[1]) * (1-model.inputs[1]), model.outputs[0]], feed_dict=feed_dict)

            grads = ret[0][0]
            is_mal = (ret[1].squeeze() > 0.5).tolist()
            
            #print (done, j)

            if j == 0:
                orig_mal = is_mal
                #print is_mal > 0.5
                
            if j == max_attempts - 1:
                #from IPython import embed; embed()

                for orig, peturbed in zip(orig_mal, is_mal):
                    if orig:
                        total = total + 1
                        if not peturbed:
                            flipped = flipped + 1
                            
                print orig_mal
                print is_mal
                print ("Flipped ", flipped, " of ", total)
                flipped_sample_batches.append((orig_mal, curr_code, curr_manifest))

            for i in xrange(end - done):
                idx = np.argmax(grads[i])
                #print (i, idx)
                curr_manifest[i,idx] = 1

            #print orig_mal, is_mal
        
        done = done + batch

    total2 = 0
    flipped2 = 0
    for (orig_mal, curr_code, curr_manifest) in flipped_sample_batches:
        feed_dict[try_on.inputs[0]] = curr_code
        feed_dict[try_on.inputs[1]] = curr_manifest

        ret = sess.run([tf.gradients(-try_on.outputs[0], try_on.inputs[1]) * (1-try_on.inputs[1]), try_on.outputs[0]], feed_dict=feed_dict)

        is_mal = (ret[1].squeeze() > 0.5).tolist()

        if True:
            if True:
                for orig, peturbed in zip(orig_mal, is_mal):
                    if orig:
                        total2 = total2 + 1
                        if not peturbed:
                            flipped2 = flipped2 + 1

                print orig_mal
                print is_mal
                print ("(Transfer attack) Flipped ", flipped2, " of ", total2)

    return (flipped, total, val_samples)
