import numpy as np
import matplotlib.pyplot as plt

from epsilon_greedy import Bandit
from optimistic_initial_values import run_experiment_oiv
from ucb import run_experiment_ucb

class BayesianBandit:
    def __init__(self, m):
        self.m = m
        # parameters for mu - prior is N(0, 1)
        self.m0 = 0
        self.lambda0 = 1
        self.sum_x = 0 # for convenience
        self.tau = 1
        
    def pull(self):
        return np.random.randn() + self.m

    def sample(self):
        return np.random.randn() / np.sqrt(self.lambda0) + self.m0

    def update(self, x):
        # assume tau is 1
        self.lambda0 += 1
        self.sum_x += x
        self.m0 = self.tau * self.sum_x / self.lambda0
        
def run_experiment_decaying_epsilon(m1, m2, m3, N):
    bandits = [Bandit(m1), Bandit(m2), Bandit(m3)]
               
    data = np.empty(N)
    
    for i in range(N):
        # epsilon greedy
        p = np.random.random()
        if p < 1.0 / (i + 1):
            j = np.random.choice(3)
        else:
            j = np.argmax([b.mean for b in bandits])
        x = bandits[j].pull()
        bandits[j].update(x)
        
        # for the plot
        data[i] = x
    cumulative_average = np.cumsum(data) / (np.arange(N) + 1)
    
    # plot moving average ctr
    plt.plot(cumulative_average)
    plt.plot(np.ones(N) * m1)
    plt.plot(np.ones(N) * m2)
    plt.plot(np.ones(N) * m3)
    plt.xscale('log')
    plt.show()
    
    for b in bandits:
        print(b.mean)
        
    return cumulative_average
    
def run_experiment_bayes(m1, m2, m3, N):
    bandits = [BayesianBandit(m1), 
               BayesianBandit(m2),
               BayesianBandit(m3)]

    data = np.empty(N)
    
    for i in range(N):
        # bayesian sampling
        j = np.argmax([b.sample() for b in bandits])
        x = bandits[j].pull()
        bandits[j].update(x)
        
        # for the plot
        data[i] = x
    cumulative_average = np.cumsum(data) / (np.arange(N) + 1)

    # plot moving average ctr
    plt.plot(cumulative_average)
    plt.plot(np.ones(N) * m1)
    plt.plot(np.ones(N) * m2)
    plt.plot(np.ones(N) * m3)
    plt.xscale('log')
    plt.show()
        
    return cumulative_average
        
if __name__ == '__main__':
    eps = run_experiment_decaying_epsilon(1.0, 2.0, 3.0, 100000)
    oiv = run_experiment_oiv(1.0, 2.0, 3.0, 100000)
    ucb = run_experiment_ucb(1.0, 2.0, 3.0, 100000)
    bayes = run_experiment_bayes(1.0, 2.0, 3.0, 100000)
    
    # log scale plot
    plt.plot(eps, label = 'decaying epsilon greedy')
    plt.plot(oiv, label = 'optimistic initial values')
    plt.plot(ucb, label = 'ucb1')
    plt.plot(bayes, label = 'bayesian')
    plt.legend(loc = 'lower right')
    plt.xscale('log')
    plt.show()
    
    # linear plot
    plt.plot(eps, label = 'decaying epsilon greedy')
    plt.plot(oiv, label = 'optimistic initial values')
    plt.plot(ucb, label = 'ucb1')
    plt.plot(bayes, label = 'bayesian')
    plt.legend(loc = 'lower right')
    plt.show()
