# 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 tf.layers.convolutional."""

import numpy as np

from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.keras.legacy_tf_layers import convolutional as conv_layers
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_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


class ConvTest(test.TestCase):

  def testInvalidDataFormat(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 3), seed=1)
    with self.assertRaisesRegex(ValueError, 'data_format'):
      conv_layers.conv2d(images, 32, 3, data_format='invalid')

  def testInvalidStrides(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 3), seed=1)
    with self.assertRaisesRegex(ValueError, 'strides'):
      conv_layers.conv2d(images, 32, 3, strides=(1, 2, 3))

    with self.assertRaisesRegex(ValueError, 'strides'):
      conv_layers.conv2d(images, 32, 3, strides=None)

  def testInvalidKernelSize(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 3), seed=1)
    with self.assertRaisesRegex(ValueError, 'kernel_size'):
      conv_layers.conv2d(images, 32, (1, 2, 3))

    with self.assertRaisesRegex(ValueError, 'kernel_size'):
      conv_layers.conv2d(images, 32, None)

  def testCreateConv2D(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 4))
    layer = conv_layers.Conv2D(32, [3, 3], activation=nn_ops.relu)
    output = layer.apply(images)
    if not context.executing_eagerly():
      self.assertEqual(output.op.name, 'conv2d/Relu')
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height - 2, width - 2, 32])
    self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 3, 4, 32])
    self.assertListEqual(layer.bias.get_shape().as_list(), [32])

  def testConv2DFloat16(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 4), dtype='float16')
    output = conv_layers.conv2d(images, 32, [3, 3], activation=nn_ops.relu)
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height - 2, width - 2, 32])

  def testCreateConv2DIntegerKernelSize(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 4))
    layer = conv_layers.Conv2D(32, 3)
    output = layer.apply(images)
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height - 2, width - 2, 32])
    self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 3, 4, 32])
    self.assertListEqual(layer.bias.get_shape().as_list(), [32])

  def testCreateConv2DChannelsFirst(self):
    with ops.Graph().as_default():
      height, width = 7, 9
      images = random_ops.random_uniform((5, 4, height, width))
      layer = conv_layers.Conv2D(32, [3, 3], data_format='channels_first')
      output = layer.apply(images)
      self.assertListEqual(output.get_shape().as_list(),
                           [5, 32, height - 2, width - 2])
      self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 3, 4, 32])
      self.assertListEqual(layer.bias.get_shape().as_list(), [32])

  def testUnknownInputChannels(self):
    with ops.Graph().as_default():
      images = array_ops.placeholder(dtypes.float32, (5, 7, 9, None))
      layer = conv_layers.Conv2D(32, [3, 3], activation=nn_ops.relu)
      with self.assertRaisesRegex(
          ValueError, 'The channel dimension of the inputs '
          'should be defined. Found `None`.'):
        _ = layer.apply(images)

      images = array_ops.placeholder(dtypes.float32, (5, None, 7, 9))
      layer = conv_layers.Conv2D(32, [3, 3], data_format='channels_first')
      with self.assertRaisesRegex(
          ValueError, 'The channel dimension of the inputs '
          'should be defined. Found `None`.'):
        _ = layer.apply(images)

  def testConv2DPaddingSame(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 32), seed=1)
    layer = conv_layers.Conv2D(64, images.get_shape()[1:3], padding='same')
    output = layer.apply(images)
    self.assertListEqual(output.get_shape().as_list(), [5, height, width, 64])

  def testCreateConvWithStrides(self):
    height, width = 6, 8
    # Test strides tuple
    images = random_ops.random_uniform((5, height, width, 3), seed=1)
    layer = conv_layers.Conv2D(32, [3, 3], strides=(2, 2), padding='same')
    output = layer.apply(images)
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height / 2, width / 2, 32])

    # Test strides integer
    layer = conv_layers.Conv2D(32, [3, 3], strides=2, padding='same')
    output = layer.apply(images)
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height / 2, width / 2, 32])

    # Test unequal strides
    layer = conv_layers.Conv2D(32, [3, 3], strides=(2, 1), padding='same')
    output = layer.apply(images)
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height / 2, width, 32])

  def testCreateConv1D(self):
    width = 7
    data = random_ops.random_uniform((5, width, 4))
    layer = conv_layers.Conv1D(32, 3, activation=nn_ops.relu)
    output = layer.apply(data)
    if not context.executing_eagerly():
      self.assertEqual(output.op.name, 'conv1d/Relu')
    self.assertListEqual(output.get_shape().as_list(), [5, width - 2, 32])
    self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 4, 32])
    self.assertListEqual(layer.bias.get_shape().as_list(), [32])

  def testConv1DFloat16(self):
    width = 7
    data = random_ops.random_uniform((5, width, 4), dtype='float16')
    output = conv_layers.conv1d(data, 32, 3, activation=nn_ops.relu)
    self.assertListEqual(output.get_shape().as_list(), [5, width - 2, 32])

  def testCreateConv1DChannelsFirst(self):
    with ops.Graph().as_default():
      width = 7
      data = random_ops.random_uniform((5, 4, width))
      layer = conv_layers.Conv1D(32, 3, data_format='channels_first')
      output = layer.apply(data)
      self.assertListEqual(output.get_shape().as_list(), [5, 32, width - 2])
      self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 4, 32])
      self.assertListEqual(layer.bias.get_shape().as_list(), [32])

  def testUnknownInputChannelsConv1D(self):
    with ops.Graph().as_default():
      data = array_ops.placeholder(dtypes.float32, (5, 4, None))
      layer = conv_layers.Conv1D(32, 3, activation=nn_ops.relu)
      with self.assertRaisesRegex(
          ValueError, 'The channel dimension of the inputs '
          'should be defined. Found `None`.'):
        _ = layer.apply(data)

      data = array_ops.placeholder(dtypes.float32, (5, None, 4))
      layer = conv_layers.Conv1D(32, 3, data_format='channels_first')
      with self.assertRaisesRegex(
          ValueError, 'The channel dimension of the inputs '
          'should be defined. Found `None`.'):
        _ = layer.apply(data)

  def testCreateConv3D(self):
    depth, height, width = 6, 7, 9
    volumes = random_ops.random_uniform((5, depth, height, width, 4))
    layer = conv_layers.Conv3D(32, [3, 3, 3], activation=nn_ops.relu)
    output = layer.apply(volumes)
    if not context.executing_eagerly():
      self.assertEqual(output.op.name, 'conv3d/Relu')
    self.assertListEqual(output.get_shape().as_list(),
                         [5, depth - 2, height - 2, width - 2, 32])
    self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 3, 3, 4, 32])
    self.assertListEqual(layer.bias.get_shape().as_list(), [32])

  def testUnknownInputChannelsConv3D(self):
    with ops.Graph().as_default():
      volumes = array_ops.placeholder(dtypes.float32, (5, 6, 7, 9, None))
      layer = conv_layers.Conv3D(32, [3, 3, 3], activation=nn_ops.relu)
      with self.assertRaisesRegex(
          ValueError, 'The channel dimension of the inputs '
          'should be defined. Found `None`.'):
        _ = layer.apply(volumes)

  def testConv2DKernelRegularizer(self):
    with ops.Graph().as_default():
      height, width = 7, 9
      images = random_ops.random_uniform((5, height, width, 4))
      reg = lambda x: 0.1 * math_ops.reduce_sum(x)
      layer = conv_layers.Conv2D(32, [3, 3], kernel_regularizer=reg)
      layer.apply(images)
      loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
      self.assertEqual(len(loss_keys), 1)
      self.evaluate([v.initializer for v in layer.variables])
      self.assertListEqual(
          self.evaluate(layer.losses), self.evaluate(loss_keys))

  def testConv2DBiasRegularizer(self):
    with ops.Graph().as_default():
      height, width = 7, 9
      images = random_ops.random_uniform((5, height, width, 4))
      reg = lambda x: 0.1 * math_ops.reduce_sum(x)
      layer = conv_layers.Conv2D(32, [3, 3], bias_regularizer=reg)
      layer.apply(images)
      loss_keys = ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)
      self.assertEqual(len(loss_keys), 1)
      self.evaluate([v.initializer for v in layer.variables])
      self.assertListEqual(
          self.evaluate(layer.losses), self.evaluate(loss_keys))

  def testConv2DNoBias(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 4))
    layer = conv_layers.Conv2D(
        32, [3, 3], activation=nn_ops.relu, use_bias=False)
    output = layer.apply(images)
    if not context.executing_eagerly():
      self.assertEqual(output.op.name, 'conv2d/Relu')
    self.assertListEqual(output.get_shape().as_list(),
                         [5, height - 2, width - 2, 32])
    self.assertListEqual(layer.kernel.get_shape().as_list(), [3, 3, 4, 32])
    self.assertEqual(layer.bias, None)

  def testDilatedConv2D(self):
    height, width = 7, 9
    images = random_ops.random_uniform((5, height, width, 4))
    layer = conv_layers.Conv2D(32, [3, 3], dilation_rate=3)