# 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.
# ==============================================================================

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

import os


from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.eager import backprop
from tensorflow.python.eager import def_function
from tensorflow.python.eager import wrap_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import saver as saver_lib


class WrapFunctionTest(test.TestCase):

  def testDocString(self):

    def f(x, do_add):
      v = variables.Variable(5.0)
      if do_add:
        op = v.assign_add(x)
      else:
        op = v.assign_sub(x)
      with ops.control_dependencies([op]):
        return v.read_value()

    f_add = wrap_function.wrap_function(
        f, [tensor_spec.TensorSpec((), dtypes.float32), True])

    self.assertAllEqual(f_add(1.0), 6.0)
    self.assertAllEqual(f_add(1.0), 7.0)

    # Can call tf.compat.v1.wrap_function again to get a new trace, a new set
    # of variables, and possibly different non-template arguments.
    f_sub = wrap_function.wrap_function(
        f, [tensor_spec.TensorSpec((), dtypes.float32), False])

    self.assertAllEqual(f_sub(1.0), 4.0)
    self.assertAllEqual(f_sub(1.0), 3.0)

  def testPrune(self):

    x_in = []
    x_out = []

    def f(x, y):
      x_in.append(x)
      xx = x * x
      x_out.append(xx)
      return xx, 2 * y*y

    f_wrapped = wrap_function.wrap_function(
        f, [tensor_spec.TensorSpec((), dtypes.float32)] * 2)

    f_pruned = f_wrapped.prune(x_in[0], [x_out[0]])
    self.assertAllEqual(f_pruned(ops.convert_to_tensor(2.0)), [4.0])

  def _assert_single_captured_variable_argument(self, graph_def):
    # The single FunctionDef should have one argument, a captured variable
    function_def, = graph_def.library.function
    self.assertLen(function_def.signature.input_arg, 1)
    function_arg, = function_def.signature.input_arg
    self.assertEqual(dtypes.resource, dtypes.as_dtype(function_arg.type))

  def testVariableLifting(self):
    save_prefix = os.path.join(self.get_temp_dir(), 'meta_graph_test')

    export_graph = ops.Graph()
    with export_graph.as_default():
      v = variables.Variable(1.)
      array_ops.identity(v + 1., name='output')
      saver = saver_lib.Saver([v])
      with self.test_session() as session:
        session.run(v.initializer)
        saver.save(session, save_prefix)

    def importer():
      saver_lib.import_meta_graph(save_prefix + '.meta')
      return ops.get_default_graph().as_graph_element('output:0')

    wrapped = wrap_function.wrap_function(importer, [])
    lifted_variables = list(wrapped.graph.variables)
    self.assertLen(lifted_variables, 1)
    initializer = wrapped.prune(
        [], wrapped.graph.as_gr
aph_element(v.initializer.name))
    self.assertEqual(lifted_variables, list(initializer.graph.variables))
    self.assertEqual(initializer.graph.external_captures,
                     wrapped.graph.external_captures)

    @def_function.function
    def wraps_initializer():
      initializer()

    wraps_initializer()
    s
elf.assertEqual(1., lifted_variables[0].numpy())
    wrapped_initializer_graphdef = (
        wraps_initializer.get_concrete_function().graph.as_graph_def())
    self._assert_single_captured_variable_argument(wrapped_initializer_graphdef)

    @def_function.function
    def wraps_wrapped():
      return wrapped()

    # Verify that the original graph also has the correct signature.
    wrapped_wrapped_graphdef = (
        wraps_wrapped.get_concrete_function().graph.as_graph_def())
    self._assert_single_captured_variable_argument(wrapped_wrapped_graphdef)
    # Now check that the graph runs wrapped, from eager, and when pruned.
    self.assertAllEqual(wraps_wrapped().numpy(),
                        lifted_variables[0].numpy() + 1.)
    self.assertAllEqual(wrapped().numpy(), lifted_variables[0].numpy() + 1.)
    pruned = wrapped.prune([], wrapped.graph.as_graph_element('output:0'))
    self.assertAllEqual(wrapped().numpy(), pruned().numpy())

  def testNoArguments(self):

    def f():
      return constant_op.constant(1.)

    f_wrapped = wrap_function.wrap_function(f, [])
    self.assertAllEqual(1.0, f_wrapped())

  def testPruneCaptures(self):

    v1 = variables.Variable(2.)

    def f():
      v2 = variables.Variable(3.)
      return array_ops.identity(v1 * v2 * constant_op.constant(1.), 'fetch')

    f_wrapped = wrap_function.wrap_function(f, [])
    self.assertAllEqual(6.0, f_wrapped())

    # Test pruning directly on the inputs
    pruned = f_wrapped.prune(
        feeds=f_wrapped.inputs,
        fetches=f_wrapped.graph.get_tensor_by_name('fetch:0'))
    self.assertAllEqual(6.0, pruned())

    # Test pruning with no inputs
    pruned = f_wrapped.prune(
        feeds=(),
        fetches=f_wrapped.graph.get_tensor_by_name('fetch:0'))
    self.assertAllEqual(6.0, pruned())

  def testCollectionsIsolation(self):

    v1 = variables.Variable(2.)
    v2_holder = []
    def f():
      v2 = variables.Variable(3.)
      v2_holder.append(v2)
      ops.add_to_collection(ops.GraphKeys.LOSSES, v2 * constant_op.constant(3.))
      return array_ops.identity(v1 * v2 * constant_op.constant(1.), 'fetch')

    f_wrapped = wrap_function.wrap_function(f, [])
    self.assertAllEqual(6.0, f_wrapped())
    self.assertEqual(
        len(f_wrapped.graph.get_collection(ops.GraphKeys.LOSSES)), 1)
    f_var_collection = f_wrapped.graph.get_collection(
        ops.GraphKeys.TRAINABLE_VARIABLES)
    self.assertEqual(len(f_var_collection), 1)
    self.assertIs(f_var_collection[0], v2_holder[0])

    v3_holder = []
    def g():
      v3 = variables.Variable(4.)
      v3_holder.append(v3)
      ops.add_to_collection(ops.GraphKeys.LOSSES, v3 * constant_op.constant(3.))
      return array_ops.identity(v1 * v3 * constant_op.constant(1.), 'fetch')

    g_wrapped = wrap_function.wrap_function(g, [])
    self.assertAllEqual(8.0, g_wrapped())
    self.assertEqual(
        len(g_wrapped.graph.get_collection(ops.GraphKeys.LOSSES)), 1)
    g_var_collection = g_wrapped.graph.get_collection(
        ops.GraphKeys.TRAINABLE_VARIABLES)
    self.assertEqual(len(g_var_collection), 1)
    self.assertIs(g_var_collection[0], v3_holder[0])

    # Both have only one value, and their values aren't equal. So no sharing.
    self.assertNotEqual(g_wrapped.graph.get_collection(ops.GraphKeys.LOSSES),
                        f_wrapped.graph.get_collection(ops.GraphKeys.LOSSES))

  def testGradientsOfPrune(self):

    v1 = variables.Variable(2.)
    v2_holder = []

    def f(z):
      v2 = variables.Variable(3.)
      v2_holder.append(v2)
      return array_ops.identity(v1 * v2 * z, 'fetch')

    f_wrapped = wrap_function.wrap_function(
        f, [tensor_spec.TensorSpec((), dtype=dtypes.float32)])

    x = constant_op.constant(1.)
    with backprop.GradientTape() as tape:
      tape.watch(x)
      out = f_wrapped(x)
    grads = tape.gradient(out, [x, v1, v2_holder[0]])

    self.assertAllEqual(6.0, out)
    self.assertAllEqual([6.0, 3.0, 2.0], grads)

    pruned = f_wrapp