#----------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License. See License.txt in the project root for license information.
#----------------------------------------------------------------------------------------------

import os
import numpy as np
from six.moves import xrange
import cntk as _cntk
from mmdnn.conversion.cntk.cntk_graph import CntkGraph
import mmdnn.conversion.common.IR.graph_pb2 as graph_pb2
from mmdnn.conversion.common.IR.graph_pb2 import NodeDef, GraphDef, DataType
from mmdnn.conversion.common.utils import *
from mmdnn.conversion.common.DataStructure.parser import Parser


class CntkParser(Parser):

    dtype_map = {
        0 : graph_pb2.DT_UNDEFINED,
        np.float32 : graph_pb2.DT_FLOAT32,
        np.float64 : graph_pb2.DT_FLOAT64,
        3 : graph_pb2.DT_INT32,
        4 : graph_pb2.DT_UINT8,
        5 : graph_pb2.DT_INT16,
        6 : graph_pb2.DT_INT8,
        7 : graph_pb2.DT_STRING,
        9 : graph_pb2.DT_INT64
    }


    @property
    def src_graph(self):
        return self.cntk_graph


    def __init__(self, model, dest_nodes = None):
        super(CntkParser, self).__init__()

        if not os.path.exists(model):
            raise ValueError('Cntk model [{}] can not be found!'.format(model))
        model = _cntk.Function.load(model)
        self.weight_loaded = True

        # Build network graph
        self.cntk_graph = CntkGraph(model)
        self.cntk_graph.build()


    @staticmethod
    def _convert_padding_to_IR(kernel_shape, auto_pad):
        lower = []
        upper = []
        for idx in range(0, len(kernel_shape)):
            if auto_pad[idx] == False:
                lower += [0]
                upper += [0]
            else:
                q = kernel_shape[idx] // 2
                lower += [q] if kernel_shape[idx] % 2 else [q - 1]
                upper += [q]

        return [0] + lower + [0, 0] + upper + [0]


    def _convert_identity_operation(self, source_node, start_edge=0, end_edge=None, new_op=None, shape_transpose=True):
        IR_node = self.IR_graph.node.add()
        CntkParser._copy_and_reop(source_node, IR_node, new_op, shape_transpose)
        self.convert_inedge(source_node, IR_node, start_edge, end_edge)
        return IR_node


    def gen_IR(self):
        for layer in self.src_graph.topological_sort:
            current_node = self.src_graph.get_node(layer)
            node_type = current_node.type
            if hasattr(self, "rename_" + node_type):
                func = getattr(self, "rename_" + node_type)
                func(current_node)
            else:
                self.rename_UNKNOWN(current_node)


    @staticmethod
    def _copy_and_reop(source_node, IR_node, new_op=None, shape_transpose=False):
        if new_op == None: new_op = source_node.type
        IR_node.name = source_node.real_name
        IR_node.op = new_op

        kwargs = {}

        if hasattr(source_node.layer, 'dtype'):
            assert source_node.layer.dtype in CntkParser.dtype_map, 'type [{}] is unknown.'.format(source_node.layer.dtype)
            IR_node.attr["dtype"].type = CntkParser.dtype_map[source_node.layer.dtype]

        if hasattr(source_node.layer, 'shape'):
            shape =  (-1,) + source_node.layer.shape
            if shape_transpose:
                shape = CntkParser.channel_first_shape_to_IR(shape)
            shape = list_to_shape(shape)
            kwargs['_output_shapes'] = [shape]

        assign_IRnode_values(IR_node, kwargs)


    def _fuse_bias_node(self, source_node):
        next_node = self.src_graph.get_son(source_node.name, [0])
        if next_node is None or next_node.type != 'Plus' or not next_node.layer.parameters:
            return False

        next_node.covered = True
        next_node.real_name = source_node.real_name
        B = next_node.layer.parameters[0].asarray()
        self.set_weight(source_node.name, 'bias', B)

        return True


    @staticmethod
    def _print_layer(source_node):
        print ("Layer: ", source_node.layer)
        print ("Parameters: ", source_node.layer.parameters)
        print ("Attributes: ", source_node.layer.attributes)
        for in_node in source_node.layer.inputs:
            print (in_node)


    def rename_UNKNOWN(self, source_node):
        print("Cntk Parser has not supported operator [%s] with name [%s]."
              % (source_node.type, source_node.name))

        self._print_layer(source_node)
        assert False


    @staticmethod
    def get_ndarray(variable):
        if variable.is_parameter:
            return variable.as_parameter().asarray()

        elif variable.is_constant:
            return variable.as_constant().asarray()

        else:
            raise ValueError("Unknown variable [{}].".format(variable))


    @staticmethod
    def _get_attribute(source_node, attribute_name):
        if attribute_name in source_node.attributes:
            return source_node.attributes

        node = source_node.block_root
        while not attribute_name in node.attributes:
            node = node.inputs[0].owner

        return node.attributes


    def rename_Convolution(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op="Conv")

        for input in source_node.layer.inputs:
            if input.name.endswith("W"):
                W = self.get_ndarray(input)
                break

        W = self.channel_first_conv_kernel_to_IR(W)
        self.set_weight(source_node.name, 'weights', W)

        attributes = CntkParser._get_attribute(source_node.layer, 'strides')

        kwargs = dict()
        kwargs['strides'] = [1] + list(attributes['strides'])[1:] + [1]
        kwargs['dilations'] = [1] + list(attributes['dilation'])[1:] + [1]
        kwargs['kernel_shape'] = list(W.shape)
        padding = attributes['autoPadding'][1:]

        for pad in padding:
            assert pad == padding[0]

        kwargs['auto_pad'] = 'SAME_LOWER' if padding[0] else 'VALID'
        kwargs['pads'] = self._convert_padding_to_IR(kwargs['kernel_shape'][:-2], padding)

        kwargs['use_bias'] = self._fuse_bias_node(source_node)

        assign_IRnode_values(IR_node, kwargs)


    def rename_ReLU(self, source_node):
        self._convert_identity_operation(source_node, new_op='Relu')


    def rename_Relu6(self, source_node):
        self._convert_identity_operation(source_node)


    def rename_Plus(self, source_node):
        if not source_node.covered:
            assert not source_node.layer.parameters
            IR_node = self._convert_identity_operation(source_node, new_op='Add')


    def rename_Minus(self, source_node):
        if not source_node.covered:
            assert not source_node.layer.parameters
            self._convert_binary_operator(source_node, new_op='Sub')


    def rename_Sub(self, source_node):
        self._convert_identity_operation(source_node)


    def rename_Reshape(self, source_node):
        IR_node = self._convert_identity_operation(source_node)
        new_shape = source_node.get_attr('newShape')
        kwargs = {'shape' : self.channel_first_shape_to_IR(new_shape)}
        assign_IRnode_values(IR_node, kwargs)


    def rename_Times(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op='FullyConnected')

        W = source_node.layer.parameters[0].asarray().squeeze()
        self.set_weight(source_node.name, 'weights', W)

        kwargs = dict()
        kwargs['units'] = W.shape[-1]
        kwargs['use_bias'] = self._fuse_bias_node(source_node)
        assign_IRnode_values(IR_node, kwargs)


    def rename_MaxPooling(self, source_node):
        if source_node.layer.is_block:
            source_node.layer = source_node.layer.block_root.owner

        self.rename_Pooling(source_node)


    def rename_AveragePooling(self, source_node):
        self.rename_Pooling(source_node)

    def rename_Slice(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op = 'Slice')
        kwargs = dict()
        kwargs['axis'] = source_node.get_attr('axis')[-1] + 1
        kwargs['ends'] = source_node.get_attr('endIndex')
        kwargs['starts'] = source_node.get_attr('beginIndex')
        kwargs['strides'] = source_node.get_attr('sliceStrides')
        assign_IRnode_values(IR_node, kwargs)


    def rename_Splice(self, source_node):
        if len(source_node.in_edges) == 1:
            source_node.in_edges.append(source_node.in_edges[0])
        IR_node = self._convert_identity_operation(source_node, new_op='Concat')
        assign_IRnode_values(IR_node, {'axis' : source_node.get_attr('axis')[-1] + 1})


    def rename_StableSigmoid(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op='Sigmoid')

    def rename_BinaryCrossEntropy(self, source_node):
        pass

    def rename_Pooling(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op='Pool')
        dim = len(IR_node.attr['_output_shapes'].list.shape[0].dim)
        kwargs = {}

        # strides
        kwargs['strides'] = list(source_node.get_attr('strides')) + [1]
        if len(kwargs['strides']) < dim:
            kwargs['strides'] = [1] + kwargs['strides']

        # window_shape
        kwargs['kernel_shape'] = list(source_node.get_attr('poolingWindowShape')) + [1]
        if len(kwargs['kernel_shape']) < dim:
            kwargs['kernel_shape'] = [1] + kwargs['kernel_shape']

        # pool type
        pool_type = source_node.get_attr('poolingType')
        if pool_type == _cntk.MAX_POOLING:
            kwargs['pooling_type'] = 'MAX'
        elif pool_type == _cntk.AVG_POOLING:
            kwargs['pooling_type'] = 'AVG'
        else:
            raise ValueError("Unknown pooling type [{}].".format(pool_type))

        # padding
        padding = source_node.get_attr('autoPadding')
        if len(padding) >= dim - 1:
            padding = padding[1:]
        elif len(padding) < dim - 2:
            padding.extend([padding[-1]] * (dim - len(padding) - 2))
        for pad in padding:
            assert pad == padding[-1]
        kwargs['auto_pad'] = 'SAME_LOWER' if padding[0] else 'VALID'
        kwargs['pads'] = self._convert_padding_to_IR(kwargs['kernel_shape'][1:-1], padding)

        assign_IRnode_values(IR_node, kwargs)


    def rename_DataInput(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op='DataInput')
        shape = [-1] + list(source_node.layer.shape)
        assign_IRnode_values(IR_node, {'shape' : list_to_shape(self.channel_first_shape_to_IR(shape))})


    def rename_BatchNormalization(self, source_node):
        kwargs = dict()
        kwargs['scale'] = False
        kwargs['bias'] = False
        for param in source_node.layer.inputs:
            if param.name.endswith('scale'):
                self.set_weight(source_node.name, 'scale', self.get_ndarray(param).flatten())
                kwargs['scale'] = True

            elif param.name.endswith('bias'):
                self.set_weight(source_node.name, 'bias', self.get_ndarray(param).flatten())
                kwargs['bias'] = True

            elif param.name.lower().endswith('mean'):
                self.set_weight(source_node.name, 'mean', self.get_ndarray(param).flatten())

            elif param.name.lower().endswith('variance'):
                self.set_weight(source_node.name, 'var', self.get_ndarray(param).flatten())

        IR_node = self._convert_identity_operation(source_node, end_edge=1, new_op='BatchNorm')
        kwargs['epsilon'] = source_node.get_attr('epsilon')
        kwargs['axis'] = -1
        assign_IRnode_values(IR_node, kwargs)


    def _add_constant_node(self, constant_node, IR_node):
        new_node = self.IR_graph.node.add()
        new_node.name = constant_node.uid
        new_node.op = 'Constant'
        value = np.atleast_1d(self.get_ndarray(constant_node))
        self.set_weight(new_node.name, 'value', value)
        IR_node.input.append(new_node.name)


    def _convert_binary_operator(self, source_node, new_op):
        IR_node = self._convert_identity_operation(source_node, new_op=new_op)
        for in_node in source_node.layer.inputs:
            if in_node.is_constant:
                self._add_constant_node(in_node, IR_node)


    def rename_ElementTimes(self, source_node):
        if source_node.layer.inputs[0] == source_node.layer.inputs[1]:
            # TODO: Handle square
            pass

        self._convert_binary_operator(source_node, 'Mul')


    def rename_Log(self, source_node):
        self._convert_identity_operation(source_node)


    def rename_Exp(self, source_node):
        self._convert_identity_operation(source_node)


    def rename_Reciprocal(self, source_node):
        self._convert_identity_operation(source_node)


    def rename_Dropout(self, source_node):
        # self._print_layer(source_node)
        # print (source_node.name)
        # print (self.src_graph.get_parent(source_node.name, [0]).real_name)
        # assert False
        source_node.real_name = self.src_graph.get_parent(source_node.name, [0]).real_name


    def rename_Dense(self, source_node):
        IR_node = self._convert_identity_operation(source_node, new_op='FullyConnected', shape_transpose=False)
        for param in source_node.layer.inputs:
            if param.name.endswith('W'):
                w = np.squeeze(self.get_ndarray(param))
                if w.ndim > 2:
                    w = np.transpose(w, list(range(1, w.ndim - 1)) + [0, -1])
                    w = np.reshape(w, [-1, w.shape[-1]])
                self.set_weight(source_node.name, 'weights', w)
                assign_IRnode_values(IR_node, {'units' : w.shape[-1] })

            elif param.name.endswith('b'):
                self.set_weight(source_node.name, 'bias', self.get_ndarray(param))
                assign_IRnode_values(IR_node, {'use_bias' : True })


    def rename_Convolution2D(self, source_node):
        assert source_node.layer.is_block

        # Convolution
        kwargs = dict()
        conv_IR_node = self.IR_graph.node.add()
        conv_node = source_node.layer.block_root.inputs[0].owner.inputs[0].owner

        conv_IR_node.name = conv_node.uid
        conv_IR_node.op = 'Conv'
        conv_IR_node.input.append(self.get_parent(source_node.name, [0]).real_name)

        # Kernel
        conv_weight = source_node.layer.block_root.inputs[0].owner.inputs[0].owner.inputs[0]
        conv_weight = self.get_ndarray(conv_weight)
        W = self.channel_first_conv_kernel_to_IR(conv_weight)
        self.set_weight(conv_IR_node.name, 'weights', W)

        # Attributes
        conv_attr = source_node.layer.block_root.inputs[0].owner.inputs[0].owner.attributes

        kwargs['strides'] = [1] + list(conv_attr['strides'])[1:] + [1]
        kwargs['dilations'] = [1] + list(conv_attr['dilation'])[1:] + [1]
        kwargs['kernel_shape'] = list(W.shape)
        padding = conv_attr['autoPadding'][1:]

        for pad in padding:
            assert pad == padding[0]

        kwargs['auto_pad'] = 'SAME_LOWER' if padding[0] else 'VALID'
        kwargs['pads'] = self._convert_padding_to_IR(kwargs['kernel_shape'][:-2], padding)

        kwargs['use_bias'] = True

        assign_IRnode_values(conv_IR_node, kwargs)

        # Bias
        plus = source_node.layer.block_root.inputs[0].owner.inputs[1]
        plus = np.squeeze(self.get_ndarray(plus))
        self.set_weight(conv_IR_node.name, 'bias', plus)

        # Activation
        activation = source_node.layer.block_root.owner.op_name

        activation_IR = self.IR_graph.node.add()
        activation_IR.name = source_node.name
        activation_IR.input.append(conv_IR_node.name)
        if (activation == 'ReLU'):
            activation_IR.op = 'Relu'
        else:
            raise ValueError()


    def rename_Activation(self, source_node):
        assert source_node.layer.is_block

        op = source_node.layer.root_function.owner.name

        if op.startswith('relu'):
            new_op= 'Relu'
        else:
            raise ValueError()

        self._convert_identity_operation(source_node, new_op=new_op)