import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from comparing_epsilons import Bandit
from optim_init_values import run_experiments as run_experiments_oiv
from ucb1 import run_experiments as run_experiments_ucb
import os

class BayesianBandit:
    def __init__(self, true_mean):
        # consider likelihood as N(0, 1), prior as N(0, 1)
        # in BayesianEstimation, data is fixed and the parameters have the distribution
        self.true_mean = true_mean
        self.mu0 = 0
        self.invsigma0 = 1
        self.sum = 0 # for convenience
        
    # pull the arm
    def pull(self):
        return np.random.randn() + self.true_mean
        
    def sample(self):
        return (np.random.randn() / self.invsigma0) + self.mu0
    
    # update the running mean
    def update(self, Xn):
        self.sum += Xn # running sum
        self.invsigma0 += 1 # invsigma0 + invsigma(=1) * N
        self.mu0 = self.sum / self.invsigma0 # (invsigma * runningsum + invsigma0 * mu0(=0)) / (invsigma0 + invsigma(=1) * N)  

#
# m1, m2, m3 = rewards of 3 bandits
# N = number of events/tries
# eps = for explore/exploit 
#
def run_experiments_bayesian(m1, m2, m3, N):
    # create 3 bandits
    bandits = [BayesianBandit(m1), BayesianBandit(m2), BayesianBandit(m3)]
    data = np.empty(N)
    
    for i in range(N):
        # sample from a bandit and update the mean
        j = np.argmax([b.sample() for b in bandits])
        X = bandits[j].pull()
        bandits[j].update(X)
        
        data[i] = X
        
    cumulative_avg = np.cumsum(data) / (np.arange(N) + 1)
    return cumulative_avg

#
# m1, m2, m3 = rewards of 3 bandits
# N = number of events/tries
# eps = for explore/exploit 
#
def run_experiments_epsilon(m1, m2, m3, N):
    # create 3 bandits
    bandits = [Bandit(m1), Bandit(m2), Bandit(m3)]
    data = np.empty(N)
    
    for i in range(N):
        # epsilon explore-exploit dilemma
        p = np.random.random()
        if(p < 1.0/(1+i)): # linear decrease of epsilon w.r.t time
            j = np.random.choice(3)
        else:
            j = np.argmax([b.running_mean for b in bandits])
            
        X = bandits[j].pull()
        bandits[j].update(X)
        
        data[i] = X
        
    cumulative_avg = np.cumsum(data) / (np.arange(N) + 1)
    return cumulative_avg
    
if(__name__ == '__main__'):
    
    explore_exploit_epsilon = run_experiments_epsilon(1., 2., 3., 100000)
    explore_exploit_oiv = run_experiments_oiv(1., 2., 3., 100000, 0, 10)
    explore_exploit_ucb = run_experiments_ucb(1., 2., 3., 100000, 0, isUCB=True)
    explore_exploit_bayesian = run_experiments_bayesian(1., 2., 3., 100000)
    
    plt.plot(explore_exploit_epsilon, label='decaying epsilon')
    plt.plot(explore_exploit_oiv, label='optim init values')
    plt.plot(explore_exploit_ucb, label='ucb')
    plt.plot(explore_exploit_bayesian, label='bayesian')
    plt.legend()
    plt.xscale('log')
    plt.savefig('./pics/'+ os.path.splitext(__file__)[0] + '_logplot.jpg')
    plt.clf()
    
    plt.plot(explore_exploit_epsilon, label='decaying epsilon')
    plt.plot(explore_exploit_oiv, label='optim init values')
    plt.plot(explore_exploit_ucb, label='ucb')
    plt.plot(explore_exploit_bayesian, label='bayesian')
    plt.legend()
    plt.savefig('./pics/'+ os.path.splitext(__file__)[0] + '_plot.jpg')  
