# Copyright 2016 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.
# ==============================================================================

# pylint: disable=unused-import
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from tensorflow.contrib.training.python.training import sampling_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import logging_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import coordinator
from tensorflow.python.training import queue_runner_impl


class StratifiedSampleTest(test.TestCase):

  def testGraphBuildAssertionFailures(self):
    val = [array_ops.zeros([1, 3]), array_ops.ones([1, 5])]
    label = constant_op.constant([1], shape=[1])  # must have batch dimension
    probs = [.2] * 5
    init_probs = [.1, .3, .1, .3, .2]
    batch_size = 16

    # Label must have only batch dimension if enqueue_many is True.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(
          val,
          array_ops.zeros([]),
          probs,
          batch_size,
          init_probs,
          enqueue_many=True)
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(
          val,
          array_ops.zeros([1, 1]),
          probs,
          batch_size,
          init_probs,
          enqueue_many=True)

    # Label must not be one-hot.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(val,
                                     constant_op.constant([0, 1, 0, 0, 0]),
                                     probs, batch_size, init_probs)

    # Data must be list, not singleton tensor.
    with self.assertRaises(TypeError):
      sampling_ops.stratified_sample(
          array_ops.zeros([1, 3]), label, probs, batch_size, init_probs)

    # Data must have batch dimension if enqueue_many is True.
    with self.ass
ertRaises(ValueError):
      sampling_ops.stratified_sample(
          val,
          constant_op.constant(1),
          probs,
          batch_size,
          init_probs,
          enqueue_many=True)

    # Batch dimensions on data and labels should be equal.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(

          [array_ops.zeros([2, 1])],
          label,
          probs,
          batch_size,
          init_probs,
          enqueue_many=True)

    # Probabilities must be numpy array, python list, or tensor.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(val, label, 1, batch_size, init_probs)

    # Probabilities shape must be fully defined.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(
          val,
          label,
          array_ops.placeholder(
              dtypes.float32, shape=[None]),
          batch_size,
          init_probs)

    # In the rejection sampling case, make sure that probability lengths are
    # the same.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(
          val, label, [.1] * 10, batch_size, init_probs=[.2] * 5)

    # In the rejection sampling case, make sure that zero initial probability
    # classes also have zero target probability.
    with self.assertRaises(ValueError):
      sampling_ops.stratified_sample(
          val, label, [.2, .4, .4], batch_size, init_probs=[0, .5, .5])

  def testRuntimeAssertionFailures(self):
    valid_probs = [.2] * 5
    valid_labels = [1, 2, 3]
    vals = [array_ops.zeros([3, 1])]

    illegal_labels = [
        [0, -1, 1],  # classes must be nonnegative
        [5, 1, 1],  # classes must be less than number of classes
        [2, 3],  # data and label batch size must be the same
    ]

    illegal_probs = [
        [.1] * 5,  # probabilities must sum to one
        [-.5, .5, .5, .4, .1],  # probabilities must be non-negative
    ]

    # Set up graph with illegal label vector.
    label_ph = array_ops.placeholder(dtypes.int32, shape=[None])
    probs_ph = array_ops.placeholder(
        dtypes.float32, shape=[5])  # shape must be defined
    val_tf, lbl_tf, prob_tf = sampling_ops._verify_input(  # pylint: disable=protected-access
        vals, label_ph, [probs_ph])

    for illegal_label in illegal_labels:
      # Run session that should fail.
      with self.test_session() as sess:
        with self.assertRaises(errors_impl.InvalidArgumentError):
          sess.run([val_tf, lbl_tf],
                   feed_dict={label_ph: illegal_label,
                              probs_ph: valid_probs})

    for illegal_prob in illegal_probs:
      # Run session that should fail.
      with self.test_session() as sess:
        with self.assertRaises(errors_impl.InvalidArgumentError):
          sess.run([prob_tf],
                   feed_dict={label_ph: valid_labels,
                              probs_ph: illegal_prob})

  def testCanBeCalledMultipleTimes(self):
    batch_size = 20
    val_input_batch = [array_ops.zeros([2, 3, 4])]
    lbl_input_batch = array_ops.ones([], dtype=dtypes.int32)
    probs = np.array([0, 1, 0, 0, 0])
    batches = sampling_ops.stratified_sample(
        val_input_batch, lbl_input_batch, probs, batch_size, init_probs=probs)
    batches += sampling_ops.stratified_sample(
        val_input_batch, lbl_input_batch, probs, batch_size, init_probs=probs)
    summary_op = logging_ops.merge_summary(
        ops.get_collection(ops.GraphKeys.SUMMARIES))

    with self.test_session() as sess:
      coord = coordinator.Coordinator()
      threads = queue_runner_impl.start_queue_runners(coord=coord)

      sess.run(batches + (summary_op,))

      coord.request_stop()
      coord.join(threads)

  def testRejectionBatchingBehavior(self):
    batch_size = 20
    input_batch_size = 11
    val_input_batch = [array_ops.zeros([input_batch_size, 2, 3, 4])]
    lbl_input_batch = control_flow_ops.cond(
        math_ops.greater(.5, random_ops.random_uniform([])),
        lambda: array_ops.ones([input_batch_size], dtype=dtypes.int32) * 1,
        lambda: array_ops.ones([input_batch_size], dtype=dtypes.int32) * 3)
    probs = np.array([0, .2, 0, .8, 0])
    data_batch, labels = sampling_ops.stratified_sample(
        val_input_batch,
        lbl_input_batch,
        probs,
        batch_size,
        init_probs=[0, .3, 0, .7, 0],
        enqueue_many=True)
    with self.test_session() as sess:
      coord = coordinator.Coordinator()
      threads = queue_runner_impl.start_queue_runners(coord=coord)

      sess.run([data_batch, labels])

      coord.request_stop()
      coord.join(threads)

  def testBatchDimensionNotRequired(self):
    classes = 5
    # Probs must be a tensor, since we pass it directly to _verify_input.
    probs = constant_op.constant([1.0 / classes] * classes)

    # Make sure that these vals/labels pairs don't throw any runtime exceptions.
    legal_input_pairs = [
        (np.zeros([2, 3]), [x % classes for x in range(2)]),  # batch dim 2
        (np.zeros([4, 15]), [x % classes for x in range(4)]),  # batch dim 4
        (np.zeros([10, 1]), [x % classes for x in range(10)]),  # batc