import random

# Add base directory of project to path.
import os
import sys
dir_path = os.path.dirname(os.path.realpath(__file__))
sys.path.append(dir_path + "/../..")

from reinforcement_learning.explore_exploit.base import ExploreExploitBase

class GreedyExploreExploit(ExploreExploitBase):
    """
    Will greedily select the 'best' action (action with highest expected reward).
    
    Has linear expected regret, due to often locking onto a single action - especially
    if the optimal action has a below average first reward.
    
    Parameters
    --------
    actions : array-like, shape [n_actions,]
        All of the valid actions
    
    estimator : estimator.Estimator
        Should be able to estimate + update the expected reward for each of the different actions.
        For bandits, should use MeanEstimator.
    """
    def __init__(self, actions, estimator):
        super().__init__(actions, estimator)
    
    def select_action(self):
        """
        Will return one of the possible actions.
        
        WARNING: update_from_action must be called after, otherwise may only select
        the same action.
        """
        return self._get_best_action()

    def describe(self):
        """
        Returns a formatted string to describe the implementation + parameters
        for this instance.
        """
        return "GreedyExploreExploit"

class EpsilonGreedyExploreExploit(ExploreExploitBase):
    """
    Will greedily select the 'best' action (action with highest expected reward)
    1 - epsilon % of time, and otherwise randomly select an action.   
    
    While will normally have less regret than GreedyExploreExploit but still has
    linear expected regret due to the % chance to select the non-optimal move.
    
    Parameters
    --------
    epsilon : float
        Percent chance to randomly select a move
    
    actions : array-like, shape [n_actions,]
        All of the valid actions
    
    estimator : estimator.Estimator
        Should be able to estimate + update the expected reward for each of the different actions.
        For bandits, should use MeanEstimator.
    """
    def __init__(self, epsilon, actions, estimator):
        super().__init__(actions, estimator)
        self._epsilon = epsilon
    
    def select_action(self):
        """
        Will return one of the possible actions.
        
        WARNING: update_from_action must be called after, otherwise may only select
        the same action.
        """
        if random.random() < self._epsilon:
            return random.choice(self._actions)
        return self._get_best_action()

    def describe(self):
        """
        Returns a formatted string to describe the implementation + parameters
        for this instance.
        """
        return "EpsilonGreedyExploreExploit {}".format(self._epsilon)

class GreedyOptimisticInitExploreExploit(ExploreExploitBase):
    """
    Like GreedyExploreExploit, but starts all of the actions with a weight of
    max_r with count start_count. With good starting parameters, can perform extremely
    well - see reinforcement_learning.multi_armed_bandits.examples.simple_percentage_bandits for an example.
    
    Idea from RL Course by David Silver - Lecture 9: Exploration and Exploitation
    """
    
    def __init__(self, actions, estimator, max_r, start_count):
        self._max_r = max_r
        self._start_count = start_count
        super().__init__(actions, estimator)
    
    def select_action(self):
        """
        Will return one of the possible actions.
        
        WARNING: update_from_action must be called after, otherwise may only select
        the same action.
        """
        return self._get_best_action()
    
    def _on_reset(self):
        self._set_estimates()
        
    def _set_estimates(self):
        for action in self._actions:
            self._estimator.update(action, self._max_r, self._start_count)

    def describe(self):
        """
        Returns a formatted string to describe the implementation + parameters
        for this instance.
        """
        return "GreedyOptimisticInitExploreExploit MaxR: {} Start Count: {}".format(
                self._max_r, self._start_count)
