import numpy as np

from abc import ABC, abstractmethod

class ExploreExploitBase(ABC):
    """
    Base class for all of the different options for exploration vs exploitation.
    Will determine if an optimal (exploit) choice should be made, or if a possibly
    sub-optimal choice should be made (explore).
    
    Currently only setup to handle multi-armed bandits, where there is a single state
    and a single 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):
        self._actions = actions
        self._estimator = estimator
        
        # Map from action to its index
        self._action_to_index = dict([(action, i) for i, action in enumerate(actions)])
        
        self.reset()
        
    def reset(self):
        """
        Reset the estimator. Must be called in constructor
        
        Will be running on arms with possibly different distributions.
        """
        self._estimator.reset()
        # Store rewards in array to make it far faster to retrieve.
        self._estimates = np.array([self._estimator.estimate(action)
                             for action in self._actions], dtype=np.float32)
        
        self._on_reset()
        
    def _on_reset(self):
        """
        Hook for child to handle resetting the class.
        
        Will be called in constructor.
        """
        pass
        
    @abstractmethod
    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.
        """
        pass
    
    def _get_best_action(self):
        """
        Can be called by child to get the best action - action with highest
        expected reward according to estimator.
        """
        return self._actions[self._estimates.argmax()]
    
    def update_from_action(self, action_chosen, r):
        """
        Update based on the result of the action selected.
        
        Parameters
        --------
        action_chosen : action
            Action chosen for this step.
        
        r : numeric
            Reward from choosing the current action.
        """
        self._estimator.update(action_chosen, r)
        
        self._estimates[self._action_to_index[action_chosen]] =\
            self._estimator.estimate(action_chosen)
        
        self._on_update_from_action(action_chosen, r, self._estimates)
    
    def _on_update_from_action(self, action, r, estimates):
        """
        Hook for child to update any extra state/estimates as required.
        
        Parameters
        --------
        action_chosen : 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.
        """
        pass
    
    @abstractmethod
    def describe(self):
        """
        Returns a formatted string to describe the implementation + parameters
        for this instance.
        """
        pass