# Copyright 2015 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.
# =============================================================================
"""Tests for functions."""

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

import re
import sys
import time

import numpy as np

from tensorflow.core.framework import function_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import function
from tensorflow.python.framework import graph_to_function_def
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.framework.errors import InvalidArgumentError
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import gen_logging_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import linalg_ops
from tensorflow.python.ops import logging_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging


def _OptimizerOptions():
  for cse in [False, True]:
    for inline in [False, True]:
      for cfold in [False, True]:
        cfg = config_pb2.ConfigProto(
            graph_options=config_pb2.GraphOptions(
                optimizer_options=config_pb2.OptimizerOptions(
                    opt_level=config_pb2.OptimizerOptions.L0,
                    do_common_subexpression_elimination=cse,
                    do_function_inlining=inline,
                    do_constant_folding=cfold)))
        if cse:
          cfg.graph_options.rewrite_options.arithmetic_optimization = (
              rewriter_config_pb2.RewriterConfig.ON)
        else:
          cfg.graph_options.rewrite_options.arithmetic_optimization = (
              rewriter_config_pb2.RewriterConfig.OFF)
        if inline:
          cfg.graph_options.rewrite_options.function_optimization = (
              rewriter_config_pb2.RewriterConfig.ON)
        else:
          cfg.graph_options.rewrite_options.function_optimization = (
              rewriter_config_pb2.RewriterConfig.OFF)
        if cfold:
          cfg.graph_options.rewrite_options.constant_folding = (
              rewriter_config_pb2.RewriterConfig.ON)
        else:
          cfg.graph_options.rewrite_options.constant_folding = (
              rewriter_config_pb2.RewriterConfig.OFF)
        yield cfg


class FunctionTest(test.TestCase):
  """Test methods for verifying Function support.

  These test methods are used as mix-ins in two test cases: with
  and without C API support.
  """

  def testIdentity(self):

    @function.Defun(dtypes.float32, func_name="MyIdentity")
    def MyIdentityFunc(a):
      return a

    with ops.Graph().as_default():
      call = MyIdentityFunc([18.0])
      self.assertEqual("MyIdentity", call.op.name)
      with session.Session() as sess:
        self.assertAllEqual([18.0], self.evaluate(call))

  @test_util.run_deprecated_v1
  def testIdentityImplicitDeref(self):

    @function.Defun(dtypes.float32, func_name="MyIdentity")
    def MyIdentityFunc(a):
      return a

    with ops.Graph().as_default():
      var = variables.VariableV1([18.0])
      call = MyIdentityFunc(var._ref())  # pylint: disable=protected-access
      self.assertEqual("MyIdentity", call.op.name)
      for cfg in _OptimizerOptions():
        with session.Session(config=cfg) as sess:
          self.evaluate(var.initializer)
          self.assertAllEqual([18.0], self.evaluate(call))

  def testIdentityOutputName(self):

    @function.Defun(
        dtypes.float32, func_name="MyIdentity", out_names=["my_result_name"])
    def MyIdentityFunc(a):
      return a

    with ops.Graph().as_default():
      call = MyIdentityFunc([18.0])
      self.assertEqual("MyIdentity", call.op.name)
      with session.Session() as sess:
        self.assertAllEqual([18.0], self.evaluate(call))

  def testTooManyOutputNames(self):

    @function.Defun(
        dtypes.float32,
        func_name="MyIdentity",
        out_names=["my_result1", "my_result2"])
    def MyIdentityFunc(a):
      return a

    with ops.Graph().as_default():
      with self.assertRaisesRegexp(
          errors_impl.InvalidArgumentError,
          (r"output names must be either empty or equal in size to outputs. "
           "output names size = 2 outputs size = 1")):
        MyIdentityFunc([18.0])

  def testDefineFunction2Args(self):

    @function.Defun(dtypes.float32, dtypes.float32, func_name="APlus2B")
    def APlus2B(a, b):
      return a + b * 2

    with ops.Graph().as_default():
      call = APlus2B([1.0], [2.0])
      self.assertEqual("APlus2B", call.op.name)
      with session.Session() as sess:
        self.assertAllEqual([5.0], self.evaluate(call))

  def testFunctionWithNoOutput(self):

    @function.Defun(dtypes.float32, dtypes.float32)
    def APlus2B(a, b):
      c = a + b * 2  # Create some ops to have nodes in the body
      print(c)  # Using 'print' to make lint happy

    with ops.Graph().as_default():
      # Call function. There should be no exceptions.
      APlus2B([1.0], [2.0])

  def testDefineFunction2ArgsOutputName(self):

    @function.Defun(
        dtypes.float32,
        dtypes.float32,
        func_name="APlus2B",
        out_names=["my_result_name"])
    def APlus2B(a, b):
      return a + b * 2

    # APlus2B is stateless.
    self.assertEqual([], APlus2B.stateful_ops)
    with ops.Graph().as_default():
      call = APlus2B([1.0], [2.0])
      self.assertEqual("APlus2B", call.op.name)
      with session.Session() as sess:
        self.assertAllEqual([5.0], self.evaluate(call))

  def testDefineFunctionDuplicateOutputs(self):

    @function.Defun(dtypes.float32, func_name="Duplicate")
    def Duplicate(a):
      b = a + 1.0
      return b, b

    g = ops.Graph()
    with g.as_default():
      Duplicate([3.0])
      func_sig = g.as_graph_def().library.function[0].signature
      # The names given to both outputs should be different
      # even though the same tensor is emitted to both.
      out_names = [a.name for a in func_sig.output_arg]
      self.assertEqual(2, len(out_names))
      self.assertNotEqual(out_names[0], out_names[1])

  def testGradientFunc(self):

    @function.Defun(dtypes.float32, func_name="XSquarePlusOneFn")
    def XSquarePlusOne(x):
      return x * x + 1.0

    @function.Defun(dtypes.float32, dtypes.float32)
    def XSquarePlusOneGrad(x, dy):
      dx = functional_ops.symbolic_gradient(
          input=[x, dy], Tout=[dtypes.float32], f="XSquarePlusOneFn", name="dx")
      return dx

    g = ops.Graph()
    with g.as_default():
      call_f = XSquarePlusOne([2.0])
      call_g = XSquarePlusOneGrad([2.0], [0.1])

      with session.Session() as sess:
        self.assertAllClose([5.0], self.evaluate(call_f))
        self.assertAllClose([0.4], self.evaluate(call_g))