import math
from abc import abstractmethod

# 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 UCB_Base(ExploreExploitBase):
    """
    Follows optimism in the face of uncertainity, by estimating a maximum upper
    confidence bound for each that, with high probability, is compatible
    with the given observations.
    
    Has logarithmic expected regret, but does not reach the best problem-dependent
    best lower bound.
    
    This base class can be extended to implement the wanted variant of UCB,
    and includes all of the base calculations that can be necessary.
    
    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 _on_reset(self):
        """
        Hook for ExploreExploitBase child to handle resetting the class.
        
        Will be called in constructor.
        """
        self._T = 0
        
        self._t = [0 for _ in self._actions]
        # Reset with large number to ensure everything is selected once.
        self._inverse_sqrt_t = [100 for _ in self._actions]
        
        self._previous_ucb = [0 for _ in self._actions]
        
        # Used to calculate variance, see _UpdateAndCalculateVariance
        self._sum_of_squared = [0 for _ in self._actions]
        self._sum = [0 for _ in self._actions]
        
        self._additional_reset()
    

    def _additional_reset(self):
        """
        Hook for UCB child-class to do any additional wanted resets.
        
        Will be called in constructor
        """
        pass
    
    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.
        """
        # Is greedy on best action, since the estimate includes the UCB
        return self._get_best_action()
        
    
    def _on_update_from_action(self, action, r, estimates):
        """
        Hook for child to update any extra state/estimates as required.
        
        Parameters
        --------
        action : action
            Action chosen for this step.
        
        r : numeric
            Reward from choosing the current action.
        
        estimates : array-like, shape [n_actions,]
            The estimate for each action. Can be updated to change the overall
            estimates.
        """
        self._T += 1
        
        # Update the values for current action.
        action_index = self._action_to_index[action]
        self._t[action_index] += 1
        self._inverse_sqrt_t[action_index] = 1 / math.sqrt(self._t[action_index])
        
        self._AdditionalUpdates(action_index, r)
        
        # The estimate for it was regenerated, so won't include the UCB calculation.
        self._previous_ucb[action_index] = 0
        
        log_T = math.log(self._T)
        sqrt_log_T = math.sqrt(log_T)
        
        for i in range(len(self._actions)):
            current_ucb = self._CalculateUCB(i, log_T, sqrt_log_T)
            estimates[i] = estimates[i] - self._previous_ucb[i] + current_ucb
            self._previous_ucb[i] = current_ucb
    
    def _AdditionalUpdates(self, action_index, r):
        """
        Hook for UCB specialization to update any additional state after t
        for action has been updated.
        
        Parameters
        --------
        action_index : integer
            Index of action chosen for this step.
        
        r : numeric
            Reward from choosing the current action.
        """
        pass
    
    def _UpdateAndCalculateVariance(self, action_index, r):
        """
        Can be called to update the variance for the action, and then will return
        the current variance for the action.
        
        Calculates variance from sum(xi**2) / t - mean**2, so speed is unchanged the
        more iterations are done.
        
        Warning
        ---------
        This function must be called in _AdditionalUpdates if you want to calculate
        variance - it is not normally called by UCB_Base
        
        Parameters
        ---------
        action : integer
            Index of action to calculate variance.
        
        r : numeric
            Current reward for the action in this step.
        """
        # This calculates variance = sum(xi**2) / t - mean**2
        self._sum_of_squared[action_index] += r * r
        self._sum[action_index] += r
        
        mean = self._sum[action_index] / self._t[action_index]
        
        mean_squared = mean ** 2
        
        # Rounding errors can cause the mean_squared to be > self._sum_of_squared[action_index] / self._t[action_index]
        return max(0, self._sum_of_squared[action_index] / self._t[action_index] - mean_squared)
        
    @abstractmethod
    def _CalculateUCB(self, action_index, log_T, sqrt_log_T):
        """
        Calculate the UCB for the current action.
        
        Can use _t, _inverse_sqrt_t or any other specialized calculations.
        
        Parameters
        --------
        action_index : integer
            Index of action to be recalculated.
        
        log_T : numeric
            log_T = math.log(self._T), to save on calculations between different
            actions.
        
        sqrt_log_T : numeric
            sqrt_log_T = math.sqrt(log_T), to save on calculations between
            different actions.
        """
        pass
    
class UCB(UCB_Base):
    """
    Follows optimism in the face of uncertainity, by estimating a maximum upper
    confidence bound for each that, with high probability, is compatible
    with the given observations.
    
    Has logarithmic expected regret, but does not reach the best problem-dependent
    best lower bound.
    
    Formula is: B = mean + c * sqrt(logT / t), where c is a constant provided,
    t is the number of pulls for the current arm, and T is total number of pulls.
    
    The original paper for UCB recommended c = sqrt(2), but other constants like
    c = sqrt(3/2) have been suggested in the literature.
    
    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.
        
    c
        Factor for the exploration factor in B estimate. As grows larger, will
        explore actions more.
        
    References
    --------
        Remi Munos. From Bandits to Monte-Carlo Tree Search: The Optimistic Principle Applied to
        Optimization and Planning. 2012.
    """
    def __init__(self, actions, estimator, c=math.sqrt(2)):
        super().__init__(actions, estimator)
        self._c = c
        
        
    def _CalculateUCB(self, action_index, log_T, sqrt_log_T):
        """
        Calculate the UCB for the current action.
        
        Can use _t, _inverse_sqrt_t or any other specialized calculations
        
        Parameters
        --------
        action_index : integer
            Index of action to be recalculated.
        
        log_T : numeric
            log_T = math.log(self._T), to save on calculations between different
            actions.
        
        sqrt_log_T : numeric
            sqrt_log_T = math.sqrt(log_T), to save on calculations between
            different actions.
        """
        return self._c * sqrt_log_T * self._inverse_sqrt_t[action_index]
        
    
    def describe(self):
        """
        Returns a formatted string to describe the implementation + parameters
        for this instance.
        """
        return "UCB {:.3f}".format(self._c)

class UCB_V(UCB_Base):
    """
    Follows optimism in the face of uncertainity, by estimating a maximum upper
    confidence bound for each that, with high probability, is compatible
    with the given observations.
    
    Has logarithmic expected regret, but does not reach the best problem-dependent
    best lower bound. Does achieve a lower regret than standard UCB.
    
    Formula is: B = mean + c * sqrt(Variance * log(N_factor * N) / t) + d * log(N_factor * N) / t,
    where c, d, and N_factor are provided constants that can be optimized, t is
    the number of pulls for the current arm, variance is experimental variance for the current arm,
    and N is total number of pulls.
    
    In the paper describing UCB-V, the log(N_factor * N) was a generic function depending
    on N or t, but log is the standard function used.
    
    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.
        
    c : numeric
        Factor for the variance based exploration factor. The larger it is, the more exploration
        will be done.
        
    d : numeric
        Factor for the second term of the exploration factor. The larger it is, the more exploration
        will be done.
    
     N_factor : numeric
        Value to scale N (total number of steps) by inside the log factor.
        The larger it is, the more exploration will be done. Probably should be >= 1.
    
    References
    --------
        Remi Munos. From Bandits to Monte-Carlo Tree Search: The Optimistic Principle Applied to
        Optimization and Planning. 2012.
        
        J.Y. Audibert, R. Munos, and Cs. Szepesvari. Exploration-exploitation tradeoff using variance ´
        estimates in multi-armed bandits. Theoretical Computer Science, 410(19):1876–1902, 2009.
    """
    def __init__(self, actions, estimator, c=math.sqrt(2), d=1, N_factor=1):
        super().__init__(actions, estimator)
        self._c = c
        self._d = d
        self._N_factor = N_factor

    def _additional_reset(self):
        """
        Hook for UCB child-class to do any additional wanted resets.
        
        Will be called in constructor
        """
        self._variance = [1 for _ in self._actions]
        self._sqrt_variance = [1 for _ in self._actions]
    
    def _AdditionalUpdates(self, action_index, r):
        """
        Hook for UCB specialization to update any additional state after t
        for action has been updated.
        
        Parameters
        --------
        action_index : integer
            Index of action chosen for this step.
        
        r : numeric
            Reward from choosing the current action.
        """
        self._variance[action_index] = self._UpdateAndCalculateVariance(action_index, r)
        
        self._sqrt_variance[action_index] = math.sqrt(self._variance[action_index])
        
        
    def _CalculateUCB(self, action_index, log_T, sqrt_log_T):
        """
        Calculate the UCB for the current action.
        
        Can use _t, _inverse_sqrt_t or any other specialized calculations
        
        Parameters
        --------
        action_index : integer
            Index of action to be recalculated.
        
        log_T : numeric
            log_T = math.log(self._T), to save on calculations between different
            actions.
        
        sqrt_log_T : numeric
            sqrt_log_T = math.sqrt(log_T), to save on calculations between
            different actions.
        """
        current_ucb = self._c * self._sqrt_variance[action_index] * self._inverse_sqrt_t[action_index] * sqrt_log_T
        # d * log(N_factor * N) / t
        second_ucb_factor = self._d * log_T
        if self._t[action_index] > 0:
            second_ucb_factor /= self._t[action_index]
        
        return current_ucb + second_ucb_factor
        
    def describe(self):
        """
        Returns a formatted string to describe the implementation + parameters
        for this instance.
        """
        return "UCB-V,  {0:.3f} * sqrt(Variance log({2} * N) / t) + {1} * log({2} * N) / t".format(self._c, self._d, self._N_factor)
