text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  """Tests for common.utils.
  """
  
  from collections import Counter
  import random
  import tempfile
  import numpy as np
  import tensorflow as tf
  
  from common import utils  # brain coder
  
  
  class UtilsTest(tf.test.TestCase):
  
    def testStackPad(self):
      # 1D.
      tensors = [[1, 2, 3], [4, 5, 6, 7, 8], [9]]
      result = utils.stack_pad(tensors, pad_axes=0, pad_to_lengths=6)
      self.assertTrue(np.array_equal(
          result,
          np.asarray([[1, 2, 3, 0, 0, 0],
                      [4, 5, 6, 7, 8, 0],
                      [9, 0, 0, 0, 0, 0]], dtype=np.float32)))
  
      # 3D.
      tensors = [[[[1, 2, 3], [4, 5, 6]]],
                 [[[7, 8, 9], [0, 1, 2]], [[3, 4, 5], [6, 7, 8]]],
                 [[[0, 1, 2]], [[3, 4, 5]]]]
      result = utils.stack_pad(tensors, pad_axes=[0, 1], pad_to_lengths=[2, 2])
      self.assertTrue(np.array_equal(
          result,
          np.asarray([[[[1, 2, 3], [4, 5, 6]],
                       [[0, 0, 0], [0, 0, 0]]],
                      [[[7, 8, 9], [0, 1, 2]],
                       [[3, 4, 5], [6, 7, 8]]],
                      [[[0, 1, 2], [0, 0, 0]],
                       [[3, 4, 5], [0, 0, 0]]]], dtype=np.float32)))
  
    def testStackPadNoAxes(self):
      # 2D.
      tensors = [[[1, 2, 3], [4, 5, 6]],
                 [[7, 8, 9], [1, 2, 3]],
                 [[4, 5, 6], [7, 8, 9]]]
      result = utils.stack_pad(tensors)
      self.assertTrue(np.array_equal(
          result,
          np.asarray(tensors)))
  
    def testStackPadNoneLength(self):
      # 1D.
      tensors = [[1, 2, 3], [4, 5, 6, 7, 8], [9]]
      result = utils.stack_pad(tensors, pad_axes=0, pad_to_lengths=None)
      self.assertTrue(np.array_equal(
          result,
          np.asarray([[1, 2, 3, 0, 0],
                      [4, 5, 6, 7, 8],
                      [9, 0, 0, 0, 0]], dtype=np.float32)))
  
      # 3D.
      tensors = [[[[1, 2, 3], [4, 5, 6]]],
                 [[[7, 8, 9], [0, 1, 2]], [[3, 4, 5], [6, 7, 8]]],
                 [[[0, 1, 2]], [[3, 4, 5]]]]
      result = utils.stack_pad(tensors, pad_axes=[0, 1], pad_to_lengths=None)
      self.assertTrue(np.array_equal(
          result,
          np.asarray([[[[1, 2, 3], [4, 5, 6]],
                       [[0, 0, 0], [0, 0, 0]]],
                      [[[7, 8, 9], [0, 1, 2]],
                       [[3, 4, 5], [6, 7, 8]]],
                      [[[0, 1, 2], [0, 0, 0]],
                       [[3, 4, 5], [0, 0, 0]]]], dtype=np.float32)))
  
      # 3D with partial pad_to_lengths.
      tensors = [[[[1, 2, 3], [4, 5, 6]]],
                 [[[7, 8, 9], [0, 1, 2]], [[3, 4, 5], [6, 7, 8]]],
                 [[[0, 1, 2]], [[3, 4, 5]]]]
      result = utils.stack_pad(tensors, pad_axes=[0, 1], pad_to_lengths=[None, 3])
      self.assertTrue(np.array_equal(
          result,
          np.asarray([[[[1, 2, 3], [4, 5, 6], [0, 0, 0]],
                       [[0, 0, 0], [0, 0, 0], [0, 0, 0]]],
                      [[[7, 8, 9], [0, 1, 2], [0, 0, 0]],
                       [[3, 4, 5], [6, 7, 8], [0, 0, 0]]],
                      [[[0, 1, 2], [0, 0, 0], [0, 0, 0]],
                       [[3, 4, 5], [0, 0, 0], [0, 0, 0]]]], dtype=np.float32)))
  
    def testStackPadValueError(self):
      # 3D.
      tensors = [[[[1, 2, 3], [4, 5, 6]]],
                 [[[7, 8, 9], [0, 1, 2]], [[3, 4, 5], [6, 7, 8]]],
                 [[[0, 1, 2]], [[3, 4, 5]]],
                 [[[1, 2, 3, 4]]]]
  
      # Not all tensors have the same shape along axis 2.
      with self.assertRaises(ValueError):
        utils.stack_pad(tensors, pad_axes=[0, 1], pad_to_lengths=[2, 2])
  
    def testRecord(self):
      my_record = utils.make_record('my_record', ['a', 'b', 'c'], {'b': 55})
      inst = my_record(a=1, b=2, c=3)
      self.assertEqual(1, inst.a)
      self.assertEqual(2, inst.b)
      self.assertEqual(3, inst.c)
      self.assertEqual(1, inst[0])
      self.assertEqual(2, inst[1])
      self.assertEqual(3, inst[2])
      self.assertEqual([1, 2, 3], list(iter(inst)))
      self.assertEqual(3, len(inst))
  
      inst.b = 999
      self.assertEqual(999, inst.b)
      self.assertEqual(999, inst[1])
  
      inst2 = my_record(1, 999, 3)
      self.assertTrue(inst == inst2)
      inst2[1] = 3
      self.assertFalse(inst == inst2)
  
      inst3 = my_record(a=1, c=3)
      inst.b = 55
      self.assertEqual(inst, inst3)
  
    def testRecordUnique(self):
      record1 = utils.make_record('record1', ['a', 'b', 'c'])
      record2 = utils.make_record('record2', ['a', 'b', 'c'])
      self.assertNotEqual(record1(1, 2, 3), record2(1, 2, 3))
      self.assertEqual(record1(1, 2, 3), record1(1, 2, 3))
  
    def testTupleToRecord(self):
      my_record = utils.make_record('my_record', ['a', 'b', 'c'])
      inst = utils.tuple_to_record((5, 6, 7), my_record)
      self.assertEqual(my_record(5, 6, 7), inst)
  
    def testRecordErrors(self):
      my_record = utils.make_record('my_record', ['a', 'b', 'c'], {'b': 10})
  
      with self.assertRaises(ValueError):
        my_record(c=5)  # Did not provide required argument 'a'.
      with self.assertRaises(ValueError):
        my_record(1, 2, 3, 4)  # Too many arguments.
  
    def testRandomQueue(self):
      np.random.seed(567890)
      queue = utils.RandomQueue(5)
      queue.push(5)
      queue.push(6)
      queue.push(7)
      queue.push(8)
      queue.push(9)
      queue.push(10)
      self.assertTrue(5 not in queue)
      sample = queue.random_sample(1000)
      self.assertEqual(1000, len(sample))
      self.assertEqual([6, 7, 8, 9, 10], sorted(np.unique(sample).tolist()))
  
    def testMaxUniquePriorityQueue(self):
      queue = utils.MaxUniquePriorityQueue(5)
      queue.push(1.0, 'string 1')
      queue.push(-0.5, 'string 2')
      queue.push(0.5, 'string 3')
      self.assertEqual((-0.5, 'string 2', None), queue.pop())
      queue.push(0.1, 'string 4')
      queue.push(1.5, 'string 5')
      queue.push(0.0, 'string 6')
      queue.push(0.2, 'string 7')
      self.assertEqual((1.5, 'string 5', None), queue.get_max())
      self.assertEqual((0.1, 'string 4', None), queue.get_min())
      self.assertEqual(
          [('string 5', None), ('string 1', None), ('string 3', None),
           ('string 7', None), ('string 4', None)],
          list(queue.iter_in_order()))
  
    def testMaxUniquePriorityQueue_Duplicates(self):
      queue = utils.MaxUniquePriorityQueue(5)
      queue.push(0.0, 'string 1')
      queue.push(0.0, 'string 2')
      queue.push(0.0, 'string 3')
      self.assertEqual((0.0, 'string 1', None), queue.pop())
      self.assertEqual((0.0, 'string 2', None), queue.pop())
      self.assertEqual((0.0, 'string 3', None), queue.pop())
      self.assertEqual(0, len(queue))
      queue.push(0.1, 'string 4')
      queue.push(1.5, 'string 5')
      queue.push(0.3, 'string 6')
      queue.push(0.2, 'string 7')
      queue.push(0.0, 'string 8')
      queue.push(1.5, 'string 5')
      queue.push(1.5, 'string 5')
      self.assertEqual((1.5, 'string 5', None), queue.get_max())
      self.assertEqual((0.0, 'string 8', None), queue.get_min())
      self.assertEqual(
          [('string 5', None), ('string 6', None), ('string 7', None),
           ('string 4', None), ('string 8', None)],
          list(queue.iter_in_order()))
  
    def testMaxUniquePriorityQueue_ExtraData(self):
      queue = utils.MaxUniquePriorityQueue(5)
      queue.push(1.0, 'string 1', [1, 2, 3])
      queue.push(0.5, 'string 2', [4, 5, 6])
      queue.push(0.5, 'string 3', [7, 8, 9])
      queue.push(0.5, 'string 2', [10, 11, 12])
      self.assertEqual((0.5, 'string 2', [4, 5, 6]), queue.pop())
      self.assertEqual((0.5, 'string 3', [7, 8, 9]), queue.pop())
      self.assertEqual((1.0, 'string 1', [1, 2, 3]), queue.pop())
      self.assertEqual(0, len(queue))
      queue.push(0.5, 'string 2', [10, 11, 12])
      self.assertEqual((0.5, 'string 2', [10, 11, 12]), queue.pop())
  
    def testRouletteWheel(self):
      random.seed(12345678987654321)
      r = utils.RouletteWheel()
      self.assertTrue(r.is_empty())
      with self.assertRaises(RuntimeError):
        r.sample()  # Cannot sample when empty.
      self.assertEqual(0, r.total_weight)
      self.assertEqual(True, r.add('a', 0.1))
      self.assertFalse(r.is_empty())
      self.assertEqual(0.1, r.total_weight)
      self.assertEqual(True, r.add('b', 0.01))
      self.assertEqual(0.11, r.total_weight)
      self.assertEqual(True, r.add('c', 0.5))
      self.assertEqual(True, r.add('d', 0.1))
      self.assertEqual(True, r.add('e', 0.05))
      self.assertEqual(True, r.add('f', 0.03))
      self.assertEqual(True, r.add('g', 0.001))
      self.assertEqual(0.791, r.total_weight)
      self.assertFalse(r.is_empty())
  
      # Check that sampling is correct.
      obj, weight = r.sample()
      self.assertTrue(isinstance(weight, float), 'Type: %s' % type(weight))
      self.assertTrue((obj, weight) in r)
      for obj, weight in r.sample_many(100):
        self.assertTrue(isinstance(weight, float), 'Type: %s' % type(weight))
        self.assertTrue((obj, weight) in r)
  
      # Check that sampling distribution is correct.
      n = 1000000
      c = Counter(r.sample_many(n))
      for obj, w in r:
        estimated_w = c[(obj, w)] / float(n) * r.total_weight
        self.assertTrue(
            np.isclose(w, estimated_w, atol=1e-3),
            'Expected %s, got %s, for object %s' % (w, estimated_w, obj))
  
    def testRouletteWheel_AddMany(self):
      random.seed(12345678987654321)
      r = utils.RouletteWheel()
      self.assertTrue(r.is_empty())
      with self.assertRaises(RuntimeError):
        r.sample()  # Cannot sample when empty.
      self.assertEqual(0, r.total_weight)
      count = r.add_many(
          ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
          [0.1, 0.01, 0.5, 0.1, 0.05, 0.03, 0.001])
      self.assertEqual(7, count)
      self.assertFalse(r.is_empty())
      self.assertEqual(0.791, r.total_weight)
  
      # Adding no items is allowed.
      count = r.add_many([], [])
      self.assertEqual(0, count)
      self.assertFalse(r.is_empty())
      self.assertEqual(0.791, r.total_weight)
  
      # Check that sampling is correct.
      obj, weight = r.sample()
      self.assertTrue(isinstance(weight, float), 'Type: %s' % type(weight))
      self.assertTrue((obj, weight) in r)
      for obj, weight in r.sample_many(100):
        self.assertTrue(isinstance(weight, float), 'Type: %s' % type(weight))
        self.assertTrue((obj, weight) in r)
  
      # Check that sampling distribution is correct.
      n = 1000000
      c = Counter(r.sample_many(n))
      for obj, w in r:
        estimated_w = c[(obj, w)] / float(n) * r.total_weight
        self.assertTrue(
            np.isclose(w, estimated_w, atol=1e-3),
            'Expected %s, got %s, for object %s' % (w, estimated_w, obj))
  
    def testRouletteWheel_AddZeroWeights(self):
      r = utils.RouletteWheel()
      self.assertEqual(True, r.add('a', 0))
      self.assertFalse(r.is_empty())
      self.assertEqual(4, r.add_many(['b', 'c', 'd', 'e'], [0, 0.1, 0, 0]))
      self.assertEqual(
          [('a', 0.0), ('b', 0.0), ('c', 0.1), ('d', 0.0), ('e', 0.0)],
          list(r))
  
    def testRouletteWheel_UniqueMode(self):
      random.seed(12345678987654321)
      r = utils.RouletteWheel(unique_mode=True)
      self.assertEqual(True, r.add([1, 2, 3], 1, 'a'))
      self.assertEqual(True, r.add([4, 5], 0.5, 'b'))
      self.assertEqual(False, r.add([1, 2, 3], 1.5, 'a'))
      self.assertEqual(
          [([1, 2, 3], 1.0), ([4, 5], 0.5)],
          list(r))
      self.assertEqual(1.5, r.total_weight)
      self.assertEqual(
          2,
          r.add_many(
              [[5, 6, 2, 3], [1, 2, 3], [8], [1, 2, 3]],
              [0.1, 0.2, 0.1, 2.0],
              ['c', 'a', 'd', 'a']))
      self.assertEqual(
          [([1, 2, 3], 1.0), ([4, 5], 0.5), ([5, 6, 2, 3], 0.1), ([8], 0.1)],
          list(r))
      self.assertTrue(np.isclose(1.7, r.total_weight))
      self.assertEqual(0, r.add_many([], [], []))  # Adding no items is allowed.
      with self.assertRaises(ValueError):
        # Key not given.
        r.add([7, 8, 9], 2.0)
      with self.assertRaises(ValueError):
        # Keys not given.
        r.add_many([[7, 8, 9], [10]], [2.0, 2.0])
      self.assertEqual(True, r.has_key('a'))
      self.assertEqual(True, r.has_key('b'))
      self.assertEqual(False, r.has_key('z'))
      self.assertEqual(1.0, r.get_weight('a'))
      self.assertEqual(0.5, r.get_weight('b'))
  
      r = utils.RouletteWheel(unique_mode=False)
      self.assertEqual(True, r.add([1, 2, 3], 1))
      self.assertEqual(True, r.add([4, 5], 0.5))
      self.assertEqual(True, r.add([1, 2, 3], 1.5))
      self.assertEqual(
          [([1, 2, 3], 1.0), ([4, 5], 0.5), ([1, 2, 3], 1.5)],
          list(r))
      self.assertEqual(3, r.total_weight)
      self.assertEqual(
          4,
          r.add_many(
              [[5, 6, 2, 3], [1, 2, 3], [8], [1, 2, 3]],
              [0.1, 0.2, 0.1, 0.2]))
      self.assertEqual(
          [([1, 2, 3], 1.0), ([4, 5], 0.5), ([1, 2, 3], 1.5),
           ([5, 6, 2, 3], 0.1), ([1, 2, 3], 0.2), ([8], 0.1), ([1, 2, 3], 0.2)],
          list(r))
      self.assertTrue(np.isclose(3.6, r.total_weight))
      with self.assertRaises(ValueError):
        # Key is given.
        r.add([7, 8, 9], 2.0, 'a')
      with self.assertRaises(ValueError):
        # Keys are given.
        r.add_many([[7, 8, 9], [10]], [2.0, 2.0], ['a', 'b'])
  
    def testRouletteWheel_IncrementalSave(self):
      f = tempfile.NamedTemporaryFile()
      r = utils.RouletteWheel(unique_mode=True, save_file=f.name)
      entries = [
          ([1, 2, 3], 0.1, 'a'),
          ([4, 5], 0.2, 'b'),
          ([6], 0.3, 'c'),
          ([7, 8, 9, 10], 0.25, 'd'),
          ([-1, -2], 0.15, 'e'),
          ([-3, -4, -5], 0.5, 'f')]
  
      self.assertTrue(r.is_empty())
      for i in range(0, len(entries), 2):
        r.add(*entries[i])
        r.add(*entries[i + 1])
        r.incremental_save()
  
        r2 = utils.RouletteWheel(unique_mode=True, save_file=f.name)
        self.assertEqual(i + 2, len(r2))
        count = 0
        for j, (obj, weight) in enumerate(r2):
          self.assertEqual(entries[j][0], obj)
          self.assertEqual(entries[j][1], weight)
          self.assertEqual(weight, r2.get_weight(entries[j][2]))
          count += 1
        self.assertEqual(i + 2, count)
  
  if __name__ == '__main__':
    tf.test.main()
  