import os, sys, random, time, math
import numpy as np
import tensorflow as tf
from tensorflow.python.training.gradient_descent import GradientDescentOptimizer
from tensorflow.examples.tutorials.mnist import input_data
import keras
from cleverhans.utils_tf import model_train, model_eval
import cleverhans.attacks_tf as attack
from cleverhans.attacks import SaliencyMapMethod
import scipy.misc

def weight_variable(shape, trainable=True):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial, trainable=trainable)

def bias_variable(shape, trainable=True):
    initial = tf.constant(0.0, shape=shape)
    return tf.Variable(initial, trainable=trainable)

# define classifier
layers_cls = [784, 1000, 500, 10]
n = len(layers_cls)
w, b = [0 for i in range(n)], [0 for i in range(n)]
with tf.variable_scope('classifier'):
    for i in range(1, n):
        w[i] = weight_variable([layers_cls[i - 1], layers_cls[i]])
        b[i] = bias_variable([layers_cls[i]])

def classifier(x):
    n = len(layers_cls)
    z, h = [0 for i in range(n)], [0 for i in range(n)]
    for i in range(n):
        if i == 0:
            z[i] = x 
            h[i] = z[i]
        if i > 0 and i < n - 1:
            z[i] = tf.matmul(h[i - 1], w[i]) + b[i]
            h[i] = tf.nn.relu(z[i])
        if i == n - 1:
            z[i] = tf.matmul(h[i - 1], w[i]) + b[i]
            h[i] = z[i]
    logits = h[n - 1]
    return logits

def save_images(images, size, path):
    img = images
    h, w = img.shape[1], img.shape[2]
    merge_img = np.zeros((h * size[0], w * size[1]))
    for idx, image in enumerate(images):
        i = idx % size[1]
        j = idx // size[1]
        merge_img[j*h:j*h+h, i*w:i*w+w] = image
    return scipy.misc.imsave(path, merge_img)

def sample_Z(m, n):
    return np.random.uniform(-1., 1., size=[m, n])

def softmax_loss(label, logits):
    return tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=label, logits=logits))

def accuracy(label, logits):
    return tf.reduce_mean(tf.cast(tf.equal(tf.argmax(logits, 1), tf.argmax(label, 1)),tf.float32))

def get_acc(x, label):
    y = classifier(x)
    return accuracy(label, y)

def avg_dist(x1, x2):
    return tf.reduce_mean(tf.reduce_sum(abs(x1 - x2), axis = 1))

def main():
    random_seed = 1024
    train_size, test_size = 55000, 10000
    batch_size = 100
    learning_rate = 0.05
    epochs = 10
    steps = epochs * 550
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = 0.8
    sess = tf.InteractiveSession(config=config)
    global_step = tf.Variable(0, name = 'global_step', trainable = False)
    tf.set_random_seed(random_seed)

    x = tf.placeholder(tf.float32, [None, 784])  # input
    label = tf.placeholder(tf.float32, [None, 10])  # true label

    x_perturb = x + tf.random_normal(shape = tf.shape(x), mean=0.0, stddev= 0.5, dtype=tf.float32)
    x_perturb = tf.clip_by_value(x_perturb, 0, 1)

    y = classifier(x)
    y_ = classifier(x_perturb)
    y_perturb = y_ + tf.random_normal(shape = tf.shape(y), mean=0.0, stddev= 300.0, dtype=tf.float32)

    #fgsm 
    x_fgsm = attack.fgsm(x, y, eps = 0.2, clip_min=0, clip_max=1)
    y_fgsm = classifier(x_fgsm)

    x_perturb2 = x + tf.random_normal(shape = tf.shape(x), mean=0.0, stddev= 0.2, dtype=tf.float32)
    y_2 = classifier(x_perturb2)
    x_fgsm_rd = attack.fgsm(x_perturb2, y_2, eps = 0.2, clip_min=0, clip_max=1)
    y_fgsm_rd = classifier(x_fgsm_rd)
    
    loss_cls = softmax_loss(label, y_perturb) 
    
    all_vars = tf.trainable_variables()
    c_vars = [var for var in all_vars if 'classifier' in var.name]
    train_op_classifier = GradientDescentOptimizer(learning_rate = learning_rate) \
        .minimize(loss_cls, var_list = c_vars, global_step = global_step)

    x1 = tf.placeholder(tf.float32, [None, 784])
    x2 = tf.placeholder(tf.float32, [None, 784])
    
    # train
    saver = tf.train.Saver()
    mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

    x_pre_fgsm_data = np.load(os.path.join('data','x_fgsm_mnist.npy'))

    sess.run(tf.global_variables_initializer())
    with sess.as_default():
        print('train classifier')
        for t in range(1, steps + 1):
            batch = mnist.train.next_batch(batch_size)
            f_dict = {x: batch[0], label: batch[1]}
            sess.run(train_op_classifier, feed_dict=f_dict)
            if t % 550 == 0:
                epoch = int(t / 550)
                acc = {} 
                dist = {}
                x_data = mnist.test.images
                label_data = mnist.test.labels

                acc['benign'] = sess.run(get_acc(x, label), feed_dict = {x: mnist.test.images, label: mnist.test.labels})
                acc['pre'] = sess.run(get_acc(x, label), feed_dict = {x: x_pre_fgsm_data, label: mnist.test.labels})
                dist['pre'] = sess.run(avg_dist(x1,x2), feed_dict = {x1: mnist.test.images, x2: x_pre_fgsm_data})

                x_fgsm_data = sess.run(x_fgsm, feed_dict = {x: mnist.test.images, label: mnist.test.labels})
                acc['fgsm'] = sess.run(get_acc(x, label), feed_dict={x: x_fgsm_data, label: mnist.test.labels})
                dist['fgsm'] = sess.run(avg_dist(x1,x2), feed_dict = {x1: mnist.test.images, x2: x_fgsm_data})

                x_perturb_data = x_data + np.random.normal(loc = 0.0, scale = 0.2, size = [10000, 784])
                x_perturb_data = np.clip(x_perturb_data, 0, 1)
                x_rd_fgsm_data = sess.run(x_fgsm, feed_dict = {x: x_perturb_data, label: mnist.test.labels})
                acc['rd'] = sess.run(get_acc(x, label), feed_dict={x: x_rd_fgsm_data, label: mnist.test.labels})
                dist['rd'] = sess.run(avg_dist(x1,x2), feed_dict = {x1: mnist.test.images, x2: x_rd_fgsm_data})

                x_it_fgsm_data = np.copy(mnist.test.images)
                for _ in range(10):
                    grad = sess.run(tf.gradients(loss_cls, x)[0], feed_dict = {x: x_it_fgsm_data, label: mnist.test.labels})
                    x_it_fgsm_data += np.sign(grad) * 0.01
                    x_it_fgsm_data = np.clip(x_it_fgsm_data, 0, 1)
                acc['it'] = sess.run(get_acc(x, label), feed_dict = {x: x_it_fgsm_data, label: mnist.test.labels})
                dist['it'] = sess.run(avg_dist(x1,x2), feed_dict = {x1: mnist.test.images, x2: x_it_fgsm_data})

                x_it_rd_fgsm_data = np.copy(mnist.test.images) + np.random.normal(loc = 0.0, scale = 0.2, size = [10000, 784])
                x_it_rd_fgsm_data = np.clip(x_it_rd_fgsm_data, 0, 1)
                for _ in range(10):
                    grad = sess.run(tf.gradients(loss_cls, x)[0], feed_dict = {x: x_it_rd_fgsm_data, label: mnist.test.labels})
                    x_it_rd_fgsm_data += np.sign(grad) * 0.01
                    x_it_rd_fgsm_data = np.clip(x_it_rd_fgsm_data, 0, 1)
                acc['it rd'] = sess.run(get_acc(x, label), feed_dict = {x: x_it_rd_fgsm_data, label: mnist.test.labels})
                dist['it rd'] = sess.run(avg_dist(x1,x2), feed_dict = {x1: mnist.test.images, x2: x_it_rd_fgsm_data})
                
                print(epoch)
                print(acc)
                print(dist)


if __name__ == "__main__":
    main()