# 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 tensorflow.ops.nn_ops.Pad."""

import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.platform import test


class PadOpTest(test.TestCase):

  def setUp(self):
    super().setUp()
    # Create mapping between TensorFlow quantized types and numpy types.
    self._quint8 = np.dtype([("quint8", np.uint8)])
    self._qint8 = np.dtype([("qint8", np.int8)])
    self._qint32 = np.dtype([("qint32", np.int32)])

  def _npPad(self, inp, paddings, mode, constant_values=0):
    mode = mode.lower()
    if mode == "constant":
      return np.pad(inp, paddings, mode=mode, constant_values=constant_values)
    else:
      return np.pad(inp, paddings, mode=mode)

  def testNpPad(self):
    self.assertAllEqual(
        np.array([[0, 0, 0, 0, 0, 0],
                  [0, 3, 3, 0, 0, 0],
                  [0, 4, 4, 0, 0, 0],
                  [0, 5, 5, 0, 0, 0],
                  [0, 0, 0, 0, 0, 0],
                  [0, 0, 0, 0, 0, 0]]),
        self._npPad(
            np.array([[3, 3], [4, 4], [5, 5]]),
            [[1, 2], [1, 3]],
            mode="constant"))

    self.assertAllEqual(
        np.array([[1, 1, 1, 1, 1, 1],
                  [1, 3, 3, 1, 1, 1],
                  [1, 4, 4, 1, 1, 1],
                  [1, 5, 5, 1, 1, 1],
                  [1, 1, 1, 1, 1, 1],
                  [1, 1, 1, 1, 1, 1]]),
        self._npPad(
            np.array([[3, 3], [4, 4], [5, 5]]),
            [[1, 2], [1, 3]],
            mode="constant", constant_values=1))

    self.assertAllEqual(
        np.array([[4, 3, 4, 9, 4, 3],
                  [1, 0, 1, 2, 1, 0],
                  [4, 3, 4, 9, 4, 3],
                  [1, 0, 1, 2, 1, 0]]),
        self._npPad(
            np.array([[0, 1, 2], [3, 4, 9]]),
            [[1, 1], [1, 2]],
            mode="reflect"))

    self.assertAllEqual(
        np.array([[0, 0, 1, 2, 2, 1],
                  [0, 0, 1, 2, 2, 1],
                  [3, 3, 4, 9, 9, 4],
                  [3, 3, 4, 9, 9, 4]]),
        self._npPad(
            np.array([[0, 1, 2], [3, 4, 9]]),
            [[1, 1], [1, 2]],
            mode="symmetric"))

  def _testPad(self, np_inputs, paddings, mode, constant_values):
    np_val = self._npPad(np_inputs, paddings, mode=mode,
                         constant_values=constant_values)
    for use_gpu in [True, False]:
      with test_util.device(use_gpu=use_gpu):
        tf_val = array_ops.pad(np_inputs, paddings, mode=mode,
                               constant_values=constant_values)
        out = self.evaluate(tf_val)

        if np_inputs.dtype in [self._qint8, self._quint8, self._qint32]:
          # Cast quantized types back to their numpy equivalents.
          np_val = np_val.astype(np_inputs.dtype[0])

      self.assertAllEqual(np_val, out)
      self.assertShapeEqual(np_val, tf_val)

  def _testGradient(self,
                    x,
                    a,
                    mode,
                    constant_values,
                    paddings_dtype=dtypes.int32):

    def pad(x):
      return array_ops.pad(
          x,
          ops.convert_to_tensor(a, paddings_dtype),
          mode=mode,
          constant_values=constant_values)

    with self.cached_session():
      jacob_t, jacob_n = gradient_checker_v2.compute_gradient(pad, [x])
      tol = 1e-3 if x.dtype == np.float16 else 4e-5
      self.assertAllClose(jacob_t, jacob_n, rtol=tol, atol=tol)

  def _testAll(self, np_inputs, paddings, constant_values):
    for mode in ("CONSTANT", "REFLECT", "SYMMETRIC", "reflect", "symmetric",
                 "constant"):
      # Zero-sized input is not allowed for REFLECT mode, but we still want
      # zero-sized input test cases for the other modes.
      if np_inputs.size or mode.upper() != "REFLECT":
        self._testPad(np_inputs, paddings, mode=mode,
                      constant_values=constant_values)
        if np_inputs.dtype == np.float32:
          self._testGradient(np_inputs, paddings, mode=mode,
                             constant_values=constant_values)

  def testInputDims(self):
    with test_util.use_gpu():
      with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                  "Shape must be rank 1 but is rank 6|"
                                  "paddings must be the rank of inputs"):
        array_ops.pad(array_ops.reshape(
            [1, 2], shape=[1, 2, 1, 1, 1, 1]),
                      array_ops.reshape(
                          [1, 2], shape=[1, 2]))

  def testPaddingsDim(self):
    with test_util.use_gpu():
      with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                  "Shape must be rank 2 but is rank 1|"
                                  "paddings must be a matrix with 2 columns"):
        array_ops.pad(array_ops.reshape(
            [1, 2], shape=[1, 2]),
                      array_ops.reshape(
                          [1, 2], shape=[2]))

  def testPaddingsDim2(self):
    with test_util.use_gpu():
      with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                  "Dimension must be 2 but is 1|"
                                  "paddings must be a matrix with 2 columns"):