# Copyright 2017 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.
# ==============================================================================
"""Test cases for operators with > 3 or arbitrary numbers of arguments."""

import unittest

import numpy as np

from tensorflow.compiler.tests import xla_test
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import googletest


class NAryOpsTest(xla_test.XLATestCase):

  def _testNAry(self, op, args, expected, equality_fn=None):
    with self.session() as session:
      with self.test_scope():
        placeholders = [
            array_ops.placeholder(dtypes.as_dtype(arg.dtype), arg.shape)
            for arg in args
        ]
        feeds = {placeholders[i]: args[i] for i in range(0, len(args))}
        output = op(placeholders)
      result = session.run(output, feeds)
      if not equality_fn:
        equality_fn = self.assertAllClose
      equality_fn(result, expected, rtol=1e-3)

  def _nAryListCheck(self, results, expected, **kwargs):
    self.assertEqual(len(results), len(expected))
    for (r, e) in zip(results, expected):
      self.assertAllClose(r, e, **kwargs)

  def _testNAryLists(self, op, args, expected):
    self._testNAry(op, args, expected, equality_fn=self._nAryListCheck)

  def testFloat(self):
    self._testNAry(math_ops.add_n,
                   [np.array([[1, 2, 3]], dtype=np.float32)],
                   expected=np.array([[1, 2, 3]], dtype=np.float32))

    self._testNAry(math_ops.add_n,
                   [np.array([1, 2], dtype=np.float32),
                    np.array([10, 20], dtype=np.float32)],
                   expected=np.array([11, 22], dtype=np.float32))
    self._testNAry(math_ops.add_n,
                   [np.array([-4], dtype=np.float32),
                    np.array([10], dtype=np.float32),
                    np.array([42], dtype=np.float32)],
                   expected=np.array([48], dtype=np.float32))

  def testComplex(self):
    for dtype in self.complex_types:
      self._testNAry(
          math_ops.add_n, [np.array([[1 + 2j, 2 - 3j, 3 + 4j]], dtype=dtype)],
          expected=np.array([[1 + 2j, 2 - 3j, 3 + 4j]], dtype=dtype))

      self._testNAry(
          math_ops.add_n, [
              np.array([1 + 2j, 2 - 3j], dtype=dtype),
              np.array([10j, 20], dtype=dtype)
          ],
          expected=np.array([1 + 12j, 22 - 3j], dtype=dtype))
      self._testNAry(
          math_ops.add_n, [
              np.array([-4, 5j], dtype=dtype),
              np.array([2 + 10j, -2], dtype=dtype),
              np.array([42j, 3 + 3j], dtype=dtype)
          ],
          expected=np.array([-2 + 52j, 1 + 8j], dtype=dtype))

  @unittest.skip("IdentityN is temporarily CompilationOnly as workaround")
  def testIdentityN(self):
    self._testNAryLists(array_ops.identity_n,
                        [np.array([[1, 2, 3]], dtype=np.float32)],
                        expected=[np.array([[1, 2, 3]], dtype=np.float32)])
    self._testNAryLists(array_ops.identity_n,
                        [np.array([[1, 2], [3, 4]], dtype=np.float32),
                         np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)],
                        expected=[
                            np.array([[1, 2], [3, 4]], dtype=np.float32),
                            np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)])
    self._testNAryLists(array_ops.identity_n,
                        [np.array([[1], [2], [3], [4]], dtype=np.int32),
                         np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)],
                        expected=[
                            np.array([[1], [2], [3], [4]], dtype=np.int32),
                            np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)])

  def testConcat(self):
    self._testNAry(
        lambda x: array_ops.concat(x, 0), [
            np.array(
                [[1, 2, 3], [4, 5, 6]], dtype=np.float32), np.array(
                    [[7, 8, 9], [10, 11, 12]], dtype=np.float32)
        ],
        expected=np.array(
            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=np.float32))

    self._testNAry(
        lambda x: array_ops.concat(x, 1), [
            np.array(
                [[1, 2, 3], [4, 5, 6]], dtype=np.float32), np.array(
                    [[7, 8, 9], [10, 11, 12]], dtype=np.float32)
        ],
        expected=np.array(
            [[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]], dtype=np.float32))

  def testOneHot(self):
    with self.session() as session, self.test_scope():
      indices = array_ops.constant(np.array([[2, 3], [0, 1]], dtype=np.int32))
      op = array_ops.one_hot(indices,
                             np.int32(4),
                             on_value=np.float32(7), off_value=np.float32(3))
      output = session.run(op)
      expected = np.array([[[3, 3, 7, 3], [3, 3, 3, 7]],
                           [[7, 3, 3, 3], [3, 7, 3, 3]]],
                          dtype=np.float32)
      self.assertAllEqual(output, expected)

      op = array_ops.one_hot(indices,
                             np.int32(4),
                             on_value=np.int32(2), off_value=np.int32(1),
                             axis=1)
      output = session.run(op)
      expected = np.array([[[1, 1], [1, 1], [2, 1], [1, 2]],
                           [[2, 1], [1, 2], [1, 1], [1, 1]]],
                          dtype=np.int32)
      self.assertAllEqual(output, expected)

  def testSplitV(self):
    with self.session() as session:
      with self.test_scope():
        output = session.run(
            array_ops.split(np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 0, 1, 2]],
                                     dtype=np.float32),
                            [2, 2], 1))
        expected = [np.array([[1, 2], [5, 6], [9, 0]], dtype=np.float32),
                    np.array([[3, 4], [7, 8], [1, 2]], dtype=np.float32)]
        self.assertAllEqual(output, expected)

  def testSplitVNegativeSizes(self):
    with self.session() as session:
      with self.test_scope():
        with self.assertRaisesRegexp(
            (ValueError, errors.InvalidArgumentError),
            "Split size at index 1 must be >= .*. Got: -2"):
          _ = session.run(
              array_ops.split(np.array([1, 2, 3], dtype=np.float32), [-1, -2],
                              axis=0))

  def testStridedSlice(self):
    self._testNAry(lambda x: array_ops.strided_slice(*x),
                   [np.array([[], [], []], dtype=np.float32),
                    np.array([1, 0], dtype=np.int32),
                    np.array([3, 0], dtype=np.int32),
                    np.array([1, 1], dtype=np.int32)],
                   expected=np.array([[], []], dtype=np.float32))

    if np.int64 in self.int_types:
      self._testNAry(
          lambda x: array_ops.strided_slice(*x), [
              np.array([[], [], []], dtype=np.float32), np.array(
                  [1, 0], dtype=np.int64), np.array([3, 0], dtype=np.int64),
              np.array([1, 1], dtype=np.int64)
          ],
          expected=np.array([[], []], dtype=np.float32))

    self._testNAry(lambda x: array_ops.strided_slice(*x),
                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
                             dtype=np.float32),
                    np.array([1, 1], dtype=np.int32),
                    np.array([3, 3], dtype=np.int32),
                    np.array([1, 1], dtype=np.int32)],
                   expected=np.array([[5, 6], [8, 9]], dtype=np.float32))

    self._testNAry(lambda x: array_ops.strided_slice(*x),
                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
                             dtype=np.float32),
                    np.array([0, 2], dtype=np.int32),
                    np.array([2, 0], dtype=np.int32),
                    np.array([1, -1], dtype=np.int32)],
                   expected=np.array([[3, 2], [6, 5]], dtype=np.float32))

    self._testNAry(lambda x: x[0][0:2, array_ops.newaxis, ::-1],
                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
                             dtype=np.float32)],
                   expected=np.array([[[3, 2, 1]], [[6, 5, 4]]],
                                     dtype=np.float32))

    self._testNAry(lambda x: x[0][1, :, array_ops.newaxis],
                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
                             dtype=np.float32)],
                   expected=np.array([[4], [5], [6]], dtype=np.float32))

  def testStridedSliceGrad(self):
    # Tests cases where input shape is empty.
    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
                   [np.array([], dtype=np.int32),
                    np.array([], dtype=np.int32),
                    np.array([], dtype=np.int32),
                    np.array([], dtype=np.int32),
                    np.float32(0.5)],
                   expected=np.array(np.float32(0.5), dtype=np.float32))

    # Tests case where input shape is non-empty, but gradients are empty.
    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
                   [np.array([3], dtype=np.int32),
                    np.array([0], dtype=np.int32),
                    np.array([0], dtype=np.int32),
                    np.array([1], dtype=np.int32),
                    np.array([], dtype=np.float32)],
                   expected=np.array([0, 0, 0], dtype=np.float32))

    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
                   [np.array([3, 0], dtype=np.int32),
                    np.array([1, 0], dtype=np.int32),
                    np.array([3, 0], dtype=np.int32),
                    np.array([1, 1], dtype=np.int32),
                    np.array([[], []], dtype=np.float32)],
                   expected=np.array([[], [], []], dtype=np.float32))

    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
                   [np.array([3, 3], dtype=np.int32),
                    np.array([1, 1], dtype=np.int32),
                    np.array([3, 3], dtype=np.int32),
                    np.array([1, 1], dtype=np.int32),
                    np.array([[5, 6], [8, 9]], dtype=np.float32)],
                   expected=np.array([[0, 0, 0], [0, 5, 6], [0, 8, 9]],
                                     dtype=np.float32))

    def ssg_test(x):
      return array_ops.strided_slice_grad(*x, shrink_axis_mask=0x4,
                                          new_axis_mask=0x1)

    self._testNAry(ssg_test,
                   [np.array([3, 1, 3], dtype=np.int32),
                    np.array([0, 0, 0, 2], dtype=np.int32),
                    np.array([0, 3, 1, -4], dtype=np.int32),
                    np.array([1, 2, 1, -3], dtype=np.int32),
                    np.array([[[1], [2]]], dtype=np.float32)],
                   expected=np.array([[[0, 0, 1]], [[0, 0, 0]], [[0, 0, 2]]],
                                     dtype=np.float32))

    ssg_test2 = lambda x: array_ops.strided_slice_grad(*x, new_axis_mask=0x15)
    self._testNAry(ssg_test2,
                   [np.array([4, 4], dtype=np.int32),
                    np.array([0, 0, 0, 1, 0], dtype=np.int32),
                    np.array([0, 3, 0, 4, 0], dtype=np.int32),
                    np.array([1, 2, 1, 2, 1], dtype=np.int32),
                    np.array([[[[[1], [2]]], [[[3], [4]]]]], dtype=np.float32)],
                   expected=np.array([[0, 1, 0, 2], [0, 0, 0, 0], [0, 3, 0, 4],
                                      [0, 0, 0, 0]], dtype=np.float32))

    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
                   [np.array([3, 3], dtype=np.int32),
                    np.array([0, 2], dtype=np.int32),
                    np.array([2, 0], dtype=np.int32),
                    np.array([1, -1], dtype=np.int32),
                    np.array([[1, 2], [3, 4]], dtype=np.float32)],
                   expected=np.array([[0, 2, 1], [0, 4, 3], [0, 0, 0]],
                                     dtype=np.float32))

    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
                   [np.array([3, 3], dtype=np.int32),
                    np.array([2, 2], dtype=np.int32),
                    np.array([0, 1], dtype=np.int32),
                    np.array([-1, -2], dtype=np.int32),
                    np.array([[1], [2]], dtype=np.float32)],
                   expected=np.array([[0, 0, 0], [0, 0, 2], [0, 0, 1]],
                                     dtype=np.float32))

if __name__ == "__main__":
  googletest.main()
