"""Functional tests for pooling operations."""
from __future__ import print_function
import tensorflow.python.platform

import numpy as np
import tensorflow as tf

from tensorflow.python.kernel_tests import gradient_checker as gc
from tensorflow.python.ops import gen_nn_ops


def GetInceptionMaxPoolShapes():
  """Iterator for some of the max pool ops in the Inception 2015 model.

  Yields:
    Tuple (name, input_size, filter_size, out_size, strides, padding)
  """
  names = ["maxpool2", "maxpool3", "maxpool4", "maxpool5"]
  input_sizes = [[32, 71, 71, 192],
                 [32, 35, 35, 288], [32, 17, 17, 1248], [32, 8, 8, 2048]]
  filter_sizes = [[1, 3, 3, 1], [1, 3, 3, 1],
                  [1, 3, 3, 1], [1, 3, 3, 1]]
  output_sizes = [[32, 35, 35, 192], [32, 17, 17, 288],
                  [32, 8, 8, 1248], [32, 8, 8, 2048]]
  strides = [[1, 2, 2, 1], [1, 2, 2, 1], [1, 2, 2, 1],
             [1, 1, 1, 1]]
  paddings = ["VALID", "VALID", "VALID", "SAME"]
  for n, i, f, o, s, p in zip(names, input_sizes, filter_sizes, output_sizes,
                              strides, paddings):
    yield n, i, f, o, s, p


class PoolingTest(tf.test.TestCase):

  def _VerifyValues(self, pool_func, input_sizes, ksize, strides, padding,
                    expected, use_gpu):
    """Verifies the output values of the pooling function.

    Args:
      pool_func: Function to be called, co.MaxPool, co.AvgPool,
        or the Lua version.
      input_sizes: Input tensor dimensions.
      ksize: The kernel size dimensions
      strides: The stride dimensions
      padding: Padding type.
      expected: An array containing the expected operation outputs.
      use_gpu: Whether we are running on GPU.
    """
    total_size = 1
    for s in input_sizes:
      total_size *= s
    # Initializes the input tensor with array containing incrementing
    # numbers from 1.
    x = [f * 1.0 for f in range(1, total_size + 1)]
    with self.test_session(use_gpu=use_gpu) as sess:
      t = tf.constant(x, shape=input_sizes)
      t = pool_func(t, ksize=ksize, strides=strides, padding=padding)
      actual = t.eval()
      self.assertAllClose(expected, actual.flatten())
      self.assertShapeEqual(actual, t)

  def _testAvgPoolValidPadding(self, use_gpu):
    expected_output = [7.0, 8.0, 9.0]
    self._VerifyValues(tf.nn.avg_pool, input_sizes=[1, 3, 3, 3],
                       ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                       padding="VALID",
                       expected=expected_output, use_gpu=use_gpu)

  def _testAvgPoolSamePadding(self, use_gpu):
    expected_output = [8.5, 9.5, 10.5, 14.5, 15.5, 16.5]
    self._VerifyValues(tf.nn.avg_pool, input_sizes=[1, 2, 4, 3],
                       ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],