from agent import Agent 
from util import *
import numpy as np

class ExploreThenExploit(Agent):
    best_arm = 0
    current_arm = 0
    explore_counter = 0
    #
    def __init__(self,time_horizon,number_of_arms,number_of_exploration_per_arm):
        self.number_of_exploration_per_arm = number_of_exploration_per_arm
        self.time_horizon = time_horizon
        self.number_of_arms = number_of_arms
        self.reset()
        self.name = "explore-then-exploit"        
    #
    def reset(self):
        self.mu_estimators = [0]*self.number_of_arms
        self.arms_visited = [0]*self.number_of_arms
        self.best_arm = 0
        self.current_arm = 0
        self.explore_counter = 0
        print('explore-then-exploit agent started to decide.')
    #
    def decide(self,time,time_horizon,number_of_arms):
        #explore 
        if (time < number_of_arms * self.number_of_exploration_per_arm):
            self.best_arm = self.current_arm 
            self.explore_counter += 1
            if self.explore_counter == self.number_of_exploration_per_arm:
                self.current_arm += 1
                self.explore_counter = 0
        else:
            self.best_arm = np.argmax(self.mu_estimators)
        self.arms_visited[self.best_arm] += 1
        return self.best_arm
         
    #
    def observe(self,time,arm,reward):
        if (time < self.number_of_arms * self.number_of_exploration_per_arm):
            self.mu_estimators[arm] = incrementalAvg(self.arms_visited[arm],self.mu_estimators[arm],reward)