'''
@author: Adrian Hoffmann
'''
from .deepzono_nodes import *
from .deeppoly_nodes import *
from functools import reduce
import numpy as np


operations_for_neuron_count = ["Relu", "Sigmoid", "Tanh", "MaxPool"]


class Optimizer:
    def __init__(self, operations, resources):
        """
        Arguments
        ---------
        operations : list
            list of dicts, each dict contains a mapping from a domain (like deepzono, refinezono or deeppoly) to a tuple with resources (like matrices, biases ...)
        resources : list
            list of str, each one being a type of operation (like "MatMul", "Conv2D", "Add" ...)
        """
        self.operations = operations
        self.resources  = resources

    def get_neuron_count(self):
        total_neurons = 0
        for op, res in zip(self.operations, self.resources):
            if op in operations_for_neuron_count:
                total_neurons += np.prod(res['deepzono'][-1])
        return total_neurons

    
    
    def get_deepzono(self, nn, specLB, specUB = None):
        """
        This function will go through self.operations and self.resources and creates a list of Deepzono-Nodes which then can be run by an Analyzer object.
        It is assumed that self.resources[i]['deepzono'] holds the resources for the operation of type self.operations[i]                
        
        Arguments
        ---------
        specLB : numpy.ndarray
            1D array with the lower bound of the input spec
        specUB : numpy.ndarray
            1D array with the upper bound of the input spec
        
        Return
        ------
        execute_list : list
            list of Deepzono-Nodes that can be run by an Analyzer object
        """        
        execute_list = []
        output_info = []
        domain = 'deepzono'
        nbr_op = len(self.operations)
        
        i = 0
        while i < nbr_op:
            if self.operations[i] == "Placeholder":
                input_names, output_name, output_shape = self.resources[i][domain]
                if specUB is None:
                    execute_list.append(DeepzonoInputZonotope(specLB, input_names, output_name, output_shape))
                else:
                    execute_list.append(DeepzonoInput(specLB, specUB, input_names, output_name, output_shape))
                i += 1
            elif self.operations[i] == "MatMul":
                if i != nbr_op-1 and self.operations[i+1] in ["Add", "BiasAdd"]:
                    matrix,  m_input_names, _, _           = self.resources[i][domain]
                    bias, _, output_name, b_output_shape = self.resources[i+1][domain]
                    
                    nn.weights.append(matrix)
                    nn.biases.append(bias)
                    nn.layertypes.append('Affine')
                    nn.numlayer+= 1
                    execute_list.append(DeepzonoAffine(matrix, bias, m_input_names, output_name, b_output_shape))
                    i += 2
                else:
                    #self.resources[i][domain].append(refine)
                    execute_list.append(DeepzonoMatmul(*self.resources[i][domain]))
                    i += 1
            elif self.operations[i] == "Gemm":
                matrix, bias, m_input_names, b_output_name, b_output_shape = self.resources[i][domain]

                nn.weights.append(matrix)
                nn.biases.append(bias)
                nn.layertypes.append('Affine')
                nn.numlayer+= 1
                execute_list.append(DeepzonoAffine(matrix, bias, m_input_names, b_output_name, b_output_shape))
                i += 1
            elif self.operations[i] == "Conv2D":
                if i != nbr_op-1 and self.operations[i+1] == "BiasAdd":
                    filters, image_shape, strides, pad_top, pad_left, c_input_names, _, _ = self.resources[i][domain]
                    bias, _, b_output_name, b_output_shape = self.resources[i+1][domain]
                    nn.numfilters.append(filters.shape[3])
                    nn.filter_size.append([filters.shape[0], filters.shape[1]])
                    nn.input_shape.append([image_shape[0],image_shape[1],image_shape[2]])
                    nn.strides.append([strides[0],strides[1]])
                    nn.padding.append([pad_top, pad_left])
                    nn.out_shapes.append(b_output_shape)
                    nn.filters.append(filters)

                    nn.biases.append(bias)
                    execute_list.append(DeepzonoConvbias(image_shape, filters, bias, strides, pad_top, pad_left, c_input_names, b_output_name, b_output_shape))
                    i += 2
                else:
                    filters, image_shape, strides, pad_top, pad_left, input_names, output_name, output_shape = self.resources[i][domain]
                    nn.numfilters.append(filters.shape[3])
                    nn.filter_size.append([filters.shape[0], filters.shape[1]])
                    nn.input_shape.append([image_shape[0],image_shape[1],image_shape[2]])
                    nn.strides.append([strides[0],strides[1]])
                    nn.padding.append([pad_top, pad_left])
                    nn.out_shapes.append(output_shape)
                    nn.filters.append(filters)

                    execute_list.append(DeepzonoConv(image_shape, filters, strides, pad_top, pad_left, input_names, output_name, output_shape))
                    i += 1
                nn.layertypes.append('Conv2DNoReLU')
                nn.numlayer+=1
            elif self.operations[i] == "Conv":
                filters, bias, image_shape, strides, pad_top, pad_left, c_input_names, output_name, b_output_shape = self.resources[i][domain]
                nn.numfilters.append(filters.shape[3])
                nn.filter_size.append([filters.shape[0], filters.shape[1]])
                nn.input_shape.append([image_shape[0],image_shape[1],image_shape[2]])
                nn.strides.append([strides[0],strides[1]])
                nn.out_shapes.append(b_output_shape)
                nn.padding.append([pad_top, pad_left])
                nn.filters.append(filters)

                nn.biases.append(bias)
                nn.layertypes.append('Conv2DNoReLU')
                nn.numlayer+=1
                execute_list.append(DeepzonoConvbias(image_shape, filters, bias, strides, pad_top, pad_left, c_input_names, output_name, b_output_shape))
                i += 1
            elif self.operations[i] == "Add":
                #self.resources[i][domain].append(refine)
                execute_list.append(DeepzonoAdd(*self.resources[i][domain]))
                nn.layertypes.append('Add')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "Sub":
                #self.resources[i][domain].append(refine)
                execute_list.append(DeepzonoSub(*self.resources[i][domain]))
                nn.layertypes.append('Sub')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "Mul":
                #self.resources[i][domain].append(refine)
                execute_list.append(DeepzonoMul(*self.resources[i][domain]))
                nn.layertypes.append('Mul')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "MaxPool":
                image_shape, window_size, strides, pad_top, pad_left, input_names, output_name, output_shape = self.resources[i][domain]
                nn.pool_size.append(window_size)
                nn.input_shape.append([image_shape[0],image_shape[1],image_shape[2]])
                nn.strides.append([strides[0],strides[1]])
                nn.out_shapes.append(output_shape)
                nn.padding.append([pad_top, pad_left])
                nn.layertypes.append('MaxPooling2D')
                nn.numlayer+=1
                execute_list.append(DeepzonoMaxpool(image_shape, window_size, strides, pad_top, pad_left, input_names, output_name, output_shape))
                i += 1
            elif self.operations[i] == "Resadd":
                #self.resources[i][domain].append(refine)
                execute_list.append(DeepzonoResadd(*self.resources[i][domain]))
                nn.layertypes.append('Resaddnorelu')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "Relu":
                #self.resources[i][domain].append(refine)
                if nn.layertypes[-1]=='Affine':
                    nn.layertypes[-1] = 'ReLU'
                if nn.layertypes[-1][-6:].lower() == 'norelu':
                    nn.layertypes[-1] = nn.layertypes[-1][:-6]
                execute_list.append(DeepzonoRelu(*self.resources[i][domain]))
                i += 1
            elif self.operations[i] == "Sigmoid":
                execute_list.append(DeepzonoSigmoid(*self.resources[i][domain]))
                nn.layertypes.append('Sigmoid')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "Tanh":
                execute_list.append(DeepzonoTanh(*self.resources[i][domain]))
                nn.layertypes.append('Tanh')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "Gather":
                image_shape, indexes, axis,  input_names, output_name, output_shape = self.resources[i][domain]
                calculated_indexes = self.get_gather_indexes(image_shape, indexes, axis)
                execute_list.append(DeepzonoGather(calculated_indexes, input_names, output_name, output_shape))
                nn.layertypes.append('Gather')
                nn.numlayer += 1
                i += 1
            elif self.operations[i] == "Reshape":
                execute_list.append(DeepzonoGather(*self.resources[i][domain]))
                nn.layertypes.append('Gather')
                nn.numlayer += 1
                i += 1
            else:
                assert 0, "the optimizer for Deepzono doesn't know of the operation type " + self.operations[i]

            # for testing, getting the corresponding layer in the tensorflow or onnx model
            output_info.append(self.resources[i-1][domain][-2:])

        use_dict = self.deepzono_get_dict(execute_list)
        self.set_predecessors(nn, execute_list)
        execute_list   = self.deepzono_forward_pass(execute_list, use_dict)

        return execute_list, output_info


    def deepzono_get_dict(self, ir_list):
        """
        Returns a dict mapping output-names to the number of times that output is used by the nodes in the ir_list.
        This functions is a helper function for organizing the sections of an abstract elements when we have a ResNet or later an RNN.

        Arguments
        ---------
        ir_list : iterable
            list of Deepzono-Nodes

        Return
        ------
        use_dict : dict
            mapping from a name to the number of times that node-output is used
        """
        use_dict = {}
        for node in ir_list:
            for input_name in node.input_names:
                use_dict[input_name] += 1
            use_dict[node.output_name] = 0
        return use_dict


    def deepzono_forward_pass(self, ir_list, use_dict):
        """
        This function plans which Deepzono-Node-output occupies which section of an abstract element. If a DeepzonoDuplicate-Node should be needed, then this function will add it.
        This is needed when we have a ResNet or later RNNs.

        Arguments
        ---------
        ir_list : list
            list of Nodes, where each node has the fields output_length, input_names, and output_name (see DeepzonoNodes.py for examples)
        use_dict : dict
            maps the output_name of each node in ir_list to the number of times the node's output will be used

        Return
        ------
        ir_list : list
            the ir_list with updated and potentionally added nodes
        """
        def get_index(active_abstracts, in_name, index_store):
            index = 0
            while True:
                index = index + active_abstracts[index:].index(in_name)
                if not index in index_store:
                    break
                index += 1
            return index


        active_abstracts = []
        abstract_length  = []

        i = 0
        while i < len(ir_list):
            node        = ir_list[i]
            index_store = []
            node.abstract_information = []
            for in_name in node.input_names:
                index  = get_index(active_abstracts, in_name, index_store)
                length = abstract_length[index]
                offset = reduce(lambda x,y: x+y, abstract_length[:index], 0)
                node.abstract_information += [offset, length]
                index_store.append(index)

            if len(index_store) != 0:
                active_abstracts[index_store[0]] = node.output_name
                abstract_length[index_store[0]]  = node.output_length
                for j in range(1,len(index_store)):
                    index = index_store[j]
                    del active_abstracts[index]
                    del abstract_length[index]
            else:
                active_abstracts.append(node.output_name)
                abstract_length.append(node.output_length)
                node.abstract_information = [0, node.output_length]

            i += 1

            if use_dict[node.output_name] > 1:
                for j in range(1, use_dict[node.output_name]):
                    ir_list.insert(i, DeepzonoDuplicate(node.abstract_information[0], node.output_length))
                    active_abstracts.append(node.output_name)
                    abstract_length.append(node.output_length)
                    i += 1

        return ir_list




    def get_deeppoly(self, nn, specLB, specUB, lexpr_weights, lexpr_cst, lexpr_dim, uexpr_weights, uexpr_cst, uexpr_dim, expr_size):
        """
        This function will go through self.operations and self.resources and create a list of Deeppoly-Nodes which then can be run by an Analyzer object.
        It is assumed that self.resources[i]['deeppoly'] holds the resources for an operation of type self.operations[i].
        self.operations should only contain a combination of the following 4 basic sequences:
            - Placholder         (only at the beginning)
                - MatMul -> Add -> Relu
                - Conv2D -> Add -> Relu    (not as last layer)
                - MaxPool         (only as intermediate layer)

        Arguments
        ---------
        specLB : numpy.ndarray
            1D array with the lower bound of the input spec
        specUB : numpy.ndarray
            1D array with the upper bound of the input spec

        Return
        ------
        execute_list : list
            list of Deeppoly-Nodes that can be run by an Analyzer object
        """
        execute_list = []
        output_info = []
        domain = 'deeppoly'

        i = 0
        while i < len(self.operations):
            #print(self.operations[i])
            if self.operations[i] == "Placeholder":
                input_names, output_name, output_shape = self.resources[i][domain]
                execute_list.append(DeeppolyInput(specLB, specUB, input_names, output_name, output_shape,
                                            lexpr_weights, lexpr_cst, lexpr_dim, uexpr_weights, uexpr_cst, uexpr_dim, expr_size))
                i += 1

            # Tensorflow operation
            elif self.operations[i] == "MatMul" and self.operations[i+1] in ["Add", "BiasAdd"]:
                matrix, input_names, _,_ = self.resources[i][domain]
                bias,_, output_name, output_shape = self.resources[i+1][domain]
                if i != len(self.operations) - 2 and self.operations[i + 2] in ["Relu", "Sigmoid", "Tanh"]:
                    _,output_name, output_shape = self.resources[i+2][domain]
                nn.weights.append(matrix)
                nn.biases.append(bias)
                nn.layertypes.append('Affine')
                nn.numlayer+= 1

                if i == 1:
                    if (self.operations[i + 2] == "Relu"):
                        nn.layertypes[-1] = 'ReLU'
                        execute_list.append(DeeppolyReluNodeFirst(matrix, bias, input_names, output_name, output_shape))
                    elif (self.operations[i + 2] == "Sigmoid"):
                        execute_list.append(DeeppolySigmoidNodeFirst(matrix, bias, input_names, output_name, output_shape))
                    elif (self.operations[i + 2] == "Tanh"):
                        execute_list.append(DeeppolyTanhNodeFirst(matrix, bias, input_names, output_name, output_shape))
                    else:
                        assert 0
                    i += 3

                elif i == len(self.operations) - 2:
                    execute_list.append(DeeppolyReluNodeLast(matrix, bias, False, input_names, output_name, output_shape))
                    i += 2

                elif i == len(self.operations) - 3:
                    if(self.operations[i+2] == "Relu"):
                        nn.layertypes[-1] = 'ReLU'
                        execute_list.append(DeeppolyReluNodeLast(matrix, bias, True, input_names, output_name, output_shape))
                    elif(self.operations[i+2] == "Sigmoid"):
                        execute_list.append(DeeppolySigmoidNodeLast(matrix, bias, True, input_names, output_name, output_shape))
                    elif(self.operations[i+2] == "Tanh"):
                        execute_list.append(DeeppolyTanhNodeLast(matrix, bias, True, input_names, output_name, output_shape))
                    else:
                        assert 0, self.operations[i+2]
                    i += 3

                else:
                    if(self.operations[i+2] == "Relu"):
                        nn.layertypes[-1] = 'ReLU'
                        execute_list.append(DeeppolyReluNodeIntermediate(matrix, bias, input_names, output_name, output_shape))
                    elif(self.operations[i+2] == "Sigmoid"):
                        execute_list.append(DeeppolySigmoidNodeIntermediate(matrix, bias, input_names, output_name, output_shape))
                    elif(self.operations[i+2]=="Tanh"):
                        execute_list.append(DeeppolyTanhNodeIntermediate(matrix, bias, input_names, output_name, output_shape))
                    else:
                        assert 0
                    i += 3


            # ONNX operation
            elif self.operations[i] == "Gemm" :
                matrix, bias, input_names,output_name,output_shape = self.resources[i][domain]
                if i != len(self.operations) - 1:
                    _,output_name, output_shape = self.resources[i+1][domain]
                nn.weights.append(matrix)
                nn.biases.append(bias)
                nn.layertypes.append('Affine')
                nn.numlayer+= 1

                if i == 1:
                    if (self.operations[i + 1] == "Relu"):
                        nn.layertypes[-1] = 'ReLU'
                        execute_list.append(DeeppolyReluNodeFirst(matrix, bias, input_names, output_name, output_shape))
                    elif (self.operations[i + 1] == "Sigmoid"):
                        execute_list.append(DeeppolySigmoidNodeFirst(matrix, bias, input_names, output_name, output_shape))
                    elif (self.operations[i + 1] == "Tanh"):
                        execute_list.append(DeeppolyTanhNodeFirst(matrix, bias, input_names, output_name, output_shape))
                    else:
                        assert 0
                    i += 2

                elif i == len(self.operations) - 2:
                    if(self.operations[i+1] == "Relu"):
                        nn.layertypes[-1] = 'ReLU'
                        execute_list.append(DeeppolyReluNodeLast(matrix, bias, True, input_names, output_name, output_shape))
                    elif(self.operations[i+1] == "Sigmoid"):
                        execute_list.append(DeeppolySigmoidNodeLast(matrix, bias, True, input_names, output_name, output_shape))
                    elif(self.operations[i+1] == "Tanh"):
                        execute_list.append(DeeppolyTanhNodeLast(matrix, bias, True, input_names, output_name, output_shape))
                    else:
                        assert 0
                    i += 2

                elif i == len(self.operations) - 1:
                    execute_list.append(DeeppolyReluNodeLast(matrix, bias, False, input_names, output_name, output_shape))
                    i += 1

                else:
                    if(self.operations[i+1] == "Relu"):
                        nn.layertypes[-1] = 'ReLU'
                        execute_list.append(DeeppolyReluNodeIntermediate(matrix, bias, input_names, output_name, output_shape))
                    elif(self.operations[i+1] == "Sigmoid"):
                        execute_list.append(DeeppolySigmoidNodeIntermediate(matrix, bias, input_names, output_name, output_shape))
                    elif(self.operations[i+1]=="Tanh"):
                        execute_list.append(DeeppolyTanhNodeIntermediate(matrix, bias, input_names, output_name, output_shape))
                    else:
                        assert 0
                    i += 2

            elif self.operations[i] == "MaxPool":
                image_shape, window_size, strides, pad_top, pad_left, input_names, output_name, output_shape = self.resources[i][domain]
                nn.pool_size.append(window_size)
                nn.input_shape.append([image_shape[0],image_shape[1],image_shape[2]])
                nn.strides.append([strides[0],strides[1]])
                nn.out_shapes.append(output_shape)
                nn.padding.append([pad_top, pad_left])
                nn.layertypes.append('MaxPooling2D')
                nn.numlayer+=1
                execute_list.append(DeeppolyMaxpool(image_shape, window_size, strides, input_names, output_name, output_shape))
                i += 1

                # Tensorflow operation
            elif self.operations[i] == "Conv2D" and self.operations[i + 1] == "BiasAdd":

                filters, image_shape, strides, pad_top, pad_left, input_names, _, _ = self.resources[i][domain]
                bias, _, output_name, output_shape = self.resources[i + 1][domain]
                nn.numfilters.append(filters.shape[3])
                nn.filter_size.append([filters.shape[0], filters.shape[1]])
                nn.input_shape.append([image_shape[0], image_shape[1], image_shape[2]])
                nn.strides.append([strides[0], strides[1]])
                nn.padding.append([pad_top, pad_left])
                nn.out_shapes.append(output_shape)
                nn.filters.append(filters)

                nn.biases.append(bias)
                has_relu = self.operations[i + 2] == "Relu"
                if has_relu:
                    nn.layertypes.append('Conv2D')
                    _, output_name, output_shape = self.resources[i + 2][domain]
                    increment = 3
                else:
                    nn.layertypes.append('Conv2DNoReLU')
                    increment = 2
                nn.numlayer += 1
                if i == 1:
                    execute_list.append(
                        DeeppolyConv2dNodeFirst(filters, strides, pad_top, pad_left, bias, image_shape, input_names, output_name,
                                                output_shape, has_relu))
                else:
                    execute_list.append(DeeppolyConv2dNodeIntermediate(filters, strides, pad_top, pad_left, bias, image_shape, input_names,
                                                             output_name, output_shape, has_relu))
                i += increment

            # ONNX operation
            elif self.operations[i] == "Conv":
                filters, bias, image_shape, strides, pad_top, pad_left, input_names,output_name,out_shape = self.resources[i][domain]
                has_relu = self.operations[i+1] == "Relu"
                if has_relu:
                    nn.layertypes.append('Conv2D')
                    _,output_name,output_shape = self.resources[i+1][domain]
                    increment = 2
                else:
                    nn.layertypes.append('Conv2DNoReLU')
                    increment = 1
                nn.numfilters.append(filters.shape[3])
                nn.filter_size.append([filters.shape[0], filters.shape[1]])
                nn.input_shape.append([image_shape[0],image_shape[1],image_shape[2]])
                nn.strides.append([strides[0],strides[1]])
                nn.padding.append([pad_top, pad_left])
                nn.out_shapes.append(out_shape)
                nn.filters.append(filters)

                nn.biases.append(bias)
                nn.numlayer+=1
                if i == 1:
                    execute_list.append(DeeppolyConv2dNodeFirst(filters, strides, pad_top, pad_left, bias, image_shape, input_names, output_name, output_shape, has_relu))
                else:
                    execute_list.append(DeeppolyConv2dNodeIntermediate(filters, strides, pad_top, pad_left, bias, image_shape, input_names, output_name, output_shape, has_relu))
                i += increment

            # Residual layer
            elif self.operations[i] == "Resadd":
                if self.operations[i+1] == "Relu":
                    input_names,_,_ = self.resources[i][domain]
                    _,output_name,output_shape = self.resources[i+1][domain]
                    execute_list.append(DeeppolyResadd(input_names,output_name,output_shape, True))
                    nn.layertypes.append('Resadd')
                    i += 2
                else:
                    input_names,output_name,output_shape = self.resources[i][domain]
                    execute_list.append(DeeppolyResadd(input_names,output_name,output_shape, False))
                    nn.layertypes.append('Resaddnorelu')
                    i += 1
                nn.numlayer+=1

            # Gather
            elif self.operations[i] == "Gather":
                image_shape, indexes, axis, input_names,output_name,output_shape = self.resources[i][domain]
                calculated_indexes = self.get_gather_indexes(image_shape, indexes, axis)
                execute_list.append(DeeppolyGather(calculated_indexes, input_names,output_name,output_shape))
                i+=1

            elif self.operations[i] == "Reshape":
                indexes, input_names, output_name, output_shape = self.resources[i][domain]
                execute_list.append(DeeppolyGather(indexes, [input_names[0]], output_name, output_shape))
                i+=1

            elif self.operations[i] == "Sub":
                if i == 1:
                    execute_list.append(DeeppolySubNodeFirst(*self.resources[i][domain]))
                else:
                    execute_list.append(DeeppolySubNodeIntermediate(*self.resources[i][domain]))
                nn.layertypes.append('Sub')
                nn.numlayer += 1
                i += 1

            elif self.operations[i] == "Mul":
                if i == 1:
                    execute_list.append(DeeppolyMulNodeFirst(*self.resources[i][domain]))
                else:
                    execute_list.append(DeeppolyMulNodeIntermediate(*self.resources[i][domain]))
                nn.layertypes.append('Mul')
                nn.numlayer += 1
                i += 1
            else:
                assert 0, "the Deeppoly analyzer doesn't support the operation: '" + self.operations[i] + "' of this network: " + str(self.operations)

            # for testing, getting the corresponding layer in the tensorflow or onnx model
            output_info.append(self.resources[i-1][domain][-2:])

        self.set_predecessors(nn, execute_list)
        return execute_list, output_info

    def set_predecessors(self, nn, output):
        output_index_store = {}
        index_o = 0
        for node in output:
            if isinstance(node, DeepzonoRelu):
                output_index_store[node.output_name] = index_o - 1
            else:
                output_index_store[node.output_name] = index_o
                index_o += 1
        for node in output:
            predecessors = (c_size_t * len(node.input_names))()
            i = 0
            for input_name in node.input_names:
                predecessors[i] = output_index_store[input_name]
                i += 1
            node.predecessors = predecessors
            if not isinstance(node, DeepzonoRelu):
                nn.predecessors.append(predecessors)

    def get_gather_indexes(self, input_shape, indexes, axis):
        size = np.prod(input_shape)
        base_indexes = np.arange(size).reshape(input_shape)
        return np.take(base_indexes, indexes, axis=axis)