from Learners.Learner import *


class TS_Learner(Learner):
    #This class extends the superclass Learner so in the construction we have to pass
    #the same arguments

    def __init__(self, n_arms, prices):
        #We have 2 parameters for each arms
        super().__init__(n_arms, prices)
        self.beta_parameters = np.ones((n_arms,2))

    def pull_arm(self):
        beta = np.random.beta(self.beta_parameters[:,0], self.beta_parameters[:,1])
        idx = np.argmax(beta*self.prices)
        return idx

    def update(self, pulled_arm, reward):
         self.t+= 1
         self.update_observations(pulled_arm, reward)
         self.beta_parameters[pulled_arm,0] = self.beta_parameters[pulled_arm,0] + reward
         self.beta_parameters[pulled_arm,1] = self.beta_parameters[pulled_arm,1] + 1.0 - reward
