import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
np.random.seed(1)

#%%
BANDIT_PROBABILITIES = [0.2, 0.55, 0.75]
NUM_TRIALS = 10_000

class Bandit:
    def __init__(self, true_mean):
        self.true_mean = true_mean
        # parameters for mu - prior is N(0,1)
        self.m = 0
        self.lambda_ = 1
        self.tau = 1
        self.N = 0

    def pull(self):
        return np.random.randn() / np.sqrt(self.tau) + self.true_mean

    def sample(self):
        return np.random.randn() / np.sqrt(self.lambda_) + self.m

    def update(self, x):
        self.m = (self.tau * x + self.lambda_ * self.m) / (self.tau + self.lambda_)
        self.lambda_ += self.tau
        self.N += 1

def plot(bandits, trial):
    plt.figure()
    x = np.linspace(-3, 6, 200)
    for b in bandits:
        y = norm.pdf(x, b.m, np.sqrt(1. / b.lambda_))
        plt.plot(x, y, label=f"real mean: {b.true_mean:.4f}, num plays: {b.N}")
    plt.title(f"Bandit distributions after {trial} trials")
    plt.legend()
    plt.show()

def experiment():
    bandits = [Bandit(p) for p in BANDIT_PROBABILITIES]
    optimal_j = np.argmax(BANDIT_PROBABILITIES)
    sample_points = [5, 10, 20, 50, 100, 200, 500, 1000, 1500, 1999]
    num_times_explored = 0
    num_times_exploited = 0
    num_optimal = 0
    rewards = np.zeros(NUM_TRIALS)
    for i in range(NUM_TRIALS):
        j = np.argmax([b.sample() for b in bandits])

        # plot the posteriors
        if i in sample_points:
            plot(bandits, i)

        if j==optimal_j:
            num_optimal += 1

        x = bandits[j].pull()
        bandits[j].update(x)
        rewards[i] = x

    cumulative_average = np.cumsum(rewards) / (np.arange(NUM_TRIALS) + 1)
    # plot moving average ctr
    plt.plot(cumulative_average)
    for m in BANDIT_PROBABILITIES:
        plt.plot(np.ones(NUM_TRIALS) * m)
    plt.show()

    print("Number times explored: ",num_times_explored)
    print("Number times exploited : ", num_times_exploited )
    print("Total reward earned: ", rewards.sum())
    print("Win rate: ", rewards.sum()/NUM_TRIALS)
    print("Selected optimal bandit: ", num_optimal)

    return cumulative_average


if __name__ == "__main__":
    experiment()
