# Copyright 2015 Google Inc. 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 time

import numpy as np
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

from tensorflow.python.framework import function
from tensorflow.python.ops import functional_ops


def _OptimizerOptions():
  for cse in [False, True]:
    for inline in [False, True]:
      for cfold in [False, True]:
        yield tf.ConfigProto(
            graph_options=tf.GraphOptions(optimizer_options=tf.OptimizerOptions(
                opt_level=tf.OptimizerOptions.L0,
                do_common_subexpression_elimination=cse,
                do_function_inlining=inline,
                do_constant_folding=cfold)))


class FunctionTest(tf.test.TestCase):

  def _mat(self, x):
    return np.array([x]).astype("float32").reshape([1, 1])

  def testBasic(self):
    g = tf.Graph()

    # Define a function
    #   foo(a:float, b:float, c:float)->u:float,v:float,w:float
    #     u = matmul(a, b) + c
    #     v = u^2
    #     w = u + v
    # TODO(zhifengc): replaces w/ a nicer @decorator sugar.
    foo = tf.Graph()
    with foo.as_default():
      a = tf.placeholder(tf.float32, name="a")
      b = tf.placeholder(tf.float32, name="b")
      c = tf.placeholder(tf.float32, name="c")
      u = tf.add(tf.matmul(a, b), c, name="u")
      v = tf.square(u, name="v")
      w = tf.add_n([u, v], name="w")
    fdef = function.graph_to_function_def(foo, "foo", [a, b, c], [u, v, w])
    g._add_function(fdef)

    # Compute 2 * 3 + 4 and its square.
    with g.as_default(), tf.Session() as sess:
      two = tf.constant(self._mat(2.0), name="two")
      three = tf.constant(self._mat(3.0), name="three")
      four = tf.constant(self._mat(4.0), name="four")
      # TODO(zhifengc): w/ @decorator sugar, we will just do:
      #   y, s, t = foo_func(two, three, four)

      # The graph contains two ops each of which calls foo.
      u0, v0, w0 = g.create_op("foo",
                               [two, three, four],
                               [tf.float32, tf.float32, tf.float32],
                               compute_shapes=False).outputs
      u1, v1, w1 = g.create_op("foo",
                               [four, two, three],
                               [tf.float32, tf.float32, tf.float32],
                               compute_shapes=False).outputs

      # Checks some property of the graph def.
      gdef = g.as_graph_def()
      self.assertEqual(len(gdef.node), 5)  # 5 nodes added.
      self.assertEqual(len(gdef.library.function), 1)  # 1 function is defined.

      for _ in xrange(10):
        # Run the graph, which is basicly two function calls.
        ans_u0, ans_v0, ans_w0, ans_u1, ans_v1, ans_w1 = sess.run([u0, v0, w0,
                                                                   u1, v1, w1])
        self.assertAllEqual(ans_u0, self._mat(10.0))  # 2 * 3 + 4 = 10
        self.assertAllEqual(ans_v0, self._mat(100.0))  # 10^2 = 100
        self.assertAllEqual(ans_w0, self._mat(110.0))  # 100 + 10 = 110
        self.assertAllEqual(ans_u1, self._mat(11.0))  # 4 * 2 + 3 = 11
        self.assertAllEqual(ans_v1, self._mat(121.0))  # 11^2 = 121
        self.assertAllEqual(ans_w1, self._mat(132.0))  # 11 + 121 = 132

  def testDefineFunction2Args(self):

    def APlus2B(a, b):
      return a + b * 2

    with tf.Graph().as_default():
      f_def = function.define_function(APlus2B, {"a": tf.float32,
                                                 "b": tf.float32})
      one = tf.constant([1.0])
      two = tf.constant([2.0])
      call = function.call_function(f_def, one, two)
      self.assertEquals("APlus2B", call.op.name)
      with tf.Session() as sess:
        self.assertAllEqual([5.0], sess.run(call))

  def testGradientFunc(self):

    def XSquarePlusOne(x):
      return x * x + 1.0

    def XSquarePlusOneGrad(x, dy):
      dx = functional_ops._symbolic_gradient(input=[x, dy],
                                             Tout=[tf.float32],
                                             f="XSquarePlusOne",
                                             name="dx")
      return dx

    g = tf.Graph()
    with g.as_default():
      f = function.define_function(XSquarePlusOne, {"x": tf.float32})
      g = function.define_function(XSquarePlusOneGrad, {"x": tf.float32,
                                                        "dy": tf.float32})
      epsilon = tf.constant([0.1])
      two = tf.constant([2.0])
      call_f = function.call_function(f, two)
      call_g = function.call_function(g, two, epsilon)

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

  def testTanhSymGrad(self):
    g = tf.Graph()
    with g.as_default():
      @function.Defun(tf.float32)
      def Forward(x):
        return tf.reduce_sum(tf.tanh(x))
      x = tf.placeholder(tf.float32)
      y = Forward(x)
      dx = tf.gradients([y], [x])

    inp = np.array([-1, 1, 2, -2], dtype=np.float32)
    feed = {x: inp}
    cfg = tf.ConfigProto(
        graph_options=tf.GraphOptions(
            optimizer_options=tf.OptimizerOptions(
                opt_level=tf.OptimizerOptions.L1,
                do_function_inlining=True)))
    with tf.Session(graph=g, config=cfg) as sess:
      out, = sess.run(dx, feed)
    self.assertAllClose(1 - np.square(np.tanh(inp)), out)

  def testCustomGradient(self):
    g = tf.Graph()
    dtype = tf.float32
    with g.as_default():

      @function.Defun(dtype, dtype, dtype)
      def XentLossGrad(logits, labels, dloss):
        dlogits = tf.reshape(dloss, [-1, 1]) * (tf.nn.softmax(logits) - labels)
        dlabels = tf.zeros_like(labels)
        # Takes exp(dlogits) to differentiate it from the "correct" gradient.
        return tf.exp(dlogits), dlabels

      @function.Defun(dtype, dtype, grad_func="XentLossGrad")
      def XentLoss(logits, labels):
        return tf.reduce_sum(labels * tf.log(tf.nn.softmax(logits)), 1)

      logits = tf.placeholder(dtype)
      labels = tf.placeholder(dtype)
      loss = XentLoss(logits, labels)
      dlogits = tf.gradients([loss], [logits])

    x = np.random.uniform(-10., 10., size=(4, 9)).astype(np.float32)
    prob = np.exp(x) / np.sum(np.exp(x), 1, keepdims=1)
    y = np.random.uniform(-10., 10., size=(4, 9)).astype(np.float32)
    for cfg in _OptimizerOptions():
      print("cfg = ", cfg)
      with tf.Session(graph=g, config=cfg) as sess:
        out, = sess.run(dlogits, {logits: x, labels: y})
      self.assertAllClose(out, np.exp(prob - y))

  def testCustomGradientError(self):
    g = tf.Graph()
    dtype = tf.float32
    with g.as_default():

      @function.Defun(dtype, dtype, dtype)
      def Grad(x, dy, dz):
        # Should have returned 1 result.
        return x, dy + dz

      @function.Defun(dtype, grad_func="Grad")
      def Forward(x):
        return x, x

      inp = tf.placeholder(dtype)
      out = tf.add_n(Forward(inp))
      dinp = tf.gradients(out, [inp])

    x = np.random.uniform(-10., 10., size=(4, 9)).astype(np.float32)
    with tf.Session(graph=g) as sess:
      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          "SymGrad expects to return 1.*but get 2.*instead"):
        _ = sess.run(dinp, {inp: x})

  def testSymGradShape(self):
    g = tf.Graph()
    with g.as_default():
      x = tf.placeholder(tf.float32, [25, 4])
      y = tf.placeholder(tf.float32, [200, 100])
      dz = tf.placeholder(tf.float32, [1])
      # We assume Foo is a function of (x, y) -> (z) Then, Foo's
      # gradient function is (x, y, dz) -> (dx, dy).  dx's shape
      # should be the same as x's; and dy's shape should be the same
      # as y's.
      dx, dy = functional_ops._symbolic_gradient(input=[x, y, dz],
                                                 Tout=[tf.float32] * 2,
                                                 f="Foo")
      self.assertEquals(x.get_shape(), dx.get_shape())
      self.assertEquals(y.get_shape(), dy.get_shape())

  def testDefineFunctionNoArgs(self):

    def AConstant():
      return tf.constant([42])

    with tf.Graph().as_default():
      f_def = function.define_function(AConstant, {})
      call = function.call_function(f_def)
      self.assertEquals("AConstant", call.op.name)
      with tf.Session() as sess:
        self.assertAllEqual([42], sess.run(call))

  def testDefineFunctionNames(self):

    def Foo(a):
      return a + 1

    with tf.Graph().as_default():
      f_def = function.define_function(Foo, {"a": tf.float32})
      one = tf.constant([1.0])
      call1 = function.call_function(f_def, one)
      self.assertEquals("Foo", call1.op.name)
      call2 = function.call_function(f_def, one)
      self.assertEquals("Foo_1", call2.op.name)
      call3 = function.call_function(f_def, one, name="mine")
      self.assertEquals("mine", call3.op.name)
      with tf.name_scope("my"):
        call4 = function.call_function(f_def, one, name="precious")
        self.assertEquals("my/precious", call4.op.name)

  def testDefineErrors(self):

    def NoResult():
      pass

    def DefaultArg(unused_a=12):
      return tf.constant([1])

    def KwArgs(**unused_kwargs):
      return tf.constant([1])

    def PlusMinus(a, b):
      return a + b, b - a

    with tf.Graph().as_default():
      with self.assertRaisesRegexp(ValueError, "return at least one tensor"):
        function.define_function(NoResult, {})
      with self.assertRaisesRegexp(ValueError, "are not supported"):
        function.define_function(DefaultArg, {})
      with self.assertRaisesRegexp(ValueError, "are not supported"):
        function.define_function(KwArgs, {})
      with self.assertRaisesRegexp(ValueError, "specified input types"):
        function.define_function(PlusMinus, {})
      with self.assertRaisesRegexp(ValueError, "specified input types"):
        function.define_function(PlusMinus, {"c": tf.float32})
      with self.assertRaisesRegexp(ValueError, "type for argument: b"):
        function.define_function(PlusMinus, {"a": tf.float32,
                                             "c": tf.float32})
      with self.assertRaisesRegexp(ValueError, "specified input types"):
        function.define_function(PlusMinus, {"a": tf.float32,
                                             "b": tf.float32,
                                             "c": tf.float32})

  def testCallErrors(self):

    def Const():
      return tf.constant(1)

    def PlusOne(a):
      return a + 1

    def PlusMinus(a, b):
      return a + b, b - a

    with tf.Graph().as_default():
      one = tf.constant([1])
      two = tf.constant([2])
      const = function.define_function(Const, {})
      plus_one = function.define_function(PlusOne, {"a": tf.int32})
      plus_minus = function.define_function(PlusMinus, {"a": tf.int32,
                                                        "b": tf.int32})

      function.call_function(const)
      with self.assertRaisesRegexp(ValueError, "arguments: 0"):
        function.call_function(const, one)
      with self.assertRaisesRegexp(ValueError, "arguments: 0"):
        function.call_function(const, one, two)

      with self.assertRaisesRegexp(ValueError, "arguments: 1"):
        function.call_function(plus_one)
      function.call_function(plus_one, one)
      with self.assertRaisesRegexp(ValueError, "arguments: 1"):
        function.call_function(plus_one, one, two)

      with self.assertRaisesRegexp(ValueError, "arguments: 2"):
        function.call_function(plus_minus)
      with self.assertRaisesRegexp(ValueError, "arguments: 2"):
        function.call_function(plus_minus, one)
      function.call_function(plus_minus, one, two)

      function.call_function(plus_one, one, name="p1")
      with self.assertRaisesRegexp(ValueError, "Unknown keyword arguments"):
        function.call_function(plus_one, one, device="/gpu:0")

  def testFunctionDecorator(self):

    with tf.Graph().as_default():

      @function.Defun(tf.float32)
      def Minus1(b):
        return b - 1.0

      two = tf.constant([2.])
      call1 = Minus1(two)
      self.assertEquals("Minus1", call1.op.name)
      # pylint: disable=unexpected-keyword-arg
      call2 = Minus1(call1, name="next")
      # pylint:enable=unexpected-keyword-arg
      self.assertEquals("next", call2.op.name)
      with tf.Session() as sess:
        self.assertAllEqual([1], sess.run(call1))
        self.assertAllEqual([0], sess.run(call2))

  def testNestedFunction(self):
    with tf.Graph().as_default():

      @function.Defun(tf.float32)
      def Cube(x):
        return x * x * x

      @function.Defun(tf.float32, tf.float32)
      def CubeXPlusY(x, y):
        return Cube(x) + y

      z = CubeXPlusY(tf.constant(3.0), tf.constant(-2.0))
      with self.test_session():
        self.assertAllEqual(z.eval(), 25.0)


class UnrollLSTMTest(tf.test.TestCase):
  BATCH_SIZE = 16
  LSTM_DIMS = 32
  NUM_UNROLL = 20

  def _Weights(self):
    dims = self.LSTM_DIMS
    return tf.random_uniform([2 * dims, 4 * dims], -1, 1, seed=123456)

  def _Input(self):
    return tf.random_uniform(
        [self.NUM_UNROLL, self.BATCH_SIZE, self.LSTM_DIMS],
        seed=654321)

  # Helper to construct a LSTM cell graph.
  @classmethod
  def LSTMCell(cls, x, mprev, cprev, weights):
    xm = tf.concat(1, [x, mprev])
    i_i, i_g, f_g, o_g = tf.split(1, 4, tf.matmul(xm, weights))
    new_c = tf.sigmoid(f_g) * cprev + tf.sigmoid(i_g) * tf.tanh(i_i)
    new_c = tf.clip_by_value(new_c, -50.0, 50.0)
    new_m = tf.sigmoid(o_g) * tf.tanh(new_c)
    return new_m, new_c

  def _BuildForward(self, weights, inp, mode="cell"):

    def Loop(cell, w, i):
      x = tf.unpack(i, self.NUM_UNROLL)
      m = tf.zeros_like(x[0])
      c = tf.zeros_like(x[0])
      for i in range(self.NUM_UNROLL):
        m, c = cell(x[i], m, c, w)
      return m

    cell = UnrollLSTMTest.LSTMCell
    if mode == "complete":
      # Constructs the complete graph in python.
      return Loop(cell, weights, inp)

    cell = function.Defun(x=tf.float32,
                          mprev=tf.float32,
                          cprev=tf.float32,
                          weights=tf.float32)(cell)
    if mode == "cell":
      # Just represent the LSTM as a function.
      return Loop(cell, weights, inp)

    if mode == "loop":
      # Wraps the whole loop as a function.
      @function.Defun(tf.float32, tf.float32)
      def LSTMLoop(w, i):
        return Loop(cell, w, i)

      return LSTMLoop(weights, inp)

    if mode == "loop10":
      # Wraps 10 lstm steps into one function, and the whole loop
      # into another calling the formers.

      # Groups 10 steps at a time.
      @function.Defun(tf.float32, tf.float32, tf.float32,
                      *([tf.float32] * 10))
      def Loop10(w, m, c, *args):
        for x in args:
          m, c = cell(x, m, c, w)
        return m, c

      @function.Defun(tf.float32, tf.float32)
      def LSTMLoop10(weights, inp):
        x = tf.unpack(inp, self.NUM_UNROLL)
        m = tf.zeros_like(x[0])
        c = tf.zeros_like(x[0])
        assert self.NUM_UNROLL % 10 == 0
        for i in range(0, self.NUM_UNROLL, 10):
          m, c = Loop10(weights, m, c, *x[i:i + 10])
        return m

      return LSTMLoop10(weights, inp)

  def testUnrollLSTM(self):
    # Run one step of the unrolled lstm graph.
    def RunForward(mode, cfg=None):
      print("mode = ", mode)
      g = tf.Graph()
      start = time.time()
      with g.as_default():
        weights = self._Weights()
        inp = self._Input()
        m = self._BuildForward(weights, inp, mode)
      gdef = g.as_graph_def()
      finish = time.time()
      print("time: ", finish - start, " txt size: ", len(str(gdef)),
            "gdef bin size: ", len(gdef.SerializeToString()))
      with g.as_default(), tf.Session(config=cfg) as sess:
        return sess.run(m)

    mv0 = RunForward("complete")
    for cfg in _OptimizerOptions():
      print("cfg = ", cfg)
      mv1 = RunForward("cell", cfg)
      mv2 = RunForward("loop", cfg)
      mv3 = RunForward("loop10", cfg)
      self.assertAllClose(mv0, mv1, rtol=1e-4)
      self.assertAllClose(mv0, mv2, rtol=1e-4)
      self.assertAllClose(mv0, mv3, rtol=1e-4)

  def testUnrollLSTMGrad(self):
    # Run one step of the unrolled lstm graph.
    def RunForwardBackward(mode, cfg=None):
      print("mode = ", mode)
      g = tf.Graph()
      start = time.time()
      with g.as_default():
        weights = self._Weights()
        inp = self._Input()
        m = self._BuildForward(weights, inp, mode)
        loss = tf.reduce_sum(tf.square(m))
        dw = tf.gradients([loss], [weights])
      gdef = g.as_graph_def()
      finish = time.time()
      print("time: ", finish - start, " txt size: ", len(str(gdef)),
            "gdef bin size: ", len(gdef.SerializeToString()))
      with g.as_default(), tf.Session(config=cfg) as sess:
        return sess.run(dw)

    d0 = RunForwardBackward("complete")
    for cfg in _OptimizerOptions():
      print("cfg = ", cfg)
      d1 = RunForwardBackward("cell", cfg)
      d2 = RunForwardBackward("loop", cfg)
      d3 = RunForwardBackward("loop10", cfg)
      self.assertAllClose(d0, d1, rtol=1e-4)
      self.assertAllClose(d0, d2, rtol=1e-4)
      self.assertAllClose(d0, d3, rtol=1e-4)


class FunctionInlineControlTest(tf.test.TestCase):

  def testFoo(self):
    dtype = tf.float32
    cfg = tf.ConfigProto(
        graph_options=tf.GraphOptions(optimizer_options=tf.OptimizerOptions(
            opt_level=tf.OptimizerOptions.L0,
            do_common_subexpression_elimination=True,
            do_function_inlining=True,
            do_constant_folding=True)))
    for noinline in [False, True]:
      g = tf.Graph()
      with g.as_default():

        @function.Defun(dtype)
        def Cell(v):
          # If v is a vector [n, 1], x is a big square matrix.
          x = tf.tanh(v + tf.transpose(v, [1, 0]))
          return tf.reduce_sum(x, 1, keep_dims=True)

        @function.Defun(dtype)
        def Forward(x):
          for _ in range(10):
            x = Cell(x, noinline=noinline)
          return tf.reduce_sum(x, [0, 1])

        x = tf.placeholder(dtype)
        y = Forward(x)
        dx, = tf.gradients([y], [x])

      np.random.seed(321)
      inp = np.random.uniform(-1, 1, [16, 1]).astype(np.float32)
      with tf.Session(graph=g, config=cfg) as sess:
        ans = sess.run([y, dx], {x: inp})
        print(ans[0], np.sum(ans[1]))
        self.assertAllClose(ans[0], 255.971, rtol=1e-3)
        self.assertAllClose(np.sum(ans[1]), 13.0408, rtol=1e-3)


if __name__ == "__main__":
  tf.test.main()
