import unittest

import reinforcement_learning.explore_exploit.greedy_variants as greedy_variants
import reinforcement_learning.explore_exploit.ucb as ucb_variants

from reinforcement_learning.explore_exploit.estimator import AverageEstimator


class ExploreExploitTester(unittest.TestCase):
    # Check simple requirements - like the explore exploit algorithms will select
    # each action once at the start, that they will select an action that did far
    # better more often than other moves.
    
    # Will only have 3 actions.
    actions = [0, 1, 2]
    
    # Tuple of (algorithm, should explore)
    algorithms = [(greedy_variants.GreedyExploreExploit(actions, AverageEstimator()), False),
                  (greedy_variants.EpsilonGreedyExploreExploit(0.1, actions, AverageEstimator()), True),
                  # Note: This only explores in the beginning, and won't do any more exploration after that.
                  (greedy_variants.GreedyOptimisticInitExploreExploit(actions, AverageEstimator(), 1, 5), True),
                  (ucb_variants.UCB(actions, AverageEstimator()), True),
                  (ucb_variants.UCB_V(actions, AverageEstimator()), True)]

    def testAllMostlySelectObviousBest(self):
        algorithms = ExploreExploitTester.algorithms
        num_moves = 100
        
        # Note: Don't expect them to be 100% best over the 100 moves.
        # But it should be selected the most
        rewards = [0, 0.5, 1]
        best_action = 2
        for algorithm_explore in algorithms:
            algorithm = algorithm_explore[0]
            explores = algorithm_explore[1]
            algorithm.reset()
            
            count = [0, 0, 0]
            
            for _ in range(num_moves):
                action = algorithm.select_action()
                count[action] += 1
                
                algorithm.update_from_action(action, rewards[action])
            
            # If algorithm will not explore, then we expect it to have almost all
            # but 2 on best. otherwise we expect it to still have most >= 80.
            self.assertGreaterEqual(count[best_action], 80)
            
            if not explores:
                self.assertEqual(count[best_action], 98) 
    
    def testAllResetProperly(self):
        algorithms = ExploreExploitTester.algorithms
        num_moves = 100
        
        # Have the action 2 be best.
        rewards = [0, 0, 1]
        self.run_algorithms(rewards, num_moves, algorithms)
        
        # Now, make action 1 be best by far
        # Then, reset to have the 1st action be best.
        rewards = [0, 1, 0]
        best_action = 1
        
        def count_check_func(count, algorithm_name):
            self.assertGreaterEqual(count[best_action], 80, msg="Algorithm " + algorithm_name)
            
        self.run_algorithms(rewards, num_moves, algorithms, count_check_func=count_check_func)
            
        
    def testWillHandleWorseningAction(self):
        algorithms = ExploreExploitTester.algorithms
        num_moves = 100
        
        # Have the action 2 be best.
        rewards = [0, 0.7, 1]
        self.run_algorithms(rewards, num_moves, algorithms)
        
        # Now, reward for 1 becomes far worse.
        rewards = [0, 0.7, 0]
        best_action = 1
        
        # Note that it will take some time to update, so allow it to take ~45 actions (normally ~55 taken)
        def count_check_func(count, algorithm_name):
            self.assertGreaterEqual(count[best_action], 45, msg="Algorithm " + algorithm_name)
        
        self.run_algorithms(rewards, num_moves, algorithms, reset=False, count_check_func=count_check_func)
        
    
    def testShouldExploreAlgorithmsWillExploreShort(self):
        # Only for ones that have [1] be true!
        # After a few moves, switch to have action 1 be best
        algorithms = ExploreExploitTester.algorithms
        
        exploring_algorithms = [algorithm_explore for algorithm_explore in algorithms if algorithm_explore[1] == True]
        
        num_moves = 20
        rewards = [0, 0.4, 0.5]
            
        # Simple run with initial rewards
        self.run_algorithms(rewards, num_moves, exploring_algorithms)
        
            
        # Now make action 1 be far better.
        rewards = [0, 0.8, 0.5]
        best_action = 1
        # Huge number of actions, since will need to decide that 1 is better than 2
        num_moves = 1000
        
        # Can expect that >= 80% of the moves will be on action 1. Often far higher, but a good minimum.
        def count_check_func(count, algorithm_name):
            self.assertGreaterEqual(count[best_action], num_moves * .8, msg="Algorithm " + algorithm_name)
            
        self.run_algorithms(rewards, num_moves, exploring_algorithms, reset=False, count_check_func=count_check_func)

    def run_algorithms(self, rewards, num_moves, algorithms, reset=True, count_check_func=None):
        """
        Will run each algorithm separately, after resetting it with the given rewards.
        
        If count_check_func is provided, will call it on each algorithm. It will
        be provided the arguments (counts, algorithm.describe()), where counts
        is how many times the algorithm selected each action.
        """
        for algorithm_explore in algorithms:
            algorithm = algorithm_explore[0]
            if reset:
                algorithm.reset()
            
            counts = [0 for _ in range(len(rewards))]
            
            for _ in range(num_moves):
                action = algorithm.select_action()
                algorithm.update_from_action(action, rewards[action])
                
                counts[action] += 1
            
            if count_check_func is not None:
                count_check_func(counts, algorithm.describe())
    