# -*- coding: utf-8 -*-
"""
Created on Sat May 16 10:09:41 2020

@author: batho
"""
from __future__ import print_function, division
from builtins import range
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import beta

NUM_TRIALS=5000
BANDIT_PROBABILITIES=[0.2,0.5,0.75]

class Bandit:
    def __init__(self,p):
        self.p=p
        self.a=1
        self.b=1
        self.N =0.
    
    def pull(self):
        return np.random.random()<self.p
    def sample(self):
        return np.random.beta(self.a,self.b)
    def update(self,x):
        self.a+=x
        self.b+=1-x
        self.N+=1
def plot(bandits,trial):
    x=np.linspace(0,1,200)
    for b in bandits:
        y=beta.pdf(x,b.a,b.b)
        plt.plot(x,y,label=f"real p: {b.p:.4f}, win rate={b.a-1}/{b.N}")
    plt.title(f"Bandit distibutions after {trial} trials")
    plt.legend()
    plt.show()
    
def experiment():
    bandits=[Bandit(p) for p in BANDIT_PROBABILITIES]
    sample_points=[5,10,20,50,100,200,500,1000,1500,1999]
    rewards=np.zeros(NUM_TRIALS)
    for i in range(NUM_TRIALS):
        j=np.argmax([b.sample() for b in bandits])
        if i in sample_points:
            plot(bandits,i)
        x=bandits[j].pull()
        rewards[i]=x
        bandits[j].update(x)
    print("total reward earned:", rewards.sum())

if __name__ == "__main__":
  experiment()
        
    