from __future__ import print_function
import numpy as np
import struct  # convert from Python values and C structs
import tensorflow as tf
import re

# import barracuda
# from barracuda import Struct
from mlagents.trainers import barracuda
from mlagents.trainers.barracuda import Struct
from google.protobuf import descriptor
from google.protobuf.json_format import MessageToJson


if __name__ == "__main__":
    # Handle command line argumengts
    args = barracuda.parse_args(
        description="Convert Tensorflow model to Barracuda binary",
        source_extension=".pb",
        help="input Tensorflow serialized .pb file",
    )
    # Te following code can be used as an example of API used from another module
    # convert() is the main entry point for converter
    import tensorflow_to_barracuda as tf2bc

    tf2bc.convert(args.source_file, args.target_file, args.trim_unused_by_output, args)


# TODO: support more than 1 LSTM layer per model - prepend scope to names and inputs
# TODO: support different activation functions in LSTM
# TODO: strip output Identity node, instead patch upstream layer names
# TODO: use ScaleBias and Pow with alpha when input is constant Tensor
# TODO: support all data format types (curretly only NHWC)
# TODO: support all data types (currently only FLOAT, INT32, BOOL)
# TODO: implement FusedResizeAndPadConv2D

# Important ProtoBuf definitions:
#    https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/types.proto
#    https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/tensor.proto
#    https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/node_def.proto
#
# Node descriptions:
#    https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/ops/nn_ops.cc
#    https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/ops/math_ops.cc
#    https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/ops/random_ops.cc
#
# Class doc:
#    https://www.tensorflow.org/api_docs/cc/
#
known_classes = {
    "Dense": Struct(
        id=1,
        out_shapes=lambda shapes: [
            [shapes[0][0], 1, 1, shapes[0][1]],  # W
            [1, 1, 1, shapes[-1][-1]],  # B
        ],
        patch_data=lambda data: [data[0], data[1]],
    ),
    "MatMul": Struct(
        id=1,
        out_shapes=lambda shapes: [
            [shapes[0][0], 1, 1, shapes[0][1]],  # W
            [1, 1, 1, shapes[0][1]],  # B
        ],
        patch_data=lambda data: [data[0], np.zeros(np.shape(data[1]))],
    ),
    "BiasAdd": Struct(
        id=51,  # implemented as ScaleBias
        out_shapes=lambda shapes: [
            [1, 1, 1, shapes[0][0]],  # ONE
            [1, 1, 1, shapes[0][0]],  # B
        ],
        patch_data=lambda data: [np.ones(np.shape(data[0])), data[0]],
    ),
    # TODO: NCHW
    "Conv2D": Struct(
        id=20,
        out_shapes=lambda shapes: [shapes[0], [1, 1, 1, shapes[-1][-1]]],  # K  # B
        patch_data=lambda data: [data[0], data[1]],
    ),
    "DepthwiseConv2dNative": Struct(  # DepthwiseConv2D
        id=21,
        out_shapes=lambda s: [
            [
                s[0][0],
                s[0][1],
                s[0][3],
                s[0][2],
            ],  # K TF:[H, W, in_channels, channel_multiplier] => [H, W, 1, in_channels]
            [1, 1, 1, s[-1][-1]] if len(s) > 1 else [1, 1, 1, s[0][2]],  # B
        ],
        patch_data=lambda data: [np.transpose(data[0], (0, 1, 3, 2)), data[1]],
    ),
    "Conv2DBackpropInput": Struct(  # Conv2DTranspose
        id=22,
        out_shapes=lambda shapes: [shapes[0], [1, 1, 1, shapes[-1][-1]]],  # K  # B
        patch_data=lambda data: [data[0], data[1]],
    ),
    # TODO: 3D
    "ResizeNearestNeighbor": 23,  # implemented as Upsample2D
    "ResizeBilinear": 23,  # implemented as Upsample2D
    "ResizeBicubic": 23,  # implemented as Upsample2D
    "MaxPool": 25,
    "AvgPool": 26,
    "GlobalAveragePool": 28,
    "Activation": 50,
    "BatchNormalization": Struct(
        id=51,  # after fusion implemented as ScaleBias
        out_shapes=lambda shapes: [
            [1, 1, 1, shapes[0][0]],  # S
            [1, 1, 1, shapes[0][0]],  # B
        ],
        patch_data=lambda data:
        # fuse [gamma, beta, mean, var, epsilon] => [scale, bias]
        # TODO: double-check if epsilon is the last data argument and not the 1st?
        barracuda.fuse_batchnorm_weights(data[0], data[1], data[2], data[3], data[4])
        if len(data) == 5
        else
        # fuse [ONE, beta, mean, var, epsilon] => [scale, bias]
        # TODO: double-check if epsilon is the last data argument and not the 1st?
        barracuda.fuse_batchnorm_weights(
            np.ones(np.shape(data[0])), data[0], data[1], data[2], data[3]
        ),
    ),
    "FusedBatchNorm": Struct(
        id=51,  # after fusion implemented as ScaleBias
        out_shapes=lambda shapes: [
            [1, 1, 1, shapes[0][0]],  # S
            [1, 1, 1, shapes[0][0]],  # B
        ],
        patch_data=lambda data, layer:
        # fuse [gamma, beta, mean, var, epsilon] => [scale, bias]
        barracuda.fuse_batchnorm_weights(
            data[0], data[1], data[2], data[3], get_epsilon(layer)
        ),
    ),
    "LRN": 53,
    "RandomStandardNormal": 64,
    "RandomUniform": 65,
    "Multinomial": 66,
    "OneHot": 67,
    # Broadcast ops
    "Add": 100,
    "Sub": 101,
    "Mul": 102,
    "RealDiv": 103,
    "Pow": 104,
    "Minimum": 110,
    "Maximum": 111,
    # Reduce ops
    "Max": 124,
    "Mean": 125,
    "Min": 126,
    "Prod": 127,
    "Sum": 128,
    "Flatten": 200,
    "Reshape": 201,
    #'Squeeze':          203,
    #'Unsqueeze':        204,
    "Concat": 210,
    "StridedSlice": 211,
}

requires_runtime_flag = {
    "Dropout": "DropoutRuntime",
    "BatchNormalization": "BatchNormalizationRuntime",
}

known_activations = {
    "Linear": 0,
    "Relu": 1,
    "Softmax": 2,
    "Tanh": 3,
    "Sigmoid": 4,
    "Elu": 5,
    "Relu6": 6,
    "LeakyRelu": 7,
    "Selu": 8,
    "Swish": 9,
    "LogSoftmax": 10,
    "Softplus": 11,
    "Softsign": 12,
    "Abs": 100,
    "Neg": 101,
    "Ceil": 102,
    "Floor": 104,
    "Sqrt": 111,
    "Exp": 113,
    "Log": 114,
    "Acos": 200,
    "Acosh": 201,
    "Asin": 202,
    "Asinh": 203,
    "Atan": 204,
    "Atanh": 205,
    "Cos": 206,
    "Cosh": 207,
    "Sin": 208,
    "Sinh": 209,
    "Tan": 210,
}

known_paddings = {"VALID": [0, 0, 0, 0], "SAME": [-1]}  # SameUpper

supported_data_formats = {"NHWC"}

known_patterns = {
    # TODO: Flatten pattern using namespace regexp
    repr(["Shape", "StridedSlice", "Pack", "Reshape"]): "Flatten",
    repr(["Shape", "StridedSlice", "Prod", "Pack", "Reshape"]): "Flatten",
    repr(
        ["Shape", "Slice", "Slice", "Prod", "ExpandDims", "ConcatV2", "Reshape"]
    ): "Flatten",
    repr(["Const", "Reshape"]): "Reshape",
    repr(["Add", "Rsqrt", "Mul", "Mul", "Sub", "Add"]): "BatchNormalization",
    repr(["Add", "Rsqrt", "Mul", "Mul", "Mul", "Sub", "Add"]): "BatchNormalization",
    repr(["ConcatV2"]): "ConcatV2",
    repr(["Mean"]): "Mean",
    repr(["Multinomial"]): "Multinomial",
    repr(["OneHot"]): "OneHot",
    repr(["Square"]): "Square",
    repr(["MatMul", "BiasAdd"]): "Dense",
    repr(["Conv2D", "BiasAdd"]): "Conv2D",
    repr(["DepthwiseConv2dNative", "BiasAdd"]): "DepthwiseConv2dNative",
    repr(["Conv2DBackpropInput", "BiasAdd"]): "Conv2DBackpropInput",
    repr(
        ["Pack", "Reshape"]
    ): "Flatten$",  # for now we assume that this combination is trivial Flatten
    # for exmaple it is used in ML-agents LSTM nets with sequence_length==1
    repr(
        [
            "StridedSlice",
            "Reshape",
            re.compile("^lstm/"),
            "Reshape",
            "ConcatV2",
            "Identity",
        ]
    ): "BasicLSTM",
    repr([re.compile("^lstm/"), "Reshape", "ConcatV2", "Identity"]): "BasicLSTM",
    repr(["Sigmoid", "Mul"]): "Swish",
    # TODO: FusedResizeAndPadConv2D
}


def by_name(args, name):
    for a in args:
        if a.name.endswith(name):
            return a


def by_op(args, op):
    for a in args:
        if a.op == op:
            return a


def order_by(args, names):
    ordered = []
    arg_count = len(args)
    for name in names:
        ordered += [a for a in args if a.endswith(name)]
        args = [a for a in args if not a.endswith(name)]
    ordered += args  # append what is left
    assert len(ordered) == arg_count
    return ordered


transform_patterns = {
    "Flatten": lambda nodes, inputs, tensors, _: Struct(op="Flatten", input=inputs),
    "Flatten$": lambda nodes, inputs, tensors, _: Struct(
        op="Flatten",
        input=[
            inputs[-1]
        ],  # take only the last input, assume all other arguments are trivial (like sequence_length==1 always in ML-agents LSTM nets)
    ),
    "Reshape": lambda nodes, inputs, tensors, _: Struct(
        op="Reshape",
        input=inputs,
        shape=[
            tensors[0].data[0],
            tensors[0].data[1],
            tensors[0].data[2],
            tensors[0].data[3],
        ]
        if len(tensors[0].data) == 4
        else [tensors[0].data[0], 1, tensors[0].data[1], tensors[0].data[2]]
        if len(tensors[0].data) == 3
        else [tensors[0].data[0], 1, 1, tensors[0].data[1]]
        # tensor.name = 'shape'
    ),
    "Multinomial": lambda nodes, inputs, tensors, _: Struct(
        op="Multinomial",
        input=inputs,
        shape=[int(by_name(tensors, "/num_samples").data[0])],
        # seed = get_attr(nodes[0], 'seed'),
    ),
    "OneHot": lambda nodes, inputs, tensors, _: Struct(
        op="OneHot",
        input=inputs,
        shape=[int(by_name(tensors, "/depth").data[0])],
        alpha=by_name(tensors, "/on_value").data[0],
        beta=by_name(tensors, "/off_value").data[0],
    ),
    "Square": lambda nodes, inputs, tensors, _: Struct(
        op="Mul", input=[i for i in inputs] + [i for i in inputs]  # input * input
    ),
    "ConcatV2": lambda nodes, inputs, tensors, _: Struct(
        op="Concat",
        input=inputs,
        # TEMPORARY: until we implemented rank detection and axis remapping (hopefully in exporter)
        # HACK: assume Concat is always for last channel
        axis=int(-1)
        # axis = int(by_name(tensors, '/axis').data[0])
    ),
    "BatchNormalization": lambda nodes, inputs, tensors, _: Struct(
        op="BatchNormalization",
        input=[i for i in inputs]
        + order_by([t.name for t in tensors], ["gamma", "beta", "mean", "variance"]),
    ),
    "Mean": lambda nodes, inputs, tensors, _: Struct(
        # TODO: use data_frmt of the input instead of hardcoded [1,2] for HW
        op="GlobalAveragePool"
        if np.array_equal(tensors[0].data, [1, 2])
        else "MeanWithUnsupportedReductionTensor",
        input=[i for i in inputs],
    ),
    "Dense": lambda nodes, inputs, tensors, _: Struct(
        op="Dense",
        input=[i for i in inputs] + [t.name for t in tensors],
        data_frmt=get_attr(
            by_op(nodes, "Dense") or by_op(nodes, "MatMul"), "data_format"
        ),
    ),
    "Conv2D": lambda nodes, inputs, tensors, _: Struct(
        op="Conv2D",
        input=[i for i in inputs] + [t.name for t in tensors],
        padding=get_attr(by_op(nodes, "Conv2D"), "padding"),
        strides=get_attr(by_op(nodes, "Conv2D"), "strides"),
        dilations=get_attr(by_op(nodes, "Conv2D"), "dilations"),
        data_frmt=get_attr(by_op(nodes, "Conv2D"), "data_format"),
    ),
    "DepthwiseConv2dNative": lambda nodes, inputs, tensors, _: Struct(
        op="DepthwiseConv2dNative",
        input=[i for i in inputs] + [t.name for t in tensors],
        padding=get_attr(by_op(nodes, "DepthwiseConv2dNative"), "padding"),
        strides=get_attr(by_op(nodes, "DepthwiseConv2dNative"), "strides"),
        dilations=get_attr(by_op(nodes, "DepthwiseConv2dNative"), "dilations"),
        data_frmt=get_attr(by_op(nodes, "DepthwiseConv2dNative"), "data_format"),
    ),
    "Conv2DBackpropInput": lambda nodes, inputs, tensors, _: Struct(
        op="Conv2DBackpropInput",
        input=[i for i in inputs] + [t.name for t in tensors],
        padding=get_attr(by_op(nodes, "Conv2DBackpropInput"), "padding"),
        strides=get_attr(by_op(nodes, "Conv2DBackpropInput"), "strides"),
        dilations=get_attr(by_op(nodes, "Conv2DBackpropInput"), "dilations"),
        data_frmt=get_attr(by_op(nodes, "Conv2DBackpropInput"), "data_format"),
    ),
    "BasicLSTM": lambda nodes, inputs, tensors, context: basic_lstm(
        nodes, inputs, tensors, context
    ),
    "Swish": lambda nodes, inputs, tensors, _: Struct(op="Swish", input=inputs),
    # TODO:'Round'
    # TODO:'Rsqrt'
}

# Parse
def get_attr(node, attr_name, default=None):
    if type(node) == Struct:
        if hasattr(node, attr_name):
            return getattr(node, attr_name)
        else:
            return default

    # See: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/attr_value.proto
    val = node.attr[attr_name]

    if val.HasField("list"):
        return val.list.i
        # NOTE: can't find way to identify type of list BUT it is almost always list(int)
        # except list(float) in FractionalAvg/MaxPool
    if val.HasField("b"):
        return val.b
    if val.HasField("i"):
        return val.i
    if val.HasField("f"):
        return val.f
    if val.HasField("s"):
        return val.s.decode("utf-8")
    if val.HasField("shape"):
        return val.shape
    if val.HasField("tensor"):
        return val.tensor
    return default


def get_epsilon(layer):
    return get_attr(
        layer, "epsilon", default=0.001
    )  # default epsilon taken from tf.layers.batch_normalization


def get_layer_shape(layer):
    shape = get_attr(layer, "shape")
    if not shape:
        return [-1, -1, -1, -1]
    shape = [dim.size for dim in shape.dim]
    if len(shape) == 1:
        return [1, 1, 1, shape[0]]
    if len(shape) == 2:
        return [shape[0], 1, 1, shape[1]]
    return shape


def get_tensor_dims(tensor):
    if isinstance(tensor, np.ndarray):
        return np.shape(tensor)

    dims = []
    if tensor.tensor_shape:
        dims = [v.size for v in tensor.tensor_shape.dim]
    if tensor.float_val:
        dims = np.shape(tensor.float_val)
    if tensor.int_val:
        dims = np.shape(tensor.int_val)
    if tensor.bool_val:
        dims = np.shape(tensor.bool_val)
    return dims


def get_tensor_dtype(tensor):
    if isinstance(tensor, np.ndarray):
        return tensor.dtype

    dataType = ""
    fields = tensor.ListFields()

    for field, value in fields:
        if (
            field.name == "dtype"
            and field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_ENUM
        ):
            dataType = field.enum_type.values_by_number.get(value, None).name

    return dataType


def get_tensor_data(tensor):
    if isinstance(tensor, np.ndarray):
        return tensor.astype(float)

    dims = get_tensor_dims(tensor)
    elems = np.product(dims)

    if tensor.tensor_content:
        # TODO: support other types
        dataType = get_tensor_dtype(tensor)
        if dataType == "DT_FLOAT":
            data = struct.unpack("<" + str(elems) + "f", tensor.tensor_content)
        elif dataType == "DT_INT32":
            data = struct.unpack("<" + str(elems) + "i", tensor.tensor_content)
        elif dataType == "DT_BOOL":
            data = struct.unpack("<" + str(elems) + "?", tensor.tensor_content)
        else:
            print("UNSUPPORTED: data type", dataType)
    if tensor.float_val:
        data = tensor.float_val
    if tensor.int_val:
        data = np.array(tensor.int_val, dtype=float)
    if tensor.bool_val:
        data = np.array(tensor.bool_val, dtype=float)
    return np.array(data).reshape(dims)


def flatten(items, enter=lambda x: isinstance(x, list)):
    # http://stackoverflow.com/a/40857703
    # https://github.com/ctmakro/canton/blob/master/canton/misc.py
    """Yield items from any nested iterable; see REF."""
    for x in items:
        if enter(x):
            yield from flatten(x)
        else:
            yield x


def replace_strings_in_list(array_of_strigs, replace_with_strings):
    "A value in replace_with_strings can be either single string or list of strings"
    potentially_nested_list = [
        replace_with_strings.get(s) or s for s in array_of_strigs
    ]
    return list(flatten(potentially_nested_list))


def remove_duplicates_from_list(array):
    "Preserves the order of elements in the list"
    output = []
    unique = set()
    for a in array:
        if a not in unique:
            unique.add(a)
            output.append(a)
    return output


#########################################################


def pool_to_HW(shape, data_frmt):
    """ Convert from NHWC|NCHW => HW
    """
    if len(shape) != 4:
        return shape  # Not NHWC|NCHW, return as is
    if data_frmt == "NCHW":
        return [shape[2], shape[3]]
    return [shape[1], shape[2]]


def strides_to_HW(shape, format):
    return pool_to_HW(shape, format)


#########################################################


def gru(nodes, inputs, tensors, context):
    assert len(inputs) == 2

    def find_tensor_by_name(name, default=None):
        nonlocal tensors
        candidates = [t for t in tensors if t.name.endswith(name)]
        return candidates[0].data if candidates else default

    input = inputs[-1]
    state = inputs[0]
    gates_kernel = find_tensor_by_name("/gates/kernel")
    gates_bias = find_tensor_by_name(
        "/gates/bias", default=np.zeros(np.shape(gates_kernel)[-1])
    )
    candidate_kernel = find_tensor_by_name("/candidate/kernel")
    candidate_bias = find_tensor_by_name(
        "/candidate/bias", default=np.zeros(np.shape(candidate_kernel)[-1])
    )
    new_state = nodes[-1].name + "_h"

    assert np.shape(gates_kernel)[-1] == np.shape(gates_bias)[-1]
    assert np.shape(candidate_kernel)[-1] == np.shape(candidate_bias)[-1]

    num_gates = 2
    seq_length = 1
    hidden_size = np.shape(gates_kernel)[-1] // num_gates

    gate_kernels = np.split(gates_kernel, num_gates, axis=-1)
    gate_biases = np.split(gates_bias, num_gates, axis=-1)

    context.model_tensors["kernel_r"] = gate_kernels[0]
    context.model_tensors["kernel_u"] = gate_kernels[1]
    context.model_tensors["kernel_c"] = candidate_kernel
    context.model_tensors["bias_r"] = gate_biases[0]
    context.model_tensors["bias_u"] = gate_biases[1]
    context.model_tensors["bias_c"] = candidate_bias

    new_layers = barracuda.gru(
        "gru",
        input,
        state,
        "kernel_r",
        "kernel_u",
        "kernel_c",
        "bias_r",
        "bias_u",
        "bias_c",
        new_state,
    )

    state_shape = [1, 1, seq_length, hidden_size]
    context.model_memories += [state_shape, state, new_state]

    # map exptected output of the replaced pattern to output from our GRU cell
    actual_output_node = nodes[-4]
    assert actual_output_node.op == "Reshape"
    context.map_ignored_layer_to_its_input[actual_output_node.name] = new_state

    return new_layers


def basic_lstm(nodes, inputs, tensors, context):
    assert len(inputs) == 2

    def find_tensor_by_name(name, default=None):
        nonlocal tensors
        candidates = [t for t in tensors if t.name.endswith(name)]
        return candidates[0].data if candidates else default

    def find_forget_bias():
        nonlocal nodes
        nonlocal tensors
        # TODO: make it more fault-tolerant
        # search for scalar float constant that is input to Add node
        # and hope it is not a constant for some complex activation function
        for t in tensors:
            if np.prod(t.shape) == 1 and get_tensor_dtype(t.obj) == "DT_FLOAT":
                for n in nodes:
                    if n.op == "Add" and t.name in n.input:
                        return t.data
        return np.zeros(1)

    input = inputs[-1]
    state_c = inputs[0] + "_c"
    state_h = inputs[0] + "_h"
    kernel = find_tensor_by_name("/kernel")
    bias = find_tensor_by_name("/bias", default=np.zeros(np.shape(kernel)[-1]))
    forget_bias = find_forget_bias()
    new_state_c = nodes[-1].name + "_c"
    new_state_h = nodes[-1].name + "_h"

    assert np.shape(kernel)[-1] == np.shape(bias)[-1]

    num_gates = 4
    seq_length = 1
    hidden_size = np.shape(kernel)[-1] // num_gates

    kernels = np.split(kernel, num_gates, axis=-1)
    biases = np.split(bias, num_gates, axis=-1)

    context.model_tensors["kernel_i"] = kernels[0]
    context.model_tensors["kernel_j"] = kernels[1]
    context.model_tensors["kernel_f"] = kernels[2]
    context.model_tensors["kernel_o"] = kernels[3]
    context.model_tensors["bias_i"] = biases[0]
    context.model_tensors["bias_j"] = biases[1]
    context.model_tensors["bias_f"] = biases[2] + forget_bias
    context.model_tensors["bias_o"] = biases[3]

    new_layers = barracuda.lstm(
        "lstm",
        input,
        state_c,
        state_h,
        "kernel_i",
        "kernel_j",
        "kernel_f",
        "kernel_o",
        "bias_i",
        "bias_j",
        "bias_f",
        "bias_o",
        new_state_c,
        new_state_h,
    )

    state_shape = [1, 1, seq_length, hidden_size]
    context.model_memories += [state_shape, state_c, new_state_c]
    context.model_memories += [state_shape, state_h, new_state_h]

    # map expected output of the replaced pattern to output from our LSTM cell
    actual_output_node = nodes[-4]
    assert actual_output_node.op == "Reshape"
    context.map_ignored_layer_to_its_input[actual_output_node.name] = new_state_h

    return new_layers


#########################################################


def process_layer(layer, context, args):
    model_tensors = context.model_tensors
    input_shapes = context.input_shapes
    map_ignored_layer_to_its_input = context.map_ignored_layer_to_its_input

    name = layer.name
    class_name = layer.op
    inputs = (
        layer.input
    )  # Tensorflow inputs are always explicit, but in case of Keras we had 'inputs = layer.input or [prev_layer_name]'
    inputs = replace_strings_in_list(inputs, map_ignored_layer_to_its_input)

    if class_name == "Const":
        model_tensors[name] = layer.attr["value"].tensor
        return

    if class_name == "Placeholder":
        assert inputs == []
        map_ignored_layer_to_its_input[name] = inputs
        input_shapes[name] = get_layer_shape(layer)
        return

    if class_name == "Identity":
        connected_to_const = len(inputs) == 1 and inputs[0] in model_tensors
        if connected_to_const:
            map_ignored_layer_to_its_input[name] = inputs
            return
        else:
            # treat Identity layer that are connected to processing nodes
            # as output from the network
            class_name = "Linear"

    # TEMPORARY: until we implemented rank detection and StidedSlice at runtime
    # HACK: skips trivial StridedSlices for rank=2 tensors
    if (
        class_name == "StridedSlice"
        and get_attr(layer, "begin_mask") == 1
        and get_attr(layer, "end_mask") == 1
    ):
        map_ignored_layer_to_its_input[name] = inputs[0]
        return

    if args.print_layers or args.verbose:
        var_tensors = [i for i in inputs if i not in model_tensors]
        const_tensors = [i for i in inputs if i in model_tensors]
        print(
            "'%s' %s Vars:%s Const:%s" % (name, class_name, var_tensors, const_tensors)
        )

    if class_name in known_activations:
        activation = class_name
        class_name = "Activation"
    else:
        activation = "Linear"

    if not class_name in known_classes:
        if class_name in requires_runtime_flag:
            print("SKIP:", class_name, "layer is used only for training")
        else:
            print("IGNORED:", class_name, "unknown layer")
        map_ignored_layer_to_its_input[name] = inputs
        return

    klass = known_classes[class_name]
    if type(klass) == int:
        klass = Struct(id=klass)

    o_l = Struct()
    o_l.type = klass.id
    o_l.class_name = class_name
    o_l.name = name

    padding = get_attr(layer, "padding")  # layer.attr['padding'].s.decode("utf-8")
    strides = get_attr(layer, "strides")  # layer.attr['strides'].list.i
    dilations = get_attr(layer, "dilations")  # layer.attr['dilations'].list.i
    pool_size = get_attr(layer, "ksize")  # layer.attr['ksize'].list.i
    shape = get_attr(layer, "shape", default=[])
    data_frmt = get_attr(
        layer, "data_format"
    )  # layer.attr['data_format'].s.decode("utf-8")
    axis = get_attr(layer, "axis")
    alpha = get_attr(layer, "alpha")
    beta = get_attr(layer, "beta")

    if activation and not activation in known_activations:
        print("IGNORED: unknown activation", activation)
    if padding and not padding in known_paddings:
        print("IGNORED: unknown padding", padding)
    if data_frmt and not data_frmt in supported_data_formats:
        print("UNSUPPORTED: data format", data_frmt)

    o_l.activation = known_activations.get(activation) or 0
    o_l.pads = known_paddings.get(padding) or [0, 0, 0, 0]
    o_l.strides = strides_to_HW(strides, data_frmt) if strides else []
    o_l.pool_size = pool_to_HW(pool_size, data_frmt) if pool_size else shape
    o_l.axis = axis or -1
    o_l.alpha = alpha or 1
    o_l.beta = beta or 0

    tensor_names = [i for i in inputs if i in model_tensors]
    o_l.tensors = [
        Struct(
            name=x,
            shape=get_tensor_dims(model_tensors[x]),
            data=get_tensor_data(model_tensors[x]),
        )
        for x in tensor_names
    ]
    # Patch shapes & data
    layer_has_model_tensors = len(o_l.tensors) > 0
    if hasattr(klass, "out_shapes") and layer_has_model_tensors:
        shapes = klass.out_shapes([x.shape for x in o_l.tensors])

        # if we have more shapes than actual tensors,
        # then create & fill missing tensors with zeros
        in_tensor_num = len(o_l.tensors)
        for index, new_shape in enumerate(shapes):
            if index >= in_tensor_num:
                new_tensor = Struct(
                    name=("%s/patch:%i") % (name, index - in_tensor_num),
                    shape=new_shape,
                    data=np.zeros(new_shape),
                )
                o_l.tensors.append(new_tensor)
        assert len(shapes) <= len(o_l.tensors)

        if hasattr(klass, "patch_data"):
            data = [x.data for x in o_l.tensors]

            patch_data_fn = klass.patch_data
            patch_data_expected_arg_count = patch_data_fn.__code__.co_argcount
            patch_data_args = (
                (data, layer) if patch_data_expected_arg_count > 1 else (data,)
            )
            tensor_data = patch_data_fn(*patch_data_args)
            o_l.tensors = o_l.tensors[
                : len(tensor_data)
            ]  # resize tensor array to match patched data - patching might reduce number of tensors
            for x, data in zip(o_l.tensors, tensor_data):
                x.data = data

        # after this point we should have equal amount of shapes and tensors
        assert len(o_l.tensors) == len(shapes)

        for x, shape in zip(o_l.tensors, shapes):
            x.shape = shape

        o_l.inputs = [i for i in inputs if i not in model_tensors]

    else:
        # no 'patch_data' lambda was specified, op does not require tensor args
        o_l.tensors = []
        o_l.inputs = inputs

    # Force all tensors to float32
    for x in o_l.tensors:
        x.data = x.data.astype(np.float32)

    # Layer is ready
    context.layers.append(o_l)


class ModelBuilderContext:
    def __init__(self):
        self.layers = []
        self.input_shapes = {}
        self.model_tensors = {}
        self.model_memories = []
        self.map_ignored_layer_to_its_input = {}


def process_model(model, args):
    o_context = ModelBuilderContext()

    # Find node patterns
    nodes_as_array = [node for node in model.node]

    node_index = 0
    while node_index < len(nodes_as_array):
        node = nodes_as_array[node_index]
        match = False
        for pattern_repr, pattern_name in known_patterns.items():
            pattern = eval(pattern_repr)
            if node_index + len(pattern) > len(nodes_as_array):
                continue  # pattern too long, skip

            require_exact_match = pattern[0] == "Const" or pattern[0] == "Identity"
            pattern_end = node_index

            def match_node(node, pattern):
                return node.op == pattern or (
                    hasattr(pattern, "match") and pattern.match(node.name)
                )

            for p in pattern:
                if not require_exact_match:
                    while (
                        pattern_end < len(nodes_as_array)
                        and nodes_as_array[pattern_end].op != p
                        and (
                            nodes_as_array[pattern_end].op == "Const"
                            or nodes_as_array[pattern_end].op == "Identity"
                        )
                    ):
                        pattern_end += 1
                if pattern_end >= len(nodes_as_array):
                    break

                match = False
                if hasattr(p, "match"):  # regexp
                    while pattern_end < len(nodes_as_array) and p.match(
                        nodes_as_array[pattern_end].name
                    ):
                        match = True
                        pattern_end += 1
                else:  # exact string
                    match = nodes_as_array[pattern_end].op == p
                    pattern_end += 1

                if not match:
                    break

            def get_tensors(pattern_nodes):
                nonlocal o_context
                map_ignored_layer_to_its_input = (
                    o_context.map_ignored_layer_to_its_input
                )

                # tensors <= all Const nodes within this pattern
                tensor_nodes = [n for n in pattern_nodes if n.op == "Const"]
                tensors = [
                    Struct(
                        name=n.name,
                        obj=n.attr["value"].tensor,
                        shape=get_tensor_dims(n.attr["value"].tensor),
                        data=get_tensor_data(n.attr["value"].tensor),
                    )
                    for n in tensor_nodes
                ]

                # TODO: unify / reuse code from process_layer
                identity_nodes = [n for n in pattern_nodes if n.op == "Identity"]
                for i in identity_nodes:
                    inputs = replace_strings_in_list(
                        i.input, map_ignored_layer_to_its_input
                    )
                    map_ignored_layer_to_its_input[i.name] = inputs

                # gather inputs from Op nodes (not Const, not Identity)
                op_nodes = [
                    n
                    for n in pattern_nodes
                    if n not in tensor_nodes and n not in identity_nodes
                ]
                inputs_to_op_nodes = list(
                    flatten([list(flatten(n.input)) for n in op_nodes])
                )
                inputs_to_op_nodes = replace_strings_in_list(
                    inputs_to_op_nodes, map_ignored_layer_to_its_input
                )
                inputs_to_op_nodes = [i.split(":")[0] for i in inputs_to_op_nodes]

                # filter only inputs that are coming from nodes that are outside this pattern
                # preserve the order
                pattern_nodes = [n.name for n in pattern_nodes]
                # inputs_from_outside_pattern = remove_duplicates_from_list([i for i in inputs_to_op_nodes if nodes_by_name[i] not in pattern_nodes])
                inputs_from_outside_pattern = remove_duplicates_from_list(
                    [i for i in inputs_to_op_nodes if i not in pattern_nodes]
                )

                return inputs_from_outside_pattern, tensors

            if match:
                nodes = nodes_as_array[node_index:pattern_end]
                name = nodes[-1].name
                var_tensors, const_tensors = get_tensors(nodes)
                if args.print_patterns or args.verbose:
                    print(
                        "PATTERN:",
                        name,
                        "~~",
                        pattern_name,
                        pattern,
                        "<-",
                        var_tensors,
                        "+",
                        [t.name for t in const_tensors],
                    )
                for n in nodes:
                    if n.op == "Const" or n.op == "Identity":
                        process_layer(n, o_context, args)

                new_layers = transform_patterns[pattern_name](
                    nodes, var_tensors, const_tensors, o_context
                )
                if not isinstance(new_layers, list):
                    if not hasattr(new_layers, name):
                        new_layers.name = name
                    new_layers = [new_layers]

                for l in new_layers:
                    # TODO: prefix new layer names with scope, patch inputs
                    # l.name = name + '/' + l.name
                    process_layer(l, o_context, args)

                node_index = pattern_end
                break  # pattern found & processed

        if not match:
            # TODO: gather tensors in the same way as patterns do
            process_layer(node, o_context, args)
            node_index += 1

    return (
        o_context.layers,
        o_context.input_shapes,
        o_context.model_tensors,
        o_context.model_memories,
    )


#########################################################


def convert(
    source_file,
    target_file,
    trim_unused_by_output="",
    verbose=False,
    compress_f16=False,
):
    """
    Converts a TensorFlow model into a Barracuda model.
    :param source_file: The TensorFlow Model
    :param target_file: The name of the file the converted model will be saved to
    :param trim_unused_by_output: The regexp to match output nodes to remain in the model. All other uconnected nodes will be removed.
    :param verbose: If True, will display debug messages
    :param compress_f16: If true, the float values will be converted to f16
    :return:
    """
    if type(verbose) == bool:
        args = Struct()
        args.verbose = verbose
        args.print_layers = verbose
        args.print_source_json = verbose
        args.print_barracuda_json = verbose
        args.print_layer_links = verbose
        args.print_patterns = verbose
        args.print_tensors = verbose
    else:
        args = verbose

    # Load Tensorflow model
    print("Converting %s to %s" % (source_file, target_file))
    f = open(source_file, "rb")
    i_model = tf.GraphDef()
    i_model.ParseFromString(f.read())

    if args.verbose:
        print("OP_TYPES:", {layer.op for layer in i_model.node})

    if args.print_source_json or args.verbose:
        for layer in i_model.node:
            if not layer.op == "Const":
                print("MODEL:", MessageToJson(layer) + ",")

    # Convert
    o_model = barracuda.Model()
    o_model.layers, o_input_shapes, o_model.tensors, o_model.memories = process_model(
        i_model, args
    )

    # Cleanup unconnected Identities (they might linger after processing complex node patterns like LSTM)
    def cleanup_layers(layers):
        all_layers = {l.name for l in layers}
        all_inputs = {i for l in layers for i in l.inputs}

        def is_unconnected_identity(layer):
            if layer.class_name == "Activation" and layer.activation == 0:  # Identity
                assert len(layer.inputs) == 1
                if layer.inputs[0] not in all_layers and layer.name not in all_inputs:
                    return True
            return False

        return [l for l in layers if not is_unconnected_identity(l)]

    o_model.layers = cleanup_layers(o_model.layers)

    all_inputs = {i for l in o_model.layers for i in l.inputs}
    embedded_tensors = {t.name for l in o_model.layers for t in l.tensors}

    # Find global tensors
    def dims_to_barracuda_shape(dims):
        shape = list(dims)
        while len(shape) < 4:
            shape = [1] + shape
        return shape

    o_model.globals = [
        t for t in o_model.tensors if t not in all_inputs and t not in embedded_tensors
    ]
    # for x in global_tensors:
    #    shape = dims_to_barracuda_shape(get_tensor_dims(o_model.tensors[x]))
    #    o_globals += [Struct(
    #        name = x,
    #        shape = shape,
    #        data = np.reshape(get_tensor_data(o_model.tensors[x]), shape).astype(np.float32))]

    # Trim
    if trim_unused_by_output:
        o_model.layers = barracuda.trim(
            o_model.layers, trim_unused_by_output, args.verbose
        )

    # Create load layers for constants
    const_tensors = [i for i in all_inputs if i in o_model.tensors]
    const_tensors += o_model.globals
    for x in const_tensors:
        shape = dims_to_barracuda_shape(get_tensor_dims(o_model.tensors[x]))

        o_l = Struct(
            type=255,  # Load
            class_name="Const",
            name=x,
            pads=[0, 0, 0, 0],
            strides=[],
            pool_size=[],
            axis=-1,
            alpha=1,
            beta=0,
            activation=0,
            inputs=[],
            tensors=[
                Struct(
                    name=x,
                    shape=shape,
                    data=np.reshape(get_tensor_data(o_model.tensors[x]), shape).astype(
                        np.float32
                    ),
                )
            ],
        )
        o_model.layers.insert(0, o_l)

    # Find model inputs & outputs
    all_layers = {l.name for l in o_model.layers}
    # global inputs => are inputs that are NOT connected to any layer in the network
    # global outputs => are outputs that are NOT feeding any layer in the network OR are coming from Identity layers
    o_model.inputs = {
        i: o_input_shapes[i]
        for l in o_model.layers
        for i in l.inputs
        if i not in all_layers and i not in o_model.memories
    }

    def is_output_layer(layer):
        if (
            layer.class_name == "Const"
        ):  # Constants never count as global output even when unconnected
            return False
        if (
            layer.name not in all_inputs
        ):  # this layer is not inputing to any other layer
            return True
        if (
            layer.class_name == "Activation" and layer.activation == 0
        ):  # Identity marks global output
            return True
        return False

    o_model.outputs = [l.name for l in o_model.layers if is_output_layer(l)]

    # Compress
    if compress_f16:
        o_model = barracuda.compress(o_model)

    # Sort model so that layer inputs are always ready upfront
    o_model.layers = barracuda.sort(
        o_model.layers, o_model.inputs, o_model.memories, args.verbose
    )

    # Summary
    barracuda.summary(
        o_model,
        print_layer_links=args.print_layer_links or args.verbose,
        print_barracuda_json=args.print_barracuda_json or args.verbose,
        print_tensors=args.print_tensors or args.verbose,
    )

    # Write to file
    barracuda.write(o_model, target_file)
    print("DONE: wrote", target_file, "file.")