"""Functional tests for slice op."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow.python.platform

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


class SliceTest(tf.test.TestCase):

  def _testEmpty(self, use_gpu):
    inp = np.random.rand(4, 4).astype("f")
    for k in xrange(4):
      with self.test_session(use_gpu=use_gpu):
        a = tf.constant(inp, shape=[4, 4], dtype=tf.float32)
        slice_t = a[2, k:k]
        slice_val = slice_t.eval()
      self.assertAllEqual(slice_val, inp[2, k:k])

  def testEmptyAll(self):
    self._testEmpty(use_gpu=False)
    self._testEmpty(use_gpu=True)

  def _testInt32(self, use_gpu):
    inp = np.random.rand(4, 4).astype("i")
    for k in xrange(4):
      with self.test_session(use_gpu=use_gpu):
        a = tf.constant(inp, shape=[4, 4], dtype=tf.int32)
        slice_t = a[2, k:k]
        slice_val = slice_t.eval()
      self.assertAllEqual(slice_val, inp[2, k:k])

  def testInt32(self):
    self._testEmpty(use_gpu=False)
    self._testEmpty(use_gpu=True)

  def _testSelectAll(self, use_gpu):
    with self.test_session(use_gpu=use_gpu):
      inp = np.random.rand(4, 4, 4, 4).astype("f")
      a = tf.constant(inp, shape=[4, 4, 4, 4],
                               dtype=tf.float32)

      slice_explicit_t = tf.slice(a, [0, 0, 0, 0], [-1, -1, -1, -1])
      slice_implicit_t = a[:, :, :, :]

      self.assertAllEqual(inp, slice_explicit_t.eval())
      self.assertAllEqual(inp, slice_implicit_t.eval())
      self.assertEqual(inp.shape, slice_explicit_t.get_shape())
      self.assertEqual(inp.shape, slice_implicit_t.get_shape())

  def testSelectAll(self):
    for _ in range(10):
      self._testSelectAll(use_gpu=False)
      self._testSelectAll(use_gpu=True)

  def _testSingleDimension(self, use_gpu):
    with self.test_session(use_gpu=use_gpu):
      inp = np.random.rand(10).astype("f")
      a = tf.constant(inp, shape=[10], dtype=tf.float32)

      hi = np.random.random_integers(0, 9)
      scalar_t = a[hi]
      scalar_val = scalar_t.eval()
      self.assertAllEqual(scalar_val, inp[hi])

      lo = np.random.random_integers(0, hi)
      slice_t = a[lo:hi]
      slice_val = slice_t.eval()
      self.assertAllEqual(slice_val, inp[lo:hi])

  def testSingleDimension(self):
    for _ in range(10):
      self._testSingleDimension(use_gpu=False)
      self._testSingleDimension(use_gpu=True)

  def _testSliceMatrixDim0(self, x, begin, size, use_gpu):
    with self.test_session(use_gpu=use_gpu):
      tf_ans = tf.slice(x, [begin, 0], [size, x.shape[1]]).eval()
    np_ans = x[begin:begin+size, :]
    self.assertAllEqual(tf_ans, np_ans)

  def testSliceMatrixDim0(self):
    for use_gpu in [False, True]:
      x = np.random.rand(8, 4).astype("f")
      self._testSliceMatrixDim0(x, 1, 2, use_gpu)
      self._testSliceMatrixDim0(x, 3, 3, use_gpu)
      y = np.random.rand(8, 7).astype("f")    # 7 * sizeof(float) is not aligned
      self._testSliceMatrixDim0(y, 1, 2, use_gpu)
      self._testSliceMatrixDim0(y, 3, 3, use_gpu)

  def _testIndexAndSlice(self, use_gpu):
    with self.test_session(use_gpu=use_gpu):