

#
# Copyright 2019 Xilinx Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import tensorflow as tf
from collections import OrderedDict

from tensorflow.python.keras import activations
from tensorflow.python.keras import layers as keras_layers
from tensorflow.python.util import nest
from tensorflow.python.util import tf_inspect

from tf_nndct import layers as nndct_layers
from tf_nndct.graph import OpTypes
from tf_nndct.graph import ops
from tf_nndct.graph import utils
from tf_nndct.quantization import utils as quant_utils
from tf_nndct.utils import generic_utils
from tf_nndct.utils import registry

_INPUT_ARG_PREFIX = '%input%_'

# call(self, input_0, input_1, ...)
_CALL_ARG_TEMPLATE = 'input_%d'

class Printer(object):

  def __init__(self, init_indent_level=0, indent_length=2):
    self._indent_level = init_indent_level
    self._init_indent_level = init_indent_level
    self._indent_length = indent_length
    self._start_of_line = False
    self._strs = []

  def indent(self):
    self._indent_level += 1

  def outdent(self):
    if self._indent_level == 0 or self._indent_level < self._init_indent_level:
      raise RuntimeError("outdent() without matching indent()")
    self._indent_level -= 1

  def current_indentation(self):
    return self._indent_length * self._indent_level

  def newline(self):
    self.print('\n')

  def print(self, text):
    text = str(text)
    if self._indent_level > 0:
      pos = 0
      for i, c in enumerate(text):
        if c == '\n':
          self._append(text[pos:i + 1])
          pos = i + 1
          self._start_of_line = True
      self._append(text[pos:])
    else:
      self._append(text)
      if text[-1] == '\n':
        self._start_of_line = True

  def get(self):
    return ''.join(self._strs)

  def _append(self, text):
    if not text:
      return
    if self._start_of_line:
      self._start_of_line = False
      self._append_indent()

    self._strs.append(text)

  def _append_indent(self):
    if self._indent_level == 0:
      return

    size = self.current_indentation()
    self._strs.append(size * ' ')

class ModuleClassWriter(object):

  def __init__(self, graph, class_spec):
    self._class_spec = class_spec

    # alias for imported module to be written.
    self._module_alias = {}

    self._graph = graph
    # TODO(yuwang): Remove these two members ?
    self._input_signature = graph.input_signature
    self._structured_output_tensors = graph.structured_output_tensors

    self._translator = GraphTranslator(graph, self._class_spec.quantized)
    self._printer = Printer()

  def write(self, filepath):
    self._printer.print('# Generated by NNDCT. Do not edit!')
    self._printer.newline()

    self.write_imports()
    self.write_class_def()

    generic_utils.mkdir_if_not_exist(os.path.dirname(filepath))
    with open(filepath, 'w') as f:
      f.write(self._printer.get())
      f.flush()
      os.fsync(f.fileno())

    layer_to_node = {}
    for node in self._graph.nodes:
      entity = self._translator.get_entity(node.name)
      if entity.need_init:
        layer_to_node[entity.name] = node

    return layer_to_node

  def write_imports(self):
    objects = [self._class_spec.base]
    for entity in self._translator.entities:
      if not entity.obj:
        continue
      objects.append(entity.obj)

    imports = set()
    for obj in objects:
      imports.add(self._generate_obj_import(obj))

    imports = sorted(imports)
    for impt in imports:
      self._printer.print(impt)
      self._printer.newline()
    self._printer.newline()

  def _generate_obj_import(self, object):
    module, pkg, name = _obj_modules(object)
    if module in self._module_alias:
      alias = self._module_alias[module]
    else:
      alias = name
      while alias in self._module_alias.values():
        alias = '_' + alias
      self._module_alias[module] = alias

    if alias == name:
      return 'from {} import {}'.format(pkg, name)
    else:
      return 'from {} import {} as {}'.format(pkg, name, alias)

  def write_class_def(self):
    module, _, _ = _obj_modules(self._class_spec.base)
    self._printer.print('class {cls_name}({base_cls}):\n'.format(
        cls_name=self._class_spec.name,
        base_cls='.'.join(
            [self._module_alias[module], self._class_spec.base.__name__])))
    self._printer.indent()
    self.write_init()
    self.write_call_fn(self._class_spec.call_fn_name)
    self._printer.outdent()

  def write_init(self):
    self._printer.print('def __init__(self):')
    self._printer.newline()
    self._printer.indent()

    self._printer.print('super({}, self).__init__(name={})'.format(
        self._class_spec.name, utils.stringfy_to_write(self._class_spec.name)))
    self._printer.newline()

    for entity in self._translator.entities:
      if not entity.need_init:
        continue

      init_str = self._init_string(entity)
      self._printer.print('self.{} = {}'.format(entity.name, init_str, entity))
      self._printer.newline()

    self._printer.newline()
    self._printer.outdent()

  def _init_string(self, entity):
    args = []
    input_index = 0
    for arg, value in entity.args.items():
      # Use input entity name to fill in placeholder.
      if arg.startswith(_INPUT_ARG_PREFIX):
        args.append(entity.inputs[input_index].name)
        input_index += 1
        continue

      values = generic_utils.to_list(value)
      value_strs = []
      for value in values:
        if isinstance(value, str):
          # Use a quote to wrap a string value.
          value = ''.join(['\'', value, '\''])
        elif isinstance(value, Entity):
          value = 'self.' + value.name
        else:
          pass
        value_strs.append(str(value))
      if len(value_strs) <= 1:
        args.append('{}={}'.format(arg, value_strs[0]))
      else:
        args.append('{}={}'.format(arg, '[' + ', '.join(value_strs) + ']'))

    module, _, _ = _obj_modules(entity.obj)
    init_str = '{module}.{object}({args})'.format(
        module=self._module_alias[module],
        object=entity.obj.__name__,
        args=', '.join(args))
    return init_str

  def write_call_fn(self, fn_name):
    call_args = ['self']
    for i in range(len(self._input_signature)):
      call_args.append(_CALL_ARG_TEMPLATE % i)
    self._printer.print('def {}({}):'.format(fn_name, ', '.join(call_args)))
    self._printer.newline()
    self._printer.indent()

    placeholders = []
    for node in self._graph.nodes:
      if node.op.type == OpTypes.INPUT:
        assert len(node.out_tensors) == 1
        placeholders.append(node)
    # [args_0, args_1_1, args_1] -> [args_0, args_1, args_1_1]
    placeholders = sorted(placeholders, key=lambda x: x.name)

    flattened_input_signature = nest.flatten(self._input_signature)
    assert len(placeholders) == len(flattened_input_signature)

    # TODO(yuwang): Maybe use more solid way to map ?
    # We sort the placeholders by their names and map the flattened input
    # signature in the sorted order. This method works now but relies on the
    # tensorflow's mechanism for generating placeholder names from input
    # argument. Specifically, the first input is named args_0, the second
    # input is named args_1 and so on. If an input is a type of sequence,
    # then the input is flattened to a list and each element in the list
    # is named by args_[input_index]_[element_index_in_the_list].
    # For example, we have a call function and the signature is like:
    #   def call(self, input, state)
    # The first argument 'input' is a tensor and TF will generate a
    # placeholder for it named args_0. The second argument 'state' is a list
    # of tensor, then TF will generate a placeholder for each element in the
    # list and the names of the placeholders are arg_1 and arg_1_1. In the
    # end, we got got three placeholders [args_0, args_1, args_1_1]. In a
    # nndct graph, the order of the placeholders may be [args_0, args_1_1,
    # args_1], that is why we sort the nodes by their names before do mapping.
    # Example:
    # placeholders: ['args_0', 'args_1', 'args_1_1']
    # args: (TensorSpec(shape=(1, 10), dtype=tf.float32, name='args_0'), [TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_1/0'), TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_1/1')])
    # args_to_placeholder = {
    #     'args_0': 'args_0',
    #     'args_1/0': 'args_1',
    #     'args_1/1': 'args_1_1',
    # }
    #
    # Example:
    # placeholders: ['args_0', 'args_0_1', 'args_0_2']
    # args: [{'y': TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/0/y'),
    #         'x': TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/0/x')},
    #        TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/1')]
    # args_to_placeholder = {
    #     'args_0/0/y': 'arg_0',
    #     'args_0/0/x': 'arg_0_1',
    #     'args_0/1': 'arg_0_2',
    # }

    args_to_placeholder = {}
    for i, arg in enumerate(flattened_input_signature):
      args_to_placeholder[arg.name] = placeholders[i].name

    def splice_retriving_path(getter, key):
      key = utils.stringfy_to_write(key)

      if getter == '[]':
        return '[{}]'.format(key)
      return '{}{}'.format(getter, key)

    def arg_retriving_path(arg, path=()):
      """
      Get retriving path of an argument.

      Args:
        arg: The input signature of an argument.

      Yield:

      See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/nest.py::_yield_sorted_items
      """
      if not nest.is_sequence(arg):
        yield path
      elif isinstance(arg, nest._collections_abc.Mapping):
        for key in nest._sorted(arg):
          for res in arg_retriving_path(arg[key], path + (('[]', key),)):
            yield res
      elif nest._is_attrs(arg):
        for item in nest._get_attrs_items(arg):
          for res in arg_retriving_path(item[1], path + (('.', item[0]),)):
            yield res
      elif nest._is_namedtuple(arg):
        for field in arg._fields:
          for res in arg_retriving_path(
              getattr(arg, field), path + (('.', field),)):
            yield res
      # Doesn't support composite_tensor comprared with _yield_sorted_items.
      elif nest._is_type_spec(arg):
        # Note: to allow CompositeTensors and their TypeSpecs to have matching
        # structures, we need to use the same key string here.
        for res in arg_retriving_path(arg._component_specs,
                                      path + (('.', arg.value_type.__name__),)):
          yield res
      else:
        for item in enumerate(arg):
          for res in arg_retriving_path(item[1], path + (('[]', item[0]),)):
            yield res

    for i, arg in enumerate(self._input_signature):
      for path in list(arg_retriving_path(arg)):
        # [('[]' 0), ('[]', 'x')]
        arg_name_path = ['args_%d' % i]
        retriving_path = []
        for getter, key in path:
          arg_name_path.append(key)
          retriving_path.append(splice_retriving_path(getter, key))
        arg_name = '/'.join(str(s) for s in arg_name_path)
        retriving = ''.join(r for r in retriving_path)
        if arg_name not in args_to_placeholder:
          continue

        placeholder_name = args_to_placeholder[arg_name]
        entity = self._translator.get_entity(placeholder_name)
        call_arg = _CALL_ARG_TEMPLATE % i
        #self._printer.print('{} = {}{}'.format(
        #    entity.outputs[0].name, call_arg, retriving))
        #self._printer.newline()
        entity.inputs.append(
            Entity('{}{}'.format(call_arg, retriving), EntityTypes.Tensor))

    for node in self._graph.nodes:
      #if node.op.type == OpTypes.INPUT:
      #  continue
      self._printer.print(self._call_string(node))
      self._printer.newline()

    output_tensors = [
        t.name for t in nest.flatten(self._structured_output_tensors)
    ]
    returns = [self._translator.get_entity(t) for t in output_tensors]
    if len(returns) == 1:
      returns = returns[0]
    else:
      returns = nest.pack_sequence_as(self._structured_output_tensors, returns)

    def return_str(returns):
      return_strs = []
      if isinstance(returns, Entity):
        s = returns.name
      elif isinstance(returns, list):
        for ret in returns:
          return_strs.append(return_str(ret))
        s = ''.join(['[', ', '.join(return_strs), ']'])
      elif isinstance(returns, tuple):
        for ret in returns:
          return_strs.append(return_str(ret))
        s = ''.join(['(', ', '.join(return_strs), ')'])
      elif isinstance(returns, dict):
        for key in returns:
          return_strs.append('{}: {}'.format(
              utils.stringfy_to_write(key), ''.join(return_str(returns[key]))))
        s = ''.join(['{', ', '.join(return_strs), '}'])
      else:
        raise NotImplementedError(
            'Can not rewrite return object of type '.format(type(returns)))
      return s

    self._printer.print('return %s' % return_str(returns))
    self._printer.newline()
    self._printer.outdent()

  def _call_string(self, node):
    entity = self._translator.get_entity(node.name)
    input_names = [i.name for i in entity.inputs]
    output_names = [o.name for o in entity.outputs]
    #if not input_names:
    #  # TODO(yuwang): Support multiple inputs
    #  input_names.append('inputs')

    input_str = ', '.join(input_names)
    output_str = ', '.join(output_names)

    entity = self._translator.get_entity(node.name)
    # Layer type
    if entity.need_init:
      forward_template = '{outputs} = self.{entity}({inputs}) # {node}'
      args_dict = {
          'outputs': output_str,
          'entity': entity.name,
          'inputs': input_str,
          'node': node.name,
      }
      forward_str = forward_template.format(**args_dict)
    else:
      # Function type
      init_str = self._init_string(entity)
      forward_str = '{} = {} # {}'.format(output_str, init_str, node.name)
    return forward_str

class EntityTypes(object):
  Layer = 'layer'
  Function = 'function'
  Tensor = 'tensor'

class Entity(object):

  def __init__(self, name, type, obj=None, args=None):
    self.name = name

    # One of `EntityTypes`.
    self.type = type

    # For EntityTypes.Layer and EntityTypes.Function,
    # obj is the corresponding layer class or function.
    # For EntityTypes.Tensor, obj is None
    self.obj = obj

    self.args = args

    # I/O entities
    self.inputs = []
    self.outputs = []

  def __str__(self):
    return self.name

  @property
  def need_init(self):
    return self.type == EntityTypes.Layer

class GraphTranslator(object):
  """Translates graph components to `Entity` objects.
  Specifically, all `Operation` and `Tensor` objects will be converted to
  `Entity` objects, not including `Node`. This is because a node may contain
  multiple ops in a nested form, so the conversion is performed at the op
  level, not the node level.
  """
  _op_to_layer = {
      OpTypes.BIDIRECTIONAL_RNN: keras_layers.Bidirectional,
      OpTypes.CONV1D: keras_layers.Conv1D,
      OpTypes.DENSE: keras_layers.Dense,
      OpTypes.EMBEDDING: keras_layers.Embedding,
      OpTypes.GRU: keras_layers.GRU,
      OpTypes.INPUT: nndct_layers.Identity,
      OpTypes.LSTM: nndct_layers.LSTM,
      OpTypes.LSTM_CELL: nndct_layers.LSTMCell,
      OpTypes.MAX_POOL1D: keras_layers.MaxPooling1D,
      OpTypes.RNN: keras_layers.RNN,
      OpTypes.SIMPLE_RNN: keras_layers.SimpleRNN,
      OpTypes.STACKED_RNN_CELLS: keras_layers.StackedRNNCells,
  }

  _op_to_func = {
      OpTypes.ADD: tf.math.add,
      OpTypes.CAST: tf.cast,
      OpTypes.BIAS_ADD: tf.nn.bias_add,
      OpTypes.IDENTITY: tf.identity,
      OpTypes.LINEAR: activations.linear,
      OpTypes.RELU: activations.relu,
      OpTypes.SIGMOID: activations.sigmoid,
      OpTypes.TANH: activations.tanh,
      OpTypes.MULTIPLY: tf.math.multiply,
  }

  def __init__(self, graph, quantized=False):
    self._quantized = quantized

    # All entities translated from ops.
    self._entities = []

    # Node name -> index in self._entities
    self._name_to_entity = {}

    # Op type -> count
    self._op_count = {}

    for node in graph.nodes:
      entities = self.translate_op_to_entities(node.op)
      self._entities.extend(entities)
      # Post-order traversing in `self.translate_op_to_entities`,
      # the last element in returned list is the entity converted from node.op.
      self._name_to_entity[node.name] = len(self._entities) - 1

      node_entity = self.get_entity(node.name)
      for i, tensor in enumerate(node.out_tensors):
        # Use node entity's name as prefix
        tensor_entity = Entity('%s_%d' % (node_entity.name, i),
                               EntityTypes.Tensor)
        self._add_entity(tensor.name, tensor_entity)
        node_entity.outputs.append(tensor_entity)

    for node in graph.nodes:
      entity = self.get_entity(node.name)
      for i, tensor in enumerate(node.in_tensors):
        entity.inputs.append(self.get_entity(tensor.name))

  def _add_entity(self, name, entity):
    """
      Args:
        name: Node or Tensor's name.
        entity: An `Entity` object.
    """
    self._name_to_entity[name] = len(self._entities)
    self._entities.append(entity)

  def _get_tf_obj(self, op):
    if self._quantized:
      obj = quant_utils.get_quant_module(op.type, None)
      if obj:
        return obj, EntityTypes.Layer

    if op.type in self._op_to_layer:
      obj = self._op_to_layer[op.type]
      entity_type = EntityTypes.Layer
    elif op.type in self._op_to_func:
      obj = self._op_to_func[op.type]
      entity_type = EntityTypes.Function
    elif op.type == OpTypes.RNN_LAYER:
      obj = op.attr['layer_class']
      entity_type = EntityTypes.Layer
    else:
      raise NotImplementedError("Unable to rewrite operation '{}'".format(
          op.type))
    return obj, entity_type

  def translate_op_to_entities(self, op):
    """Given an `Operation`, traverse its configs and find all `Operation`
    items, then:
    1. Convert these ops to `Entity` objects
    2. Set these entities back to configs to replace original ops.

    Arguments:
      op: An `Operation` object.

    Return:
      A list of entities extracted from the given op.
      The entities in the list are in post-order.

    Example:
      Given: RNN -> config('cell') -> StackedRNNCells -> config('cells') -> [LSTMCell0, LSTMCell1]
      Return: [Entity(LSTMCell0), Entity(LSTMCell1), Entity(StackedRNNCells)]
    """

    entities = []
    for name in op.configs:
      config_value = op.get_config(name)
      config_values = generic_utils.to_list(config_value)

      # Convert Operation to Entity.
      converted_values = []
      for value in config_values:
        if isinstance(value, ops.Operation):
          ents = self.translate_op_to_entities(value)
          converted_values.append(ents[-1])
          entities.extend(ents)
        else:
          converted_values.append(value)

      # Keep as the original type.
      if isinstance(config_value, tuple):
        config_value = tuple(converted_values)
      elif len(converted_values) <= 1:
        config_value = converted_values[0]
      else:
        config_value = converted_values
      op.set_config(name, config_value)

    entities.append(self._to_entity(op))
    # print('translate_op_to_entities:', [e.name for e in entities])
    return entities

  def _to_entity(self, op):
    obj, entity_type = self._get_tf_obj(op)
    init_func = obj.__init__ if entity_type == EntityTypes.Layer else obj

    argspec = tf_inspect.getfullargspec(init_func)
    default_count = 0 if not argspec.defaults else len(argspec.defaults)
    required_count = len(argspec.args) - default_count

    # Keep args order same as the original signature.
    arg_to_value = OrderedDict()
    args = []
    for arg in argspec.args[:required_count]:
      if arg == 'self':
        continue
      if arg not in op.configs:
        # Treat non-config args as inputs.
        #args.append('{inputs}')
        arg_to_value[_INPUT_ARG_PREFIX + arg] = None
      else:
        arg_to_value[arg] = op.get_config(arg)

    if default_count:
      for arg, value in zip(argspec.args[-default_count:], argspec.defaults):
        # Skip args that have default values.
        if arg not in op.configs or op.get_config(arg) == value:
          continue
        arg_to_value[arg] = op.get_config(arg)
    return Entity(self._unique_entity_name(op), entity_type, obj, arg_to_value)

  def _unique_entity_name(self, op):
    count = self._op_count.get(op.type, 0)
    self._op_count[op.type] = count + 1
    return '%s%d' % (op.type, count)

  def get_entity(self, name):
    index = self._name_to_entity[name]
    return self._entities[index]

  @property
  def entities(self):
    return self._entities

def _obj_modules(object):
  module = tf_inspect.getmodule(object)
  pkg, name = module.__name__.rsplit('.', 1)
  return (module, pkg, name)



#
# Copyright 2019 Xilinx Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import tensorflow as tf
from collections import OrderedDict

from tensorflow.python.keras import activations
from tensorflow.python.keras import layers as keras_layers
from tensorflow.python.util import nest
from tensorflow.python.util import tf_inspect

from tf_nndct import layers as nndct_layers
from tf_nndct.graph import OpTypes
from tf_nndct.graph import ops
from tf_nndct.graph import utils
from tf_nndct.quantization import utils as quant_utils
from tf_nndct.utils import generic_utils
from tf_nndct.utils import registry

_INPUT_ARG_PREFIX = '%input%_'

# call(self, input_0, input_1, ...)
_CALL_ARG_TEMPLATE = 'input_%d'

class Printer(object):

  def __init__(self, init_indent_level=0, indent_length=2):
    self._indent_level = init_indent_level
    self._init_indent_level = init_indent_level
    self._indent_length = indent_length
    self._start_of_line = False
    self._strs = []

  def indent(self):
    self._indent_level += 1

  def outdent(self):
    if self._indent_level == 0 or self._indent_level < self._init_indent_level:
      raise RuntimeError("outdent() without matching indent()")
    self._indent_level -= 1

  def current_indentation(self):
    return self._indent_length * self._indent_level

  def newline(self):
    self.print('\n')

  def print(self, text):
    text = str(text)
    if self._indent_level > 0:
      pos = 0
      for i, c in enumerate(text):
        if c == '\n':
          self._append(text[pos:i + 1])
          pos = i + 1
          self._start_of_line = True
      self._append(text[pos:])
    else:
      self._append(text)
      if text[-1] == '\n':
        self._start_of_line = True

  def get(self):
    return ''.join(self._strs)

  def _append(self, text):
    if not text:
      return
    if self._start_of_line:
      self._start_of_line = False
      self._append_indent()

    self._strs.append(text)

  def _append_indent(self):
    if self._indent_level == 0:
      return

    size = self.current_indentation()
    self._strs.append(size * ' ')

class ModuleClassWriter(object):

  def __init__(self, graph, class_spec):
    self._class_spec = class_spec

    # alias for imported module to be written.
    self._module_alias = {}

    self._graph = graph
    # TODO(yuwang): Remove these two members ?
    self._input_signature = graph.input_signature
    self._structured_output_tensors = graph.structured_output_tensors

    self._translator = GraphTranslator(graph, self._class_spec.quantized)
    self._printer = Printer()

  def write(self, filepath):
    self._printer.print('# Generated by NNDCT. Do not edit!')
    self._printer.newline()

    self.write_imports()
    self.write_class_def()

    generic_utils.mkdir_if_not_exist(os.path.dirname(filepath))
    with open(filepath, 'w') as f:
      f.write(self._printer.get())
      f.flush()
      os.fsync(f.fileno())

    layer_to_node = {}
    for node in self._graph.nodes:
      entity = self._translator.get_entity(node.name)
      if entity.need_init:
        layer_to_node[entity.name] = node

    return layer_to_node

  def write_imports(self):
    objects = [self._class_spec.base]
    for entity in self._translator.entities:
      if not entity.obj:
        continue
      objects.append(entity.obj)

    imports = set()
    for obj in objects:
      imports.add(self._generate_obj_import(obj))

    imports = sorted(imports)
    for impt in imports:
      self._printer.print(impt)
      self._printer.newline()
    self._printer.newline()

  def _generate_obj_import(self, object):
    module, pkg, name = _obj_modules(object)
    if module in self._module_alias:
      alias = self._module_alias[module]
    else:
      alias = name
      while alias in self._module_alias.values():
        alias = '_' + alias
      self._module_alias[module] = alias

    if alias == name:
      return 'from {} import {}'.format(pkg, name)
    else:
      return 'from {} import {} as {}'.format(pkg, name, alias)

  def write_class_def(self):
    module, _, _ = _obj_modules(self._class_spec.base)
    self._printer.print('class {cls_name}({base_cls}):\n'.format(
        cls_name=self._class_spec.name,
        base_cls='.'.join(
            [self._module_alias[module], self._class_spec.base.__name__])))
    self._printer.indent()
    self.write_init()
    self.write_call_fn(self._class_spec.call_fn_name)
    self._printer.outdent()

  def write_init(self):
    self._printer.print('def __init__(self):')
    self._printer.newline()
    self._printer.indent()

    self._printer.print('super({}, self).__init__(name={})'.format(
        self._class_spec.name, utils.stringfy_to_write(self._class_spec.name)))
    self._printer.newline()

    for entity in self._translator.entities:
      if not entity.need_init:
        continue

      init_str = self._init_string(entity)
      self._printer.print('self.{} = {}'.format(entity.name, init_str, entity))
      self._printer.newline()

    self._printer.newline()
    self._printer.outdent()

  def _init_string(self, entity):
    args = []
    input_index = 0
    for arg, value in entity.args.items():
      # Use input entity name to fill in placeholder.
      if arg.startswith(_INPUT_ARG_PREFIX):
        args.append(entity.inputs[input_index].name)
        input_index += 1
        continue

      values = generic_utils.to_list(value)
      value_strs = []
      for value in values:
        if isinstance(value, str):
          # Use a quote to wrap a string value.
          value = ''.join(['\'', value, '\''])
        elif isinstance(value, Entity):
          value = 'self.' + value.name
        else:
          pass
        value_strs.append(str(value))
      if len(value_strs) <= 1:
        args.append('{}={}'.format(arg, value_strs[0]))
      else:
        args.append('{}={}'.format(arg, '[' + ', '.join(value_strs) + ']'))

    module, _, _ = _obj_modules(entity.obj)
    init_str = '{module}.{object}({args})'.format(
        module=self._module_alias[module],
        object=entity.obj.__name__,
        args=', '.join(args))
    return init_str

  def write_call_fn(self, fn_name):
    call_args = ['self']
    for i in range(len(self._input_signature)):
      call_args.append(_CALL_ARG_TEMPLATE % i)
    self._printer.print('def {}({}):'.format(fn_name, ', '.join(call_args)))
    self._printer.newline()
    self._printer.indent()

    placeholders = []
    for node in self._graph.nodes:
      if node.op.type == OpTypes.INPUT:
        assert len(node.out_tensors) == 1
        placeholders.append(node)
    # [args_0, args_1_1, args_1] -> [args_0, args_1, args_1_1]
    placeholders = sorted(placeholders, key=lambda x: x.name)

    flattened_input_signature = nest.flatten(self._input_signature)
    assert len(placeholders) == len(flattened_input_signature)

    # TODO(yuwang): Maybe use more solid way to map ?
    # We sort the placeholders by their names and map the flattened input
    # signature in the sorted order. This method works now but relies on the
    # tensorflow's mechanism for generating placeholder names from input
    # argument. Specifically, the first input is named args_0, the second
    # input is named args_1 and so on. If an input is a type of sequence,
    # then the input is flattened to a list and each element in the list
    # is named by args_[input_index]_[element_index_in_the_list].
    # For example, we have a call function and the signature is like:
    #   def call(self, input, state)
    # The first argument 'input' is a tensor and TF will generate a
    # placeholder for it named args_0. The second argument 'state' is a list
    # of tensor, then TF will generate a placeholder for each element in the
    # list and the names of the placeholders are arg_1 and arg_1_1. In the
    # end, we got got three placeholders [args_0, args_1, args_1_1]. In a
    # nndct graph, the order of the placeholders may be [args_0, args_1_1,
    # args_1], that is why we sort the nodes by their names before do mapping.
    # Example:
    # placeholders: ['args_0', 'args_1', 'args_1_1']
    # args: (TensorSpec(shape=(1, 10), dtype=tf.float32, name='args_0'), [TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_1/0'), TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_1/1')])
    # args_to_placeholder = {
    #     'args_0': 'args_0',
    #     'args_1/0': 'args_1',
    #     'args_1/1': 'args_1_1',
    # }
    #
    # Example:
    # placeholders: ['args_0', 'args_0_1', 'args_0_2']
    # args: [{'y': TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/0/y'),
    #         'x': TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/0/x')},
    #        TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/1')]
    # args_to_placeholder = {
    #     'args_0/0/y': 'arg_0',
    #     'args_0/0/x': 'arg_0_1',
    #     'args_0/1': 'arg_0_2',
    # }

    args_to_placeholder = {}
    for i, arg in enumerate(flattened_input_signature):
      args_to_placeholder[arg.name] = placeholders[i].name

    def splice_retriving_path(getter, key):
      key = utils.stringfy_to_write(key)

      if getter == '[]':
        return '[{}]'.format(key)
      return '{}{}'.format(getter, key)

    def arg_retriving_path(arg, path=()):
      """
      Get retriving path of an argument.

      Args:
        arg: The input signature of an argument.

      Yield:

      See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/nest.py::_yield_sorted_items
      """
      if not nest.is_sequence(arg):
        yield path
      elif isinstance(arg, nest._collections_abc.Mapping):
        for key in nest._sorted(arg):
          for res in arg_retriving_path(arg[key], path + (('[]', key),)):
            yield res
      elif nest._is_attrs(arg):
        for item in nest._get_attrs_items(arg):
          for res in arg_retriving_path(item[1], path + (('.', item[0]),)):
            yield res
      elif nest._is_namedtuple(arg):
        for field in arg._fields:
          for res in arg_retriving_path(
              getattr(arg, field), path + (('.', field),)):
            yield res
      # Doesn't support composite_tensor comprared with _yield_sorted_items.
      elif nest._is_type_spec(arg):
        # Note: to allow CompositeTensors and their TypeSpecs to have matching
        # structures, we need to use the same key string here.
        for res in arg_retriving_path(arg._component_specs,
                                      path + (('.', arg.value_type.__name__),)):
          yield res
      else:
        for item in enumerate(arg):
          for res in arg_retriving_path(item[1], path + (('[]', item[0]),)):
            yield res

    for i, arg in enumerate(self._input_signature):
      for path in list(arg_retriving_path(arg)):
        # [('[]' 0), ('[]', 'x')]
        arg_name_path = ['args_%d' % i]
        retriving_path = []
        for getter, key in path:
          arg_name_path.append(key)
          retriving_path.append(splice_retriving_path(getter, key))
        arg_name = '/'.join(str(s) for s in arg_name_path)
        retriving = ''.join(r for r in retriving_path)
        if arg_name not in args_to_placeholder:
          continue

        placeholder_name = args_to_placeholder[arg_name]
        entity = self._translator.get_entity(placeholder_name)
        call_arg = _CALL_ARG_TEMPLATE % i
        #self._printer.print('{} = {}{}'.format(
        #    entity.outputs[0].name, call_arg, retriving))
        #self._printer.newline()
        entity.inputs.append(
            Entity('{}{}'.format(call_arg, retriving), EntityTypes.Tensor))

    for node in self._graph.nodes:
      #if node.op.type == OpTypes.INPUT:
      #  continue
      self._printer.print(self._call_string(node))
      self._printer.newline()

    output_tensors = [
        t.name for t in nest.flatten(self._structured_output_tensors)
    ]
    returns = [self._translator.get_entity(t) for t in output_tensors]
    if len(returns) == 1:
      returns = returns[0]
    else:
      returns = nest.pack_sequence_as(self._structured_output_tensors, returns)

    def return_str(returns):
      return_strs = []
      if isinstance(returns, Entity):
        s = returns.name
      elif isinstance(returns, list):
        for ret in returns:
          return_strs.append(return_str(ret))
        s = ''.join(['[', ', '.join(return_strs), ']'])
      elif isinstance(returns, tuple):
        for ret in returns:
          return_strs.append(return_str(ret))
        s = ''.join(['(', ', '.join(return_strs), ')'])
      elif isinstance(returns, dict):
        for key in returns:
          return_strs.append('{}: {}'.format(
              utils.stringfy_to_write(key), ''.join(return_str(returns[key]))))
        s = ''.join(['{', ', '.join(return_strs), '}'])
      else:
        raise NotImplementedError(
            'Can not rewrite return object of type '.format(type(returns)))
      return s

    self._printer.print('return %s' % return_str(returns))
    self._printer.newline()
    self._printer.outdent()

  def _call_string(self, node):
    entity = self._translator.get_entity(node.name)
    input_names = [i.name for i in entity.inputs]
    output_names = [o.name for o in entity.outputs]
    #if not input_names:
    #  # TODO(yuwang): Support multiple inputs
    #  input_names.append('inputs')

    input_str = ', '.join(input_names)
    output_str = ', '.join(output_names)

    entity = self._translator.get_entity(node.name)
    # Layer type
    if entity.need_init:
      forward_template = '{outputs} = self.{entity}({inputs}) # {node}'
      args_dict = {
          'outputs': output_str,
          'entity': entity.name,
          'inputs': input_str,
          'node': node.name,
      }
      forward_str = forward_template.format(**args_dict)
    else:
      # Function type
      init_str = self._init_string(entity)
      forward_str = '{} = {} # {}'.format(output_str, init_str, node.name)
    return forward_str

class EntityTypes(object):
  Layer = 'layer'
  Function = 'function'
  Tensor = 'tensor'

class Entity(object):

  def __init__(self, name, type, obj=None, args=None):
    self.name = name

    # One of `EntityTypes`.
    self.type = type

    # For EntityTypes.Layer and EntityTypes.Function,
    # obj is the corresponding layer class or function.
    # For EntityTypes.Tensor, obj is None
    self.obj = obj

    self.args = args

    # I/O entities
    self.inputs = []
    self.outputs = []

  def __str__(self):
    return self.name

  @property
  def need_init(self):
    return self.type == EntityTypes.Layer

class GraphTranslator(object):
  """Translates graph components to `Entity` objects.
  Specifically, all `Operation` and `Tensor` objects will be converted to
  `Entity` objects, not including `Node`. This is because a node may contain
  multiple ops in a nested form, so the conversion is performed at the op
  level, not the node level.
  """
  _op_to_layer = {
      OpTypes.BIDIRECTIONAL_RNN: keras_layers.Bidirectional,
      OpTypes.CONV1D: keras_layers.Conv1D,
      OpTypes.DENSE: keras_layers.Dense,
      OpTypes.EMBEDDING: keras_layers.Embedding,
      OpTypes.GRU: keras_layers.GRU,
      OpTypes.INPUT: nndct_layers.Identity,
      OpTypes.LSTM: nndct_layers.LSTM,
      OpTypes.LSTM_CELL: nndct_layers.LSTMCell,
      OpTypes.MAX_POOL1D: keras_layers.MaxPooling1D,
      OpTypes.RNN: keras_layers.RNN,
      OpTypes.SIMPLE_RNN: keras_layers.SimpleRNN,
      OpTypes.STACKED_RNN_CELLS: keras_layers.StackedRNNCells,
  }

  _op_to_func = {
      OpTypes.ADD: tf.math.add,
      OpTypes.CAST: tf.cast,
      OpTypes.BIAS_ADD: tf.nn.bias_add,
      OpTypes.IDENTITY: tf.identity,
      OpTypes.LINEAR: activations.linear,
      OpTypes.RELU: activations.relu,
      OpTypes.SIGMOID: activations.sigmoid,
      OpTypes.TANH: activations.tanh,
      OpTypes.MULTIPLY: tf.math.multiply,
  }

  def __init__(self, graph, quantized=False):
    self._quantized = quantized

    # All entities translated from ops.
    self._entities = []

    # Node name -> index in self._entities
    self._name_to_entity = {}

    # Op type -> count
    self._op_count = {}

    for node in graph.nodes:
      entities = self.translate_op_to_entities(node.op)
      self._entities.extend(entities)
      # Post-order traversing in `self.translate_op_to_entities`,
      # the last element in returned list is the entity converted from node.op.
      self._name_to_entity[node.name] = len(self._entities) - 1

      node_entity = self.get_entity(node.name)
      for i, tensor in enumerate(node.out_tensors):
        # Use node entity's name as prefix
        tensor_entity = Entity('%s_%d' % (node_entity.name, i),
                               EntityTypes.Tensor)
        self._add_entity(tensor.name, tensor_entity)
        node_entity.outputs.append(tensor_entity)

    for node in graph.nodes:
      entity = self.get_entity(node.name)
      for i, tensor in enumerate(node.in_tensors):
        entity.inputs.append(self.get_entity(tensor.name))

  def _add_entity(self, name, entity):
    """
      Args:
        name: Node or Tensor's name.
        entity: An `Entity` object.
    """
    self._name_to_entity[name] = len(self._entities)
    self._entities.append(entity)

  def _get_tf_obj(self, op):
    if self._quantized:
      obj = quant_utils.get_quant_module(op.type, None)
      if obj:
        return obj, EntityTypes.Layer

    if op.type in self._op_to_layer:
      obj = self._op_to_layer[op.type]
      entity_type = EntityTypes.Layer
    elif op.type in self._op_to_func:
      obj = self._op_to_func[op.type]
      entity_type = EntityTypes.Function
    elif op.type == OpTypes.RNN_LAYER:
      obj = op.attr['layer_class']
      entity_type = EntityTypes.Layer
    else:
      raise NotImplementedError("Unable to rewrite operation '{}'".format(
          op.type))
    return obj, entity_type

  def translate_op_to_entities(self, op):
    """Given an `Operation`, traverse its configs and find all `Operation`
    items, then:
    1. Convert these ops to `Entity` objects
    2. Set these entities back to configs to replace original ops.

    Arguments:
      op: An `Operation` object.

    Return:
      A list of entities extracted from the given op.
      The entities in the list are in post-order.

    Example:
      Given: RNN -> config('cell') -> StackedRNNCells -> config('cells') -> [LSTMCell0, LSTMCell1]
      Return: [Entity(LSTMCell0), Entity(LSTMCell1), Entity(StackedRNNCells)]
    """

    entities = []
    for name in op.configs:
      config_value = op.get_config(name)
      config_values = generic_utils.to_list(config_value)

      # Convert Operation to Entity.
      converted_values = []
      for value in config_values:
        if isinstance(value, ops.Operation):
          ents = self.translate_op_to_entities(value)
          converted_values.append(ents[-1])
          entities.extend(ents)
        else:
          converted_values.append(value)

      # Keep as the original type.
      if isinstance(config_value, tuple):
        config_value = tuple(converted_values)
      elif len(converted_values) <= 1:
        config_value = converted_values[0]
      else:
        config_value = converted_values
      op.set_config(name, config_value)

    entities.append(self._to_entity(op))
    # print('translate_op_to_entities:', [e.name for e in entities])
    return entities

  def _to_entity(self, op):
    obj, entity_type = self._get_tf_obj(op)
    init_func = obj.__init__ if entity_type == EntityTypes.Layer else obj

    argspec = tf_inspect.getfullargspec(init_func)
    default_count = 0 if not argspec.defaults else len(argspec.defaults)
    required_count = len(argspec.args) - default_count

    # Keep args order same as the original signature.
    arg_to_value = OrderedDict()
    args = []
    for arg in argspec.args[:required_count]:
      if arg == 'self':
        continue
      if arg not in op.configs:
        # Treat non-config args as inputs.
        #args.append('{inputs}')
        arg_to_value[_INPUT_ARG_PREFIX + arg] = None
      else:
        arg_to_value[arg] = op.get_config(arg)

    if default_count:
      for arg, value in zip(argspec.args[-default_count:], argspec.defaults):
        # Skip args that have default values.
        if arg not in op.configs or op.get_config(arg) == value:
          continue
        arg_to_value[arg] = op.get_config(arg)
    return Entity(self._unique_entity_name(op), entity_type, obj, arg_to_value)

  def _unique_entity_name(self, op):
    count = self._op_count.get(op.type, 0)
    self._op_count[op.type] = count + 1
    return '%s%d' % (op.type, count)

  def get_entity(self, name):
    index = self._name_to_entity[name]
    return self._entities[index]

  @property
  def entities(self):
    return self._entities

def _obj_modules(object):
  module = tf_inspect.getmodule(object)
  pkg, name = module.__name__.rsplit('.', 1)
  return (module, pkg, name)

# # Copyright 2019 Xilinx Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # call(self, input_0, input_1, ...) # alias for imported module to be written. # TODO(yuwang): Remove these two members ? # Use input entity name to fill in placeholder. # Use a quote to wrap a string value. # [args_0, args_1_1, args_1] -> [args_0, args_1, args_1_1] # TODO(yuwang): Maybe use more solid way to map ? # We sort the placeholders by their names and map the flattened input # signature in the sorted order. This method works now but relies on the # tensorflow's mechanism for generating placeholder names from input # argument. Specifically, the first input is named args_0, the second # input is named args_1 and so on. If an input is a type of sequence, # then the input is flattened to a list and each element in the list # is named by args_[input_index]_[element_index_in_the_list]. # For example, we have a call function and the signature is like: #   def call(self, input, state) # The first argument 'input' is a tensor and TF will generate a # placeholder for it named args_0. The second argument 'state' is a list # of tensor, then TF will generate a placeholder for each element in the # list and the names of the placeholders are arg_1 and arg_1_1. In the # end, we got got three placeholders [args_0, args_1, args_1_1]. In a # nndct graph, the order of the placeholders may be [args_0, args_1_1, # args_1], that is why we sort the nodes by their names before do mapping. # Example: # placeholders: ['args_0', 'args_1', 'args_1_1'] # args: (TensorSpec(shape=(1, 10), dtype=tf.float32, name='args_0'), [TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_1/0'), TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_1/1')]) # args_to_placeholder = { #     'args_0': 'args_0', #     'args_1/0': 'args_1', #     'args_1/1': 'args_1_1', # } # # Example: # placeholders: ['args_0', 'args_0_1', 'args_0_2'] # args: [{'y': TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/0/y'), #         'x': TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/0/x')}, #        TensorSpec(shape=(1, 5), dtype=tf.float32, name='args_0/1')] # args_to_placeholder = { #     'args_0/0/y': 'arg_0', #     'args_0/0/x': 'arg_0_1', #     'args_0/1': 'arg_0_2', # } Get retriving path of an argument.        Args:         arg: The input signature of an argument.        Yield:        See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/util/nest.py::_yield_sorted_items # Doesn't support composite_tensor comprared with _yield_sorted_items. # Note: to allow CompositeTensors and their TypeSpecs to have matching # structures, we need to use the same key string here. # [('[]' 0), ('[]', 'x')] #self._printer.print('{} = {}{}'.format( #    entity.outputs[0].name, call_arg, retriving)) #self._printer.newline() #if node.op.type == OpTypes.INPUT: #  continue #if not input_names: #  # TODO(yuwang): Support multiple inputs #  input_names.append('inputs') # Layer type # {node}' # Function type # {}'.format(output_str, init_str, node.name) # One of `EntityTypes`. # For EntityTypes.Layer and EntityTypes.Function, # obj is the corresponding layer class or function. # For EntityTypes.Tensor, obj is None # I/O entities Translates graph components to `Entity` objects.   Specifically, all `Operation` and `Tensor` objects will be converted to   `Entity` objects, not including `Node`. This is because a node may contain   multiple ops in a nested form, so the conversion is performed at the op   level, not the node level. # All entities translated from ops. # Node name -> index in self._entities # Op type -> count # Post-order traversing in `self.translate_op_to_entities`, # the last element in returned list is the entity converted from node.op. # Use node entity's name as prefix Args:         name: Node or Tensor's name.         entity: An `Entity` object. Given an `Operation`, traverse its configs and find all `Operation`     items, then:     1. Convert these ops to `Entity` objects     2. Set these entities back to configs to replace original ops.      Arguments:       op: An `Operation` object.      Return:       A list of entities extracted from the given op.       The entities in the list are in post-order.      Example:       Given: RNN -> config('cell') -> StackedRNNCells -> config('cells') -> [LSTMCell0, LSTMCell1]       Return: [Entity(LSTMCell0), Entity(LSTMCell1), Entity(StackedRNNCells)] # Convert Operation to Entity. # Keep as the original type. # print('translate_op_to_entities:', [e.name for e in entities]) # Keep args order same as the original signature. # Treat non-config args as inputs. #args.append('{inputs}') # Skip args that have default values.