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

from keras.layers import BatchNormalization

from python.slalom.mobileNetv2 import MobileNetV2 as MobileNetV2_sgx
from python.slalom.global_sgx import sgxutils
from python.slalom.Activation import Activation


import keras.layers as layers

def main(args):
    config = tf.ConfigProto(log_device_placement=False)
    config.allow_soft_placement = True
    config.gpu_options.per_process_gpu_memory_fraction = 0.90
    config.gpu_options.allow_growth = True
    h = w = 224
    c = 8
    input_shape = (15, h, w, c)
    bm, um, gm, igm = sgxutils.fill_parameter(internal_batch_size=3)

    with tf.Session(config=config) as sess:
        input_tensor = tf.random_uniform(input_shape, minval=-10000.0, maxval=1000.0, dtype=tf.float32)

        with tf.device("/gpu:0"):
            print("======== BEGIN TO BUILD TARGET ========")

            x1 = layers.Input(shape=(h, w, c))
            y1 = layers.Input(shape=(h, w, c))

            y_tar = Activation(axis=-1,
                             privacy=False,
                             act_mode="bnrelu",
                             sgxutils=sgxutils,
                             epsilon=0.0,
                             momentum=1.0,
                             name="target")(x1, skip_input=y1, training=True)
            print("======== BEGIN TO BUILD ENCLAV ========")
            x2 = layers.Input(shape=(h, w, c))
            y2 = layers.Input(shape=(h, w, c))

            y_sgx = Activation(
                               axis=-1,
                               act_mode="bnrelu",
                               privacy=True,
                               epsilon=0.0,
                               momentum=1.0,
                               sgxutils=sgxutils,
                               name="sfs")(x2, skip_input=y2)
            print("end of building layers")
        # beta & gamma uninitialized
        K.get_session().run(tf.initialize_all_variables())

        input_np1 = sess.run(input_tensor, feed_dict={})
        input_np2 = sess.run(input_tensor, feed_dict={})
        input_np1 = sess.run(input_tensor, feed_dict={})
        input_np2 = sess.run(input_tensor, feed_dict={})
        print("======== EVALUATING GPU ========")
        y_tar_val = sess.run(y_tar, feed_dict={x1:input_np1, y1:input_np2})
        
        print("======== EVALUATING SGX ========")
        y_sgx_val = sess.run(y_sgx, feed_dict={x2:input_np1, y2:input_np2})

        
        print((y_tar_val[0] - y_sgx_val[0]).max())
        print((y_tar_val[1] - y_sgx_val[1]).max())
        print((y_tar_val[4] - y_sgx_val[4]).max())
        print((y_tar_val[3] - y_sgx_val[3]).max())

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    args = parser.parse_args()

    main(args)