# Copyright 2018 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.
# ==============================================================================
"""Tools for deserializing `Function`s."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import re

from tensorflow.core.framework import function_pb2
from tensorflow.python.eager import def_function
from tensorflow.python.eager import function as function_lib
from tensorflow.python.framework import function_def_to_graph as function_def_lib
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.util import compat
from tensorflow.python.util import nest


def _is_tensor(t):
  return isinstance(t, (ops.Tensor, resource_variable_ops.ResourceVariable))


def _inputs_compatible(args, stored_inputs):
  """Checks whether function arguments are compatible with parameters."""
  if len(args) != len(stored_inputs):
    return False

  for arg, stored_input in zip(args, stored_inputs):
    if not function_lib.is_same_structure(arg, stored_input):
      return False

    flattened_arg = nest.flatten(arg)
    flattened_stored_input = nest.flatten(stored_input)

    for a, b in zip(flattened_arg, flattened_stored_input):
      if _is_tensor(a):
        if not isinstance(b, tensor_spec.TensorSpec):
          return False
        if a.dtype != b.dtype or not b.shape.is_compatible_with(a.shape):
          return False
      else:
        if a != b:
          return False
  return True


def _deserialize_function_spec(function_spec_proto, coder):
  """Deserialize a FunctionSpec object from its proto representation."""
  fullargspec = coder.decode_proto(function_spec_proto.fullargspec)
  is_method = function_spec_proto.is_method
  args_to_prepend = coder.decode_proto(function_spec_proto.args_to_prepend)
  kwargs_to_include = coder.decode_proto(function_spec_proto.kwargs_to_include)
  input_signature = coder.decode_proto(function_spec_proto.input_signature)
  return function_lib.FunctionSpec(fullargspec, is_method, args_to_prepend,
                                   kwargs_to_include, input_signature)


# TODO(allenl): The fact that we can't derive ConcreteFunction calling
# conventions from the serialized input spec right now is unfortunate. Merging
# these would be good, maybe by adding TensorSpec names to cache keys so renamed
# keyword arguments would yield different ConcreteFunctions.
def setup_bare_concrete_function(saved_bare_concrete_function,
                                 concrete_functions):
  """Makes a restored bare concrete function callable."""
  # Bare concrete functions accept only flat lists of Tensors with unique
  # names.
  concrete_function = concrete_functions[
      saved_bare_concrete_function.concrete_function_name]
  # pylint: disable=protected-access
  concrete_function._arg_keywords = (
      saved_bare_concrete_function.argument_keywords)
  concrete_function._num_positional_args = (
      saved_bare_concrete_function.allowed_positional_arguments)
  # pylint: enable=protected-access
  concrete_function.add_to_graph()
  return concrete_function


class RestoredFunction(def_function.Function):
  """Wrapper class for a function that has been restored from saved state.

  See `def_function.Function`.
  """

  def __init__(self, python_function, name, function_spec, concrete_functions):
    # TODO(mdan): We may enable autograph once exceptions are supported.
    super(RestoredFunction, self).__init__(
        python_function, name, autograph=False)
    self._concrete_functions = concrete_functions
    # This does not propagate to stateful and stateless functions of the
    # RestoredFunction, which will have seen only defunned
    # restored_function_body(*args, **kwargs). That's why we have to
    # canonicalize inputs inside restored_function_body.
    self._function_spec = function_spec

  def _list_all_concrete_functions_for_serialization(self):
    return self._concrete_functions


def recreate_function(saved_function, concrete_functions):
  """Creates a `Function` from a `SavedFunction`.

  Args:
    saved_function: `SavedFunction` proto.
    concrete_functions: map from function name to `ConcreteFunction`.

  Returns:
    A `Function`.
  """
  # TODO(andresp): Construct a `Function` with the cache populated
  # instead of creating a new `Function` backed by a Python layer to
  # glue things together. Current approach is nesting functions deeper for each
  # serialization cycle.

  coder = nested_structure_coder.StructureCoder()
  function_spec = _deserialize_function_spec(saved_function.function_spec,
                                             coder)

  def restored_function_body(*args, **kwargs):
    """Calls a restored function."""
    # TODO(allenl): Functions saved with input_signatures should revive with
    # input_signatures.
    try:
      canonicalized_inputs = function_spec.canonicalize_function_inputs(
          *args, **kwargs)
    except ValueError as e:
      raise ValueError(
          "Cannot canonicalize input args %r and kwargs %r. Error: %r." %
          (args, kwargs, e))

    debug_considered_signatures = []
    for concrete_function_name in saved_function.concrete_functions:
      function_obj = concrete_functions[concrete_function_name]
      canonicalized_original_inputs = (
          function_obj.graph.structured_input_signature)
      debug_considered_signatures.append(canonicalized_original_inputs)

      if _inputs_compatible(canonicalized_inputs,
                            canonicalized_original_inputs):
        flattened_inputs = nest.flatten(canonicalized_inputs)
        filtered_inputs = [t for t in flattened_inputs if _is_tensor(t)]

        result = function_obj._call_flat(filtered_inputs)  # pylint: disable=protected-access
        if isinstance(result, ops.Operation):
          return None
        return result

    raise AssertionError(
        "Could not find matching function to call for canonicalized inputs %r. "
        "Only existing signatures are %r."
        % (canonicalized_inputs, debug_considered_signatures))

  concrete_function_objects = []
  for concrete_function_name in saved_function.concrete_functions:
    concrete_function_objects.append(concrete_functions[concrete_function_name])

  return RestoredFunction(restored_function_body,
                          restored_function_body.__name__,
                          function_spec,
                          concrete_function_objects)


def load_function_def_library(library):
  """Load a set of functions as concrete functions without captured inputs.

  Functions names are manipulated during load such that they do not overlap
  with previously created ones.

  Args:
    library: FunctionDefLibrary proto message.

  Returns:
    Map of original function names in the library to instances of
    `ConcreteFunction` without captured inputs.

  Raises:
    ValueError: if functions dependencies have a cycle.
  """
  functions = {}

  for fdef in _sort_function_defs(library):
    copy = _fix_fdef(fdef, functions)

    func_graph = function_def_lib.function_def_to_graph(copy)
    for dep in _list_function_deps(fdef):
      functions[dep].add_to_graph(func_graph)
    func = function_lib.ConcreteFunction(func_graph)
    func.add_to_graph()

    functions[fdef.signature.name] = func

    # Also register the gradients in the current root context.
    with ops.init_scope():
      func._register_gradient()  # pylint: disable=protected-access

  return functions


def _sort_function_defs(library):
  """Return a topologic sort of FunctionDefs in a library."""
  edges = collections.defaultdict(list)
  in_count = collections.defaultdict(lambda: 0)

  for fdef in library.function:
    for dep in _list_function_deps(fdef):
      edges[dep].append(fdef.signature.name)
      in_count[fdef.signature.name] += 1

  ready = [
      fdef.signature.name
      for fdef in library.function
      if in_count[fdef.signature.name] == 0
  ]
  output = []
  while ready:
    node = ready.pop()
    output.append(node)
    for dest in edges[node]:
      in_count[dest] -= 1
      if not in_count[dest]:
        ready.append(dest)

  if len(output) != len(library.function):
    failed_to_resolve = sorted(set(in_count.keys()) - set(output))
    raise ValueError("There is a cyclic-dependency between functions. ",
                     "Could not resolve %r." % (failed_to_resolve,))

  reverse = {fdef.signature.name: fdef for fdef in library.function}
  return [reverse[x] for x in output]


def _fix_fdef(orig_fdef, functions):
  """Fixes a FunctionDef proto to be loaded in current context.

  In particular, when loading a function library into an eager context, one
  must rename the functions to avoid conflicts with existent functions.

  Args:
    orig_fdef: FunctionDef proto to fix. It is not modified.
    functions: map from function name to a ConcreteFunction instance.

  Returns:
    A fixed copy of the original FunctionDef.
  """
  fdef = function_pb2.FunctionDef()
  fdef.CopyFrom(orig_fdef)
  for node_def in fdef.node_def:
    if "_gradient_op_type" in node_def.attr:
      if node_def.op in ["StatefulPartitionedCall", "PartitionedCall"]:
        # TODO(andresp): This code assumes that the gradient registered for this
        # function call is the default gradient for the function and not a
        # custom one.
        fname = node_def.attr["f"].func.name
        node_def.attr["_gradient_op_type"].s = compat.as_bytes(
            functions[fname]._gradient_name)  # pylint: disable=protected-access
      else:
        logging.warning("Importing a function (%s) with ops with custom "
                        "gradients. Will likely fail if a gradient is "
                        "requested.", fdef.signature.name)
    for _, attr_value in node_def.attr.items():
      if attr_value.func.name:
        attr_value.func.name = functions[attr_value.func.name].name

    # TODO(b/124205571): Avoid accidental sharing and destruction of restored
    # resources. For now drop "shared_name" when loading functions to avoid
    # sharing.
    if "shared_name" in node_def.attr:
      del node_def.attr["shared_name"]

  fdef.signature.name = _clean_function_name(fdef.signature.name)
  return fdef


def _list_function_deps(fdef):
  # TODO(andresp): Recurse into list attributes and into NameAttrList attrs both
  # when listing deps and when fixing them. `function_def_to_graph` also
  # requires fixes.
  deps = set()
  for node_def in fdef.node_def:
    for _, attr_value in node_def.attr.items():
      if attr_value.WhichOneof("value") == "func":
        deps.add(attr_value.func.name)
  return deps


def _clean_function_name(name):
  """Vanity function to keep the function names comprehensible."""
  # Note: each time a function is wrapped into `function_lib.ConcreteFunction`
  # its name becomes "__inference_<orig>_xyz".
  match = re.search(r"^__inference_(.*)_\d+$", name)
  if match:
    return match.group(1)
  else:
    return name

tensorflow.core.framework.function_pb2.FunctionDef tensorflow.python.eager.function.is_same_structure tensorflow.python.platform.tf_logging.warning tensorflow.python.framework.ops.init_scope tensorflow.python.saved_model.nested_structure_coder.StructureCoder tensorflow.python.util.compat.as_bytes tensorflow.python.framework.function_def_to_graph.function_def_to_graph tensorflow.python.eager.function.FunctionSpec tensorflow.python.eager.function.ConcreteFunction tensorflow.python.util.nest.flatten
[(72, 'tensorflow.python.eager.function.FunctionSpec', 'function_lib.FunctionSpec', (['fullargspec', 'is_method', 'args_to_prepend', 'kwargs_to_include', 'input_signature'], {}), True, 'from tensorflow.python.eager import function as function_lib\n'), (133, 'tensorflow.python.saved_model.nested_structure_coder.StructureCoder', 'nested_structure_coder.StructureCoder', ([], {}), False, 'from tensorflow.python.saved_model import nested_structure_coder\n'), (219, 'collections.defaultdict', 'collections.defaultdict', (['list'], {}), False, 'import collections\n'), (220, 'collections.defaultdict', 'collections.defaultdict', (['(lambda : 0)'], {}), False, 'import collections\n'), (263, 'tensorflow.core.framework.function_pb2.FunctionDef', 'function_pb2.FunctionDef', ([], {}), False, 'from tensorflow.core.framework import function_pb2\n'), (308, 're.search', 're.search', (['"""^__inference_(.*)_\\\\d+$"""', 'name'], {}), False, 'import re\n'), (50, 'tensorflow.python.util.nest.flatten', 'nest.flatten', (['arg'], {}), False, 'from tensorflow.python.util import nest\n'), (51, 'tensorflow.python.util.nest.flatten', 'nest.flatten', (['stored_input'], {}), False, 'from tensorflow.python.util import nest\n'), (202, 'tensorflow.python.framework.function_def_to_graph.function_def_to_graph', 'function_def_lib.function_def_to_graph', (['copy'], {}), True, 'from tensorflow.python.framework import function_def_to_graph as function_def_lib\n'), (205, 'tensorflow.python.eager.function.ConcreteFunction', 'function_lib.ConcreteFunction', (['func_graph'], {}), True, 'from tensorflow.python.eager import function as function_lib\n'), (47, 'tensorflow.python.eager.function.is_same_structure', 'function_lib.is_same_structure', (['arg', 'stored_input'], {}), True, 'from tensorflow.python.eager import function as function_lib\n'), (211, 'tensorflow.python.framework.ops.init_scope', 'ops.init_scope', ([], {}), False, 'from tensorflow.python.framework import ops\n'), (158, 'tensorflow.python.util.nest.flatten', 'nest.flatten', (['canonicalized_inputs'], {}), False, 'from tensorflow.python.util import nest\n'), (272, 'tensorflow.python.util.compat.as_bytes', 'compat.as_bytes', (['functions[fname]._gradient_name'], {}), False, 'from tensorflow.python.util import compat\n'), (275, 'tensorflow.python.platform.tf_logging.warning', 'logging.warning', (['"""Importing a function (%s) with ops with custom gradients. Will likely fail if a gradient is requested."""', 'fdef.signature.name'], {}), True, 'from tensorflow.python.platform import tf_logging as logging\n')]