text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  """Tests for common.reward."""
  
  from math import log
  import numpy as np
  import tensorflow as tf
  
  from common import reward  # brain coder
  
  
  class RewardTest(tf.test.TestCase):
  
    def testAbsDiff(self):
      self.assertEqual(5, reward.abs_diff(15, 20))
      self.assertEqual(5, reward.abs_diff(20, 15))
  
    def testModAbsDiff(self):
      self.assertEqual(5, reward.mod_abs_diff(15, 20, 25))
      self.assertEqual(5, reward.mod_abs_diff(20, 15, 25))
      self.assertEqual(2, reward.mod_abs_diff(1, 24, 25))
      self.assertEqual(2, reward.mod_abs_diff(24, 1, 25))
  
      self.assertEqual(0, reward.mod_abs_diff(0, 0, 5))
      self.assertEqual(1, reward.mod_abs_diff(0, 1, 5))
      self.assertEqual(2, reward.mod_abs_diff(0, 2, 5))
      self.assertEqual(2, reward.mod_abs_diff(0, 3, 5))
      self.assertEqual(1, reward.mod_abs_diff(0, 4, 5))
  
      self.assertEqual(0, reward.mod_abs_diff(-1, 4, 5))
      self.assertEqual(1, reward.mod_abs_diff(-5, 4, 5))
      self.assertEqual(1, reward.mod_abs_diff(-7, 4, 5))
      self.assertEqual(1, reward.mod_abs_diff(13, 4, 5))
      self.assertEqual(1, reward.mod_abs_diff(15, 4, 5))
  
    def testAbsoluteDistance_AbsDiffMethod(self):
      self.assertEqual(
          4,
          reward.absolute_distance([0], [4], 5, scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          0,
          reward.absolute_distance([4], [4], 5, scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          0,
          reward.absolute_distance([], [], 5, scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([1], [], 5, scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([], [1], 5, scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          0,
          reward.absolute_distance([1, 2, 3], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          1,
          reward.absolute_distance([1, 2, 4], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          1,
          reward.absolute_distance([1, 2, 2], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([1, 2], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([1, 2, 3, 4], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.abs_diff))
      self.assertEqual(
          6,
          reward.absolute_distance([4, 4, 4], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.abs_diff))
  
    def testAbsoluteDistance_ModDiffMethod(self):
      self.assertEqual(
          1,
          reward.absolute_distance([0], [4], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          0,
          reward.absolute_distance([4], [4], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          0,
          reward.absolute_distance([], [], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([1], [], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([], [1], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          0,
          reward.absolute_distance([1, 2, 3], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          1,
          reward.absolute_distance([1, 2, 4], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          1,
          reward.absolute_distance([1, 2, 2], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([1, 2], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([1, 2, 3, 4], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
      self.assertEqual(
          5,
          reward.absolute_distance([4, 4, 4], [1, 2, 3], 5,
                                   scalar_diff_fn=reward.mod_abs_diff))
  
    def testLogAbsoluteDistance(self):
      def log_diff(diff, base):
        return log(diff + 1) / log(base // 2 + 2)
  
      self.assertEqual(
          log_diff(1, 5),
          reward.log_absolute_distance([0], [4], 5))
      self.assertEqual(
          log_diff(2, 5),
          reward.log_absolute_distance([1], [4], 5))
      self.assertEqual(
          log_diff(2, 5),
          reward.log_absolute_distance([2], [4], 5))
      self.assertEqual(
          log_diff(1, 5),
          reward.log_absolute_distance([3], [4], 5))
      self.assertEqual(
          log_diff(3, 5),  # max_dist = base // 2 + 1 = 3
          reward.log_absolute_distance([], [4], 5))
      self.assertEqual(
          0 + log_diff(3, 5),  # max_dist = base // 2 + 1 = 3
          reward.log_absolute_distance([4, 4], [4], 5))
      self.assertEqual(
          0,
          reward.log_absolute_distance([4], [4], 5))
      self.assertEqual(
          0,
          reward.log_absolute_distance([], [], 5))
      self.assertEqual(
          1,
          reward.log_absolute_distance([1], [], 5))
      self.assertEqual(
          1,
          reward.log_absolute_distance([], [1], 5))
  
      self.assertEqual(
          0,
          reward.log_absolute_distance([1, 2, 3], [1, 2, 3], 5))
      self.assertEqual(
          log_diff(1, 5) / 3,  # divided by target length.
          reward.log_absolute_distance([1, 2, 4], [1, 2, 3], 5))
      self.assertEqual(
          log_diff(1, 5) / 3,
          reward.log_absolute_distance([1, 2, 2], [1, 2, 3], 5))
      self.assertEqual(
          log_diff(3, 5) / 3,  # max_dist
          reward.log_absolute_distance([1, 2], [1, 2, 3], 5))
      self.assertEqual(
          log_diff(3, 5) / 3,  # max_dist
          reward.log_absolute_distance([1, 2, 3, 4], [1, 2, 3], 5))
      # Add log differences for each position.
      self.assertEqual(
          (log_diff(2, 5) + log_diff(2, 5) + log_diff(1, 5)) / 3,
          reward.log_absolute_distance([4, 4, 4], [1, 2, 3], 5))
  
    def testAbsoluteDistanceReward(self):
      self.assertEqual(
          1,
          reward.absolute_distance_reward([1, 2, 3], [1, 2, 3], 5))
      self.assertEqual(
          1 - 1 / (5 * 3.),  # 1 - distance / (base * target_len)
          reward.absolute_distance_reward([1, 2, 4], [1, 2, 3], 5))
      self.assertEqual(
          1 - 1 / (5 * 3.),
          reward.absolute_distance_reward([1, 2, 2], [1, 2, 3], 5))
      self.assertTrue(np.isclose(
          1 - 5 / (5 * 3.),
          reward.absolute_distance_reward([1, 2], [1, 2, 3], 5)))
      self.assertTrue(np.isclose(
          1 - 5 / (5 * 3.),
          reward.absolute_distance_reward([1, 2, 3, 4], [1, 2, 3], 5)))
      # Add log differences for each position.
      self.assertEqual(
          1 - (3 + 2 + 1) / (5 * 3.),
          reward.absolute_distance_reward([4, 4, 4], [1, 2, 3], 5))
      self.assertEqual(
          1,
          reward.absolute_distance_reward([], [], 5))
  
    def testAbsoluteModDistanceReward(self):
      self.assertEqual(
          1,
          reward.absolute_mod_distance_reward([1, 2, 3], [1, 2, 3], 5))
      self.assertEqual(
          1 - 1 / (5 * 3.),  # 1 - distance / (base * target_len)
          reward.absolute_mod_distance_reward([1, 2, 4], [1, 2, 3], 5))
      self.assertEqual(
          1 - 1 / (5 * 3.),
          reward.absolute_mod_distance_reward([1, 2, 2], [1, 2, 3], 5))
      self.assertTrue(np.isclose(
          1 - 5 / (5 * 3.),
          reward.absolute_mod_distance_reward([1, 2], [1, 2, 3], 5)))
      self.assertTrue(np.isclose(
          1 - 5 / (5 * 3.),
          reward.absolute_mod_distance_reward([1, 2, 3, 4], [1, 2, 3], 5)))
      # Add log differences for each position.
      self.assertTrue(np.isclose(
          1 - (2 + 2 + 1) / (5 * 3.),
          reward.absolute_mod_distance_reward([4, 4, 4], [1, 2, 3], 5)))
      self.assertTrue(np.isclose(
          1 - (1 + 2 + 2) / (5 * 3.),
          reward.absolute_mod_distance_reward([0, 1, 2], [4, 4, 4], 5)))
      self.assertEqual(
          1,
          reward.absolute_mod_distance_reward([], [], 5))
  
    def testAbsoluteLogDistanceReward(self):
      def log_diff(diff, base):
        return log(diff + 1) / log(base // 2 + 2)
  
      self.assertEqual(
          1,
          reward.absolute_log_distance_reward([1, 2, 3], [1, 2, 3], 5))
      self.assertEqual(
          1 - log_diff(1, 5) / 3,  # divided by target length.
          reward.absolute_log_distance_reward([1, 2, 4], [1, 2, 3], 5))
      self.assertEqual(
          1 - log_diff(1, 5) / 3,
          reward.absolute_log_distance_reward([1, 2, 2], [1, 2, 3], 5))
      self.assertEqual(
          1 - log_diff(3, 5) / 3,  # max_dist
          reward.absolute_log_distance_reward([1, 2], [1, 2, 3], 5))
      self.assertEqual(
          1 - log_diff(3, 5) / 3,  # max_dist
          reward.absolute_log_distance_reward([1, 2, 3, 4], [1, 2, 3], 5))
      # Add log differences for each position.
      self.assertEqual(
          1 - (log_diff(2, 5) + log_diff(2, 5) + log_diff(1, 5)) / 3,
          reward.absolute_log_distance_reward([4, 4, 4], [1, 2, 3], 5))
      self.assertEqual(
          1 - (log_diff(1, 5) + log_diff(2, 5) + log_diff(2, 5)) / 3,
          reward.absolute_log_distance_reward([0, 1, 2], [4, 4, 4], 5))
      self.assertEqual(
          1,
          reward.absolute_log_distance_reward([], [], 5))
  
    def testDeltaRewardManager(self):
      reward_manager = reward.DeltaRewardManager(
          [1, 2, 3, 4], base=5, distance_fn=reward.absolute_distance)
      self.assertEqual(-3, reward_manager([1]))
      self.assertEqual(0, reward_manager([1]))
      self.assertEqual(4 / 5., reward_manager([1, 3]))
      self.assertEqual(-4 / 5, reward_manager([1]))
      self.assertEqual(3, reward_manager([1, 2, 3, 4]))
      self.assertEqual(-1, reward_manager([1, 2, 3]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4, 3]))
      self.assertEqual(-1, reward_manager([1, 2, 3, 4, 3, 2]))
      self.assertEqual(2, reward_manager([1, 2, 3, 4]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4]))
  
    def testFloorRewardMananger(self):
      reward_manager = reward.FloorRewardManager(
          [1, 2, 3, 4], base=5, distance_fn=reward.absolute_distance)
      self.assertEqual(1, reward_manager([1]))
      self.assertEqual(0, reward_manager([1]))
      self.assertEqual(4 / 5., reward_manager([1, 3]))
      self.assertEqual(0, reward_manager([1]))
      self.assertEqual(1 / 5., reward_manager([1, 2]))
      self.assertEqual(0, reward_manager([0, 1]))
      self.assertEqual(0, reward_manager([]))
      self.assertEqual(0, reward_manager([1, 2]))
      self.assertEqual(2, reward_manager([1, 2, 3, 4]))
      self.assertEqual(0, reward_manager([1, 2, 3]))
      self.assertEqual(-1, reward_manager([1, 2, 3, 4, 3]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4, 3, 2]))
      self.assertEqual(1, reward_manager([1, 2, 3, 4]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4]))
  
      reward_manager = reward.FloorRewardManager(
          [1, 2, 3, 4], base=5, distance_fn=reward.absolute_distance)
      self.assertEqual(1, reward_manager([1]))
      self.assertEqual(-1, reward_manager([1, 0, 0, 0, 0, 0]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4, 0, 0]))
      self.assertEqual(0, reward_manager([1, 2, 3, 4, 0]))
      self.assertEqual(1, reward_manager([]))
      self.assertEqual(0, reward_manager([]))
      self.assertEqual(0, reward_manager([1]))
      self.assertEqual(1, reward_manager([1, 2]))
      self.assertEqual(-1, reward_manager([1, 2, 3, 4, 0, 0]))
      self.assertEqual(0, reward_manager([1, 1, 1, 1, 1]))
      self.assertEqual(1 + 2, reward_manager([1, 2, 3, 4]))
  
  
  if __name__ == '__main__':
    tf.test.main()
  