# -*- coding: utf-8 -*-
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Utilities for handling Keras model in graph plugin.

Two canonical types of Keras model are Functional and Sequential.
A model can be serialized as JSON and deserialized to reconstruct a model.
This utility helps with dealing with the serialized Keras model.

They have distinct structures to the configurations in shapes below:
Functional:
  config
    name: Name of the model. If not specified, it is 'model' with
          an optional suffix if there are more than one instance.
    input_layers: Keras.layers.Inputs in the model.
    output_layers: Layer names that are outputs of the model.
    layers: list of layer configurations.
      layer: [*]
        inbound_nodes: inputs to this layer.

Sequential:
  config
    name: Name of the model. If not specified, it is 'sequential' with
          an optional suffix if there are more than one instance.
    layers: list of layer configurations.
      layer: [*]

[*]: Note that a model can be a layer.
Please refer to https://github.com/tensorflow/tfjs-layers/blob/master/src/keras_format/model_serialization.ts
for more complete definition.
"""
from tensorboard.compat.proto.graph_pb2 import GraphDef
from tensorboard.compat.tensorflow_stub import dtypes


def _walk_layers(keras_layer):
  """Walks the nested keras layer configuration in preorder.
  Args:
    keras_layer: Keras configuration from model.to_json.

  Yields:
    A tuple of (name_scope, layer_config).
    name_scope: a string representing a scope name, similar to that of tf.name_scope.
    layer_config: a dict representing a Keras layer configuration.
  """
  yield ('', keras_layer)
  if keras_layer.get('config').get('layers'):
    name_scope = keras_layer.get('config').get('name')
    for layer in keras_layer.get('config').get('layers'):
      for (sub_name_scope, sublayer) in _walk_layers(layer):
        sub_name_scope = '%s/%s' % (
            name_scope, sub_name_scope) if sub_name_scope else name_scope
        yield (sub_name_scope, sublayer)


def _scoped_name(name_scope, node_name):
  """Returns scoped name for a node as a string in the form '<scope>/<node name>'.

  Args:
    name_scope: a string representing a scope name, similar to that of tf.name_scope.
    node_name: a string representing the current node name.

  Returns
    A string representing a scoped name.
  """
  if name_scope:
    return '%s/%s' % (name_scope, node_name)
  return node_name


def _is_model(layer):
  """Returns True if layer is a model.

  Args:
    layer: a dict representing a Keras model configuration.

  Returns:
    bool: True if layer is a model.
  """
  return layer.get('config').get('layers') is not None


def _norm_to_list_of_layers(maybe_layers):
  """Normalizes to a list of layers.

  Args:
    maybe_layers: A list of data[1] or a list of list of data.

  Returns:
    List of list of data.

  [1]: A Functional model has fields 'inbound_nodes' and 'output_layers' which can
  look like below:
  - ['in_layer_name', 0, 0]
  - [['in_layer_is_model', 1, 0], ['in_layer_is_model', 1, 1]]
  The data inside the list seems to describe [name, size, index].
  """
  return (maybe_layers if isinstance(maybe_layers[0], (list,))
          else [maybe_layers])


def _update_dicts(name_scope,
                  model_layer,
                  input_to_in_layer,
                  model_name_to_output,
                  prev_node_name):
  """Updates input_to_in_layer, model_name_to_output, and prev_node_name
  based on the model_layer.

  Args:
    name_scope: a string representing a scope name, similar to that of tf.name_scope.
    model_layer: a dict representing a Keras model configuration.
    input_to_in_layer: a dict mapping Keras.layers.Input to inbound layer.
    model_name_to_output: a dict mapping Keras Model name to output layer of the model.
    prev_node_name: a string representing a previous, in sequential model layout,
                    node name.

  Returns:
    A tuple of (input_to_in_layer, model_name_to_output, prev_node_name).
    input_to_in_layer: a dict mapping Keras.layers.Input to inbound layer.
    model_name_to_output: a dict mapping Keras Model name to output layer of the model.
    prev_node_name: a string representing a previous, in sequential model layout,
                    node name.
  """
  layer_config = model_layer.get('config')
  if not layer_config.get('layers'):
    raise ValueError('layer is not a model.')

  node_name = _scoped_name(name_scope, layer_config.get('name'))
  input_layers = layer_config.get('input_layers')
  output_layers = layer_config.get('output_layers')
  inbound_nodes = model_layer.get('inbound_nodes')

  is_functional_model = bool(input_layers and output_layers)
  # In case of [1] and the parent model is functional, current layer
  # will have the 'inbound_nodes' property.
  is_parent_functional_model = bool(inbound_nodes)

  if is_parent_functional_model and is_functional_model:
    for (input_layer, inbound_node) in zip(input_layers, inbound_nodes):
      input_layer_name = _scoped_name(node_name, input_layer)
      inbound_node_name = _scoped_name(name_scope, inbound_node[0])
      input_to_in_layer[input_layer_name] = inbound_node_name
  elif is_parent_functional_model and not is_functional_model:
    # Sequential model can take only one input. Make sure inbound to the
    # model is linked to the first layer in the Sequential model.
    prev_node_name = _scoped_name(name_scope, inbound_nodes[0][0][0])
  elif not is_parent_functional_model and prev_node_name and is_functional_model:
    assert len(input_layers) == 1, (
        'Cannot have multi-input Functional model when parent model '
        'is not Functional. Number of input layers: %d' % len(input_layer))
    input_layer = input_layers[0]
    input_layer_name = _scoped_name(node_name, input_layer)
    input_to_in_layer[input_layer_name] = prev_node_name

  if is_functional_model and output_layers:
    layers = _norm_to_list_of_layers(output_layers)
    layer_names = [_scoped_name(node_name, layer[0]) for layer in layers]
    model_name_to_output[node_name] = layer_names
  else:
    last_layer = layer_config.get('layers')[-1]
    last_layer_name = last_layer.get('config').get('name')
    output_node = _scoped_name(node_name, last_layer_name)
    model_name_to_output[node_name] = [output_node]
  return (input_to_in_layer, model_name_to_output, prev_node_name)


def keras_model_to_graph_def(keras_layer):
  """Returns a GraphDef representation of the Keras model in a dict form.

  Note that it only supports models that implemented to_json().

  Args:
    keras_layer: A dict from Keras model.to_json().

  Returns:
    A GraphDef representation of the layers in the model.
  """
  input_to_layer = {}
  model_name_to_output = {}
  g = GraphDef()

  # Sequential model layers do not have a field "inbound_nodes" but
  # instead are defined implicitly via order of layers.
  prev_node_name = None

  for (name_scope, layer) in _walk_layers(keras_layer):
    if _is_model(layer):
      (input_to_layer, model_name_to_output, prev_node_name) = _update_dicts(
          name_scope, layer, input_to_layer, model_name_to_output, prev_node_name)
      continue

    layer_config = layer.get('config')
    node_name = _scoped_name(name_scope, layer_config.get('name'))

    node_def = g.node.add()
    node_def.name = node_name

    if layer.get('class_name') is not None:
      keras_cls_name = layer.get('class_name').encode('ascii')
      node_def.attr['keras_class'].s = keras_cls_name

    if layer_config.get('dtype') is not None:
      tf_dtype = dtypes.as_dtype(layer_config.get('dtype'))
      node_def.attr['dtype'].type = tf_dtype.as_datatype_enum

    if layer.get('inbound_nodes') is not None:
      for maybe_inbound_node in layer.get('inbound_nodes'):
        inbound_nodes = _norm_to_list_of_layers(maybe_inbound_node)
        for [name, size, index, _] in inbound_nodes:
          inbound_name = _scoped_name(name_scope, name)
          # An input to a layer can be output from a model. In that case, the name
          # of inbound_nodes to a layer is a name of a model. Remap the name of the
          # model to output layer of the model. Also, since there can be multiple
          # outputs in a model, make sure we pick the right output_layer from the model.
          inbound_node_names = model_name_to_output.get(
              inbound_name, [inbound_name])
          node_def.input.append(inbound_node_names[index])
    elif prev_node_name is not None:
      node_def.input.append(prev_node_name)

    if node_name in input_to_layer:
      node_def.input.append(input_to_layer.get(node_name))

    prev_node_name = node_def.name

  return g
