import numpy as np
import abc
from tqdm import tqdm
from .utils import inv_sherman_morrison


class UCB(abc.ABC):
    """Base class for UBC methods.
    """

    def __init__(self,
                 bandit,
                 reg_factor=1.0,
                 confidence_scaling_factor=-1.0,
                 delta=0.1,
                 train_every=1,
                 throttle=int(1e2),
                 per=True
                 ):
        # bandit object, contains features and generated rewards
        self.bandit = bandit
        # L2 regularization strength
        self.reg_factor = reg_factor
        # Confidence bound with probability 1-delta
        self.delta = delta
        # multiplier for the confidence bound (default is bandit reward noise std dev)
        if confidence_scaling_factor == -1.0:
            confidence_scaling_factor = bandit.noise_std
        self.confidence_scaling_factor = confidence_scaling_factor

        # train approximator only every few rounds
        self.train_every = train_every

        # throttle tqdm updates
        self.throttle = throttle

        self.reset()
        self.action = None
        self.postfix = {
            'total regret': 0.0,
            '% optimal arm': 0.0,
        }
        if per:
            self.pbar = tqdm(total=self.bandit.T, postfix=self.postfix)

    def reset_upper_confidence_bounds(self):
        """Initialize upper confidence bounds and related quantities.
        """
        self.exploration_bonus = np.empty((self.bandit.T, self.bandit.n_arms))
        self.mu_hat = np.empty((self.bandit.T, self.bandit.n_arms))
        self.upper_confidence_bounds = np.ones((self.bandit.T, self.bandit.n_arms))

    def reset_regrets(self):
        """Initialize regrets.
        """
        self.regrets = np.empty(self.bandit.T)

    def reset_actions(self):
        """Initialize cache of actions.
        """
        self.actions = np.empty(self.bandit.T).astype('int')

    def reset_A_inv(self):
        """Initialize n_arms square matrices representing the inverses
        of exploration bonus matrices.
        """
        self.A_inv = np.array(
            [
                np.eye(self.approximator_dim) / self.reg_factor for _ in self.bandit.arms
            ]
        )

    def reset_grad_approx(self):
        """Initialize the gradient of the approximator w.r.t its parameters.
        """
        self.grad_approx = np.zeros((self.bandit.n_arms, self.approximator_dim))

    def sample_action(self, k):
        """Return the action to play based on current estimates
        """
        return np.argsort(self.upper_confidence_bounds[self.iteration])[-k:]

    #       return np.argmax(self.upper_confidence_bounds[self.iteration]).astype('int')

    # def sample_action_new(self, num):
    #
    #     return np.

    @abc.abstractmethod
    def reset(self):
        """Initialize variables of interest.
        To be defined in children classes.
        """
        pass

    @property
    @abc.abstractmethod
    def approximator_dim(self):
        """Number of parameters used in the approximator.
        """
        pass

    @property
    @abc.abstractmethod
    def confidence_multiplier(self):
        """Multiplier for the confidence exploration bonus.
        To be defined in children classes.
        """
        pass

    @abc.abstractmethod
    def update_output_gradient(self):
        """Compute output gradient of the approximator w.r.t its parameters.
        """
        pass

    @abc.abstractmethod
    def train(self):
        """Update approximator.
        To be defined in children classes.
        """
        pass

    @abc.abstractmethod
    def predict(self):
        """Predict rewards based on an approximator.
        To be defined in children classes.
        """
        pass

    def update_confidence_bounds(self):
        """Update confidence bounds and related quantities for all arms.
        """
        self.update_output_gradient()

        # UCB exploration bonus
        self.exploration_bonus[self.iteration] = np.array(
            [
                self.confidence_multiplier * np.sqrt(
                    np.dot(self.grad_approx[a], np.dot(self.A_inv[a], self.grad_approx[a].T))) for a in self.bandit.arms
            ]
        )

        # update reward prediction mu_hat
        self.predict()

        # estimated combined bound for reward
        self.upper_confidence_bounds[self.iteration] = self.mu_hat[self.iteration] + self.exploration_bonus[
            self.iteration]

    def update_A_inv(self):
        self.A_inv[self.action] = inv_sherman_morrison(
            self.grad_approx[self.action],
            self.A_inv[self.action]
        )

    def run(self):
        """Run an episode of bandit.
        """
        postfix = {
            'total regret': 0.0,
            '% optimal arm': 0.0,
        }
        with tqdm(total=self.bandit.T, postfix=postfix) as pbar:
            for t in range(self.bandit.T):
                # update confidence of all arms based on observed features at time t
                self.update_confidence_bounds()
                # pick action with the highest boosted estimated reward
                self.action = self.sample_action()
                self.actions[t] = self.action
                # update approximator
                if t % self.train_every == 0:
                    self.train()
                # update exploration indicator A_inv
                self.update_A_inv()
                # compute regret
                self.regrets[t] = self.bandit.best_rewards_oracle[t] - self.bandit.rewards[t, self.action]
                # increment counter
                self.iteration += 1

                # log
                postfix['total regret'] += self.regrets[t]
                n_optimal_arm = np.sum(
                    self.actions[:self.iteration] == self.bandit.best_actions_oracle[:self.iteration]
                )
                postfix['% optimal arm'] = '{:.2%}'.format(n_optimal_arm / self.iteration)

                if t % self.throttle == 0:
                    pbar.set_postfix(postfix)
                    pbar.update(self.throttle)

    def per_run_1(self, k):
        # update confidence of all arms based on observed features at time t
        self.update_confidence_bounds()
        # pick action with the highest boosted estimated reward
        self.action = self.sample_action(k)
        # self.actions[t] = self.action
        return self.action

    def per_run_2(self, t, reward, action=None):
        if action is not None:
            self.actions[t] = action
            self.action = action
        self.bandit.set_reward(t, self.actions[t], reward)
        # update approximator
        if t % self.train_every == 0:
            self.train()
        # update exploration indicator A_inv
        self.update_A_inv()
        # compute regret
        self.regrets[t] = self.bandit.best_rewards_oracle[t] - self.bandit.rewards[t, self.actions[t]]
        # increment counter
        self.iteration += 1

        # log
        self.postfix['total regret'] += self.regrets[t]
        n_optimal_arm = np.sum(
            self.actions[:self.iteration] == self.bandit.best_actions_oracle[:self.iteration]
        )
        self.postfix['% optimal arm'] = '{:.2%}'.format(n_optimal_arm / self.iteration)

        if t % self.throttle == 0:
            self.pbar.set_postfix(self.postfix)
            self.pbar.update(self.throttle)
