"""
  Copyright 2020 ETH Zurich, Secure, Reliable, and Intelligent Systems Lab

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
"""


import tensorflow as tf
import numpy as np
import re
import onnx

def product(it):
    product = 1
    for x in it:
        product *= x
    return product

def runRepl(arg, repl):
    for a in repl:
        arg = arg.replace(a+"=", "'"+a+"':")
    return eval("{"+arg+"}")

def extract_mean(text):
    mean = ''
    m = re.search('mean=\[(.+?)\]', text)
    
    if m:
        means = m.group(1)
    mean_str = means.split(',')
    num_means = len(mean_str)
    mean_array = np.zeros(num_means)
    for i in range(num_means):
         mean_array[i] = np.float64(mean_str[i])
    return mean_array

def extract_std(text):
    std = ''
    m = re.search('std=\[(.+?)\]', text)
    if m:
        stds = m.group(1)
    std_str =stds.split(',')
    num_std = len(std_str)
    std_array = np.zeros(num_std)
    for i in range(num_std):
        std_array[i] = np.float64(std_str[i])
    return std_array

def numel(x):
    return product([int(i) for i in x.shape])

def parseVec(net):
    return np.array(eval(net.readline()[:-1]))

def myConst(vec):
    return tf.constant(vec.tolist(), dtype = tf.float64)

def permutation(W, h, w, c):
    m = np.zeros((h*w*c, h*w*c))
    column = 0
    for i in range(h*w):
        for j in range(c):
            m[i+j*h*w, column] = 1
            column += 1
    
    return np.matmul(W, m) 

is_tf_version_2 = tf.__version__[0]=='2'

if is_tf_version_2: 
   tf = tf.compat.v1

tf.InteractiveSession().as_default()
tf.disable_eager_execution()

def read_tensorflow_net(net_file, in_len, is_trained_with_pytorch, is_gpupoly):
    mean = 0.0
    std = 0.0
    net = open(net_file,'r')
    x = tf.placeholder(tf.float64, [in_len], name = "x")
    y = None
    z1 = None
    z2 = None
    last_layer = None
    h,w,c = None, None, None
    is_conv = False
    
    while True:
        curr_line = net.readline()[:-1]
        if 'Normalize' in curr_line:
            mean = extract_mean(curr_line)
            std = extract_std(curr_line)
        elif 'ParSum1' in curr_line:
            z1 = x
            print("par sum1")
        elif 'ParSum2' in curr_line:
            z2 = x
            x = z1
        elif 'ParSumComplete' in curr_line:
            x = tf.add(z2,x)
        elif 'ParSumReLU' in curr_line:
            x = tf.nn.relu(tf.add(z2,x))
        elif 'SkipNet1' in curr_line:
            y = x
            print("skip net1")
        elif 'SkipNet2' in curr_line:
            print("skip net2")
            #y = tf.placeholder(tf.float64, [in_len], name = "y")
            tmp = x
            x = y
            y = tmp
        elif 'SkipCat' in curr_line:
            print("skip concatenation ",x.shape[0],x.shape[1],y.shape[0],y.shape[1])
            x = tf.concat([y,x],1)
        elif curr_line in ["ReLU", "Sigmoid", "Tanh", "Sign", "Affine", "LeakyRelu"]:
            print(curr_line)
            W = None
            
            if (last_layer in ["Conv2D", "ParSumComplete", "ParSumReLU"]) and is_trained_with_pytorch and not is_gpupoly:
                W = myConst(permutation(parseVec(net), h, w, c).transpose())
            else:
                W = myConst(parseVec(net).transpose())
            b = parseVec(net)
            #b = myConst(b.reshape([1, numel(b)]))
            b = myConst(b)
            if(curr_line=="Affine"):
                x = tf.nn.bias_add(tf.matmul(tf.reshape(x, [1, numel(x)]),W), b)
            elif(curr_line=="ReLU"):
                x = tf.nn.relu(tf.nn.bias_add(tf.matmul(tf.reshape(x, [1, numel(x)]),W), b))
            elif(curr_line=="Sigmoid"):
                x = tf.nn.sigmoid(tf.nn.bias_add(tf.matmul(tf.reshape(x, [1, numel(x)]),W), b))
            elif(curr_line=="LeakyRelu"):
                x = tf.nn.leaky_relu(tf.nn.bias_add(tf.matmul(tf.reshape(x, [1, numel(x)]),W), b))
            elif(curr_line=="Sign"):
                x = tf.math.sign(tf.nn.bias_add(tf.matmul(tf.reshape(x, [1, numel(x)]),W), b))
            else:
                x = tf.nn.tanh(tf.nn.bias_add(tf.matmul(tf.reshape(x, [1, numel(x)]),W), b))
            print("\tOutShape: ", x.shape)
            print("\tWShape: ", W.shape)
            print("\tBShape: ", b.shape)
        elif curr_line == "MaxPooling2D":
            maxpool_line = net.readline()[:-1]
            if 'stride' in maxpool_line:
                args = runRepl(maxpool_line, ["input_shape" , "pool_size", "stride", "padding"])
                stride = [1] + args['stride'] + [1]
            else:
                args = runRepl(maxpool_line, ["input_shape" , "pool_size"])
                stride = [1] + args['pool_size'] + [1]
            if("padding" in maxpool_line):
                if(args["padding"]==1):
                    padding_arg = "SAME"
                else:
                    padding_arg = "VALID"
            else:
                padding_arg = "VALID"
            ksize =  [1] + args['pool_size'] + [1]
            print("MaxPool", args)

            x = tf.nn.max_pool(tf.reshape(x, [1] + args["input_shape"]), padding=padding_arg, strides=stride, ksize=ksize)
            print("\tOutShape: ", x.shape)
        elif curr_line == "Conv2D":
            is_conv = True
            line = net.readline()
            args = None
            #print(line[-10:-3])
            start = 0
            if("ReLU" in line):
                start = 5
            elif("Sigmoid" in line):
                start = 8
            elif("Tanh" in line):
                start = 5
            elif("Sign" in line):
                start = 5
            elif("Affine" in line):
                start = 7
            if 'padding' in line:
                args =  runRepl(line[start:-1], ["filters", "input_shape", "kernel_size", "stride", "padding"])
            else:
                args = runRepl(line[start:-1], ["filters", "input_shape", "kernel_size"])

            W = myConst(parseVec(net))
            print("W shape", W.shape)
            #W = myConst(permutation(parseVec(net), h, w, c).transpose())
            b = None
            if("padding" in line):
                if(args["padding"]>=1):
                    padding_arg = "SAME"
                else:
                    padding_arg = "VALID"
            else:
                padding_arg = "VALID"

            if("stride" in line):
                stride_arg = [1] + args["stride"] + [1]
            else:
                stride_arg = [1,1,1,1]

            x = tf.nn.conv2d(tf.reshape(x, [1] + args["input_shape"]), filter=W, strides=stride_arg, padding=padding_arg)

            b = myConst(parseVec(net))
            h, w, c = [int(i) for i in x.shape ][1:]
            print("Conv2D", args, "W.shape:",W.shape, "b.shape:", b.shape)
            print("\tOutShape: ", x.shape)
            if("ReLU" in line):
                x = tf.nn.relu(tf.nn.bias_add(x, b))
            elif("Sigmoid" in line):
                x = tf.nn.sigmoid(tf.nn.bias_add(x, b))
            elif("Tanh" in line):
                x = tf.nn.tanh(tf.nn.bias_add(x, b))
            elif("Affine" in line):
                x = tf.nn.bias_add(x, b)
            elif("Sign" in line):
                x = tf.math.sign(tf.nn.bias_add(x, b))
            else:
                raise Exception("Unsupported activation: ", curr_line)
        elif curr_line == "":
            break
        else:
            raise Exception("Unsupported Operation: ", curr_line)
        last_layer = curr_line

    model = x
    return model, is_conv, mean, std


def read_onnx_net(net_file):
    onnx_model = onnx.load(net_file)
    onnx.checker.check_model(onnx_model)

    is_conv = False

    for node in onnx_model.graph.node:
        if node.op_type == 'Conv':
            is_conv = True
            break

    return onnx_model, is_conv
