import inspect
from itertools import chain

import numpy as np
from onnx import NodeProto
from onnx import TensorProto
from onnx import ValueInfoProto
from onnx import numpy_helper
from onnx.helper import make_graph
from onnx.helper import make_tensor
from onnx.helper import make_tensor_value_info
from onnx.helper import mapping
import tensorflow as tf
from tensorflow.core.framework.attr_value_pb2 import AttrValue
from tensorflow.core.framework.node_def_pb2 import NodeDef

from onnx_tf.common import attr_converter
from onnx_tf.common import attr_translator
from onnx_tf.common import CONST_MINUS_ONE_INT32
from onnx_tf.common import CONST_ONE_FP32
from onnx_tf.common import CONST_ONE_INT32
from onnx_tf.common import CONST_ZERO_INT32
from onnx_tf.common import IS_PYTHON3
from onnx_tf.common import logger
from onnx_tf.common.data_type import any_dtype_to_onnx_dtype


class TensorflowNode(object):

  def __init__(self,
               node=None,
               name=None,
               inputs=None,
               outputs=None,
               attr=None,
               domain=None,
               op_type=None):
    # storing a reference to the original protobuf object
    if node is None:
      self.node = None
      self.name = name or ""
      self.inputs = inputs or []
      self.attr = attr or {}
      self.domain = domain or ""
      self.op_type = op_type or ""
      self.outputs = outputs or self.get_outputs_names()
    elif isinstance(node, (OnnxNode, NodeProto)):
      self._load_onnx_node(node)
    elif isinstance(node, NodeDef):
      self._load_tf_node(node)

  def _load_onnx_node(self, node):
    if isinstance(node, NodeProto):
      node = OnnxNode(node)
    self.name = node.name
    self.inputs = node.inputs
    self.outputs = node.outputs
    self.attr = node.attrs
    self.domain = node.domain
    self.op_type = node.op_type

  def _load_tf_node(self, node):
    self.node = node
    self.name = node.name
    self.inputs = list(node.input)
    self.attr = {}
    for key, val in node.attr.items():
      new_val = attr_translator.translate_tf(key, val)
      if isinstance(new_val, AttrValue):
        new_val = attr_converter.convert_tf(new_val)
      self.attr[key] = new_val
    splitted_op_name = node.op.split(".")
    self.domain = "" if len(splitted_op_name) == 1 else ".".join(
        splitted_op_name[:-1])
    self.op_type = splitted_op_name[-1]
    self.outputs = self.get_outputs_names()

  def get_outputs_names(self, num=None):
    """ Helper method to get outputs names.
    e.g. tf.split: [Split, Split:1, Split:2]

    :param num: Force to get `num` outputs names.
    :return: List of outputs names.
    """
    if num is None:
      if "_output_shapes" in self.attr:
        num = len(self.attr["_output_shapes"])
      else:
        num = 1
        logger.warning("_output_shapes is not in node.attr. "
                      "The num of output is set to 1 for commonly. "
                      "It will cause problem with case of multiple outputs.")
    return [
        self.name + ":{}".format(i) if i > 0 else self.name for i in range(num)
    ]


class TensorflowGraph(object):

  def __init__(self, graph_def, outputs=(), graph_name="graph"):
    self._graph_name = graph_name
    self._graph_def = self._process_graph_def(graph_def)
    self._nodes = self._create_util_nodes() + [
        TensorflowNode(node) for node in self.graph_def.node
    ]
    self._nodes_dict = {n.name: n for n in self._nodes}
    self._outputs = outputs or self.get_output_node_names(self.graph_def)

  @staticmethod
  def _create_util_nodes():
    util_nodes = [(CONST_MINUS_ONE_INT32, np.array([-1]).astype(np.int32)),
                  (CONST_ZERO_INT32, np.array([0]).astype(np.int32)),
                  (CONST_ONE_INT32, np.array([1]).astype(np.int32))]
    return [
        TensorflowNode(
            op_type="Const",
            name=name,
            attr={
                "value": value,
                "dtype": any_dtype_to_onnx_dtype(value.dtype),
                "_output_shapes": [value.shape]
            }) for name, value in util_nodes
    ]

  def get_node_by_name(self, name):
    node = self._nodes_dict.get(name, None)
    if node is None:
      raise ValueError(
          "Node {} is not found in the graph provided".format(name))
    return node

  def _process_graph_def(self, graph_def):
    if "_output_shapes" not in TensorflowNode(graph_def.node[0]).attr:
      graph_def = self._add_infer_shapes(graph_def)
    return graph_def

  @staticmethod
  def _add_infer_shapes(graph_def):
    with tf.Graph().as_default():
      with tf.Session(
          config=tf.ConfigProto(
              graph_options=tf.GraphOptions(infer_shapes=True))) as sess:
        tf.import_graph_def(graph_def, name="")
      return sess.graph_def

  @staticmethod
  def get_output_node_names(graph_def):
    """Get output node names from GraphDef.

    Args:
      graph_def: GraphDef object.

    Returns:
      List of output node names.
    """
    input_names, output_names = set(), set()
    for node in graph_def.node:
      output_names.add(node.name)
      input_names.update(set(node.input))
    return list(output_names - input_names)

  def update_nodes(self, nodes):
    self._nodes = nodes
    self._nodes_dict = {n.name: n for n in self._nodes}

  @property
  def graph_def(self):
    return self._graph_def

  @property
  def graph_name(self):
    return self._graph_name

  @property
  def nodes(self):
    return self._nodes

  @property
  def nodes_dict(self):
    return self._nodes_dict

  @property
  def outputs(self):
    return self._outputs


# TODO: Move this into ONNX main library
class OnnxNode(object):
  """
  Reimplementation of NodeProto from ONNX, but in a form
  more convenient to work with from Python.
  """

  def __init__(self, node):
    self.name = str(node.name)
    self.op_type = str(node.op_type)
    self.domain = str(node.domain)
    self.attrs = dict([(attr.name,
                        attr_translator.translate_onnx(
                            attr.name, attr_converter.convert_onnx(attr)))
                       for attr in node.attribute])
    self.inputs = list(node.input)
    self.outputs = list(node.output)
    self.node_proto = node


class OnnxGraph(object):
  """ A helper class for making ONNX graph.
  This class holds all information ONNX graph needs.
  """

  def __init__(self, name=None, graph_proto=None):
    if graph_proto:
      self._name = graph_proto.name
      self._inputs_proto = list(graph_proto.input)
      self._outputs_proto = list(graph_proto.output)
      self._nodes_proto = list(graph_proto.node)
      self._consts_proto = list(graph_proto.initializer)
      self._value_info_proto = list(graph_proto.value_info)
      self._consts = dict([(init.name, numpy_helper.to_array(init))
                           for init in graph_proto.initializer])
    else:
      self._name = name or ""
      self._inputs_proto = []
      self._outputs_proto = []
      self._nodes_proto = []
      self._consts = {}
      self._consts_proto = []
      self._value_info_proto = []
    # Either way, data_type_cast_map is empty when initialized.
    self._data_type_cast_map = {}

    self._add_utility_constants()

  def _add_utility_constants(self):
    util_consts = {CONST_ONE_FP32: np.array([1.0]).astype(np.float32)}
    # Add a few useful utility constants:
    for name, value in util_consts.items():
      self.add_const_explicit(name=name, value=value)
      self.add_const_proto_explicit(
          name=name, value=value, np_dtype=value.dtype)
      self.add_input_proto_explicit(
          name=name, shape=value.shape, np_dtype=value.dtype)

  # This list holds the protobuf objects of type ValueInfoProto
  # representing the input to the converted ONNX graph.
  @property
  def inputs_proto(self):
    return self._inputs_proto

  @inputs_proto.setter
  def inputs_proto(self, inputs_proto):
    self._inputs_proto = inputs_proto

  @property
  def all_node_inputs(self):
    return list(chain.from_iterable(map(lambda p: p.input, self._nodes_proto)))

  @property
  def outputs(self):
    return list(map(lambda p: p.name, self._outputs_proto))

  @property
  def outputs_proto(self):
    return self._outputs_proto

  # This list holds the protobuf objects of type NodeProto
  # representing the ops in the converted ONNX graph.
  @property
  def nodes_proto(self):
    return self._nodes_proto

  @nodes_proto.setter
  def nodes_proto(self, nodes_proto):
    self._nodes_proto = nodes_proto

  # This dictionary contains a map from the name of the constant
  # op to the array of values it holds. This is useful because
  # tensorflow is less eager to know about input values at
  # graph construction time than ONNX. That is to say, some ONNX
  # attributes are input tensors in TF. This dictionary extracts
  # those values of constant tensors that are known at graph
  # construction time.
  @property
  def consts(self):
    return self._consts

  @consts.setter
  def consts(self, consts):
    self._consts = consts

  # Sometimes the constants are used as inputs to ops. This list
  # holds initializers that creates global constant tensors available
  # to be accessed by ops as inputs (as oppose to attributes which
  # is supplied by the `consts` map above).
  @property
  def consts_proto(self):
    return self._consts_proto

  @consts_proto.setter
  def consts_proto(self, consts_proto):
    self._consts_proto = consts_proto

  # A map holds nodes name and new data type. Will be used to
  # process protos to match ONNX type constraints.
  @property
  def data_type_cast_map(self):
    return self._data_type_cast_map

  @data_type_cast_map.setter
  def data_type_cast_map(self, data_type_cast_map):
    self._data_type_cast_map = data_type_cast_map

  # This list holds the protobuf objects of type ValueInfoProto
  # representing the all nodes' outputs to the converted ONNX graph.
  @property
  def value_info_proto(self):
    return self._value_info_proto

  def add_input_proto_explicit(self,
                               name,
                               shape,
                               np_dtype=None,
                               tf_dtype=None,
                               onnx_dtype=None):
    onnx_dtype = any_dtype_to_onnx_dtype(
        np_dtype=np_dtype, tf_dtype=tf_dtype, onnx_dtype=onnx_dtype)
    input_proto = make_tensor_value_info(name, onnx_dtype, shape)
    self._inputs_proto.append(input_proto)

  def add_input_proto(self, node):
    name = node.name
    onnx_dtype = node.attr["dtype"]
    shape = node.attr["shape"] if node.op_type != "Const" else node.attr[
        'value'].shape
    self.add_input_proto_explicit(name, shape, onnx_dtype=onnx_dtype)

  def add_output_proto(self, node):
    output_onnx_type = node.attr.get("T", TensorProto.BOOL)
    for i, output_shape in enumerate(node.attr["_output_shapes"]):
      output_name = node.name + ":{}".format(i) if i > 0 else node.name
      self._outputs_proto.append(
          make_tensor_value_info(output_name, output_onnx_type, output_shape))

  def add_node_proto(self, node_proto):
    if not isinstance(node_proto, (list, tuple)):
      node_proto = [node_proto]
    self._nodes_proto.extend(node_proto)

  def remove_node_proto(self, names):
    if not isinstance(names, (list, tuple)):
      names = [names]
    self._nodes_proto = list(
        filter(lambda x: x.name not in names, self._nodes_proto))

  def add_const_explicit(self, name, value):
    self._consts[name] = value

  def add_const(self, node):
    self.add_const_explicit(node.name, node.attr["value"])

  def add_const_proto_explicit(self,
                               name,
                               value,
                               np_dtype=None,
                               tf_dtype=None,
                               onnx_dtype=None):
    onnx_dtype = any_dtype_to_onnx_dtype(
        np_dtype=np_dtype, tf_dtype=tf_dtype, onnx_dtype=onnx_dtype)

    const_dim = len(value.shape)

    if const_dim == 0:
      raw_values = [value.tolist()]
      values = [value]
    else:
      raw_values = value.flatten().tolist()
      values = value

    shape = np.array(values).shape
    const_proto = make_tensor(
        name=name, data_type=onnx_dtype, dims=shape, vals=raw_values)
    self._consts_proto.append(const_proto)

  def add_const_proto(self, node):
    self.add_const_proto_explicit(
        node.name, node.attr["value"], onnx_dtype=node.attr["dtype"])

  def add_value_info_proto(self, node):
    node_onnx_type = node.attr.get("T", TensorProto.BOOL)
    for i, output_shape in enumerate(node.attr["_output_shapes"]):
      node_name = node.name + ":{}".format(i) if i > 0 else node.name
      value_info_proto = make_tensor_value_info(node_name, node_onnx_type,
                                                output_shape)
      self._value_info_proto.append(value_info_proto)

  # Remove proto in inputs_proto and consts_proto
  # if proto is not used as input or an output in ONNX
  def _clean_graph(self):
    in_out = self.all_node_inputs + self.outputs
    self._inputs_proto = list(
        filter(lambda x: x.name in in_out, self.inputs_proto))
    self._consts_proto = list(
        filter(lambda x: x.name in in_out, self.consts_proto))

  def _fix_data_type(self):
    self.inputs_proto = self._data_type_caster(self.inputs_proto,
                                               self.data_type_cast_map)
    self.consts_proto = self._data_type_caster(self.consts_proto,
                                               self.data_type_cast_map)

  @classmethod
  def _data_type_caster(cls, protos, data_type_cast_map):
    """Cast to a new data type if node name is in data_type_cast_map.
    Be used to process protos to match ONNX type constraints.

    :param protos: Target protos.
      TensorProto for inputs and ValueInfoProto for consts.
    :param data_type_cast_map: A {node.name: new_data_type} dict.
    :return: Processed protos.
    """
    if not data_type_cast_map:
      return protos
    result = []
    for proto in protos:
      new_proto = proto
      if proto.name in data_type_cast_map:
        new_data_type = data_type_cast_map[proto.name]
        if type(proto) == TensorProto and proto.data_type != new_data_type:
          field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[
              mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[proto.data_type]]
          vals = getattr(proto, field)
          new_proto = make_tensor(
              name=proto.name,
              data_type=new_data_type,
              dims=proto.dims,
              vals=vals)
        elif type(
            proto
        ) == ValueInfoProto and proto.type.tensor_type.elem_type != new_data_type:
          new_proto.type.tensor_type.elem_type = new_data_type
      result.append(new_proto)
    return result

  def make_graph_proto(self):
    self._clean_graph()
    self._fix_data_type()

    if IS_PYTHON3:
      params = list(inspect.signature(make_graph).parameters.keys())
    else:
      params = inspect.getargspec(make_graph).args

    kwargs = {
        "initializer": self.consts_proto,
        "value_info": self.value_info_proto
    }

    return make_graph(self.nodes_proto, self._name, self.inputs_proto,
                      self.outputs_proto,
                      **dict([(k, kwargs[k]) for k in kwargs if k in params]))