import numpy as np
import matplotlib.pyplot as plt


NUM_TRIALS = 100000
EPS = 0.05
BANDIT_PROBABILITIES = [0.2, 0.5, 0.75]


class Bandit:
  def __init__(self, p):
    # p: the win rate
    self.p = p
    self.p_estimate = 0.
    self.N = 0. # num samples collected so far

  def pull(self):
    # draw a 1 with probability p
    return np.random.random() < self.p

  def update(self, x):
    self.N += 1.
    self.p_estimate = ((self.N - 1)*self.p_estimate + x) / self.N


def ucb(mean, n, nj):
  return mean + np.sqrt(2*np.log(n) / nj)


def run_experiment():
  bandits = [Bandit(p) for p in BANDIT_PROBABILITIES]
  rewards = np.empty(NUM_TRIALS)
  total_plays = 0

  # initialization: play each bandit once
  for j in range(len(bandits)):
    x = bandits[j].pull()
    total_plays += 1
    bandits[j].update(x)
  
  for i in range(NUM_TRIALS):
    j = np.argmax([ucb(b.p_estimate, total_plays, b.N) for b in bandits])
    x = bandits[j].pull()
    total_plays += 1
    bandits[j].update(x)

    # for the plot
    rewards[i] = x
  cumulative_average = np.cumsum(rewards) / (np.arange(NUM_TRIALS) + 1)

  # plot moving average ctr
  plt.plot(cumulative_average)
  plt.plot(np.ones(NUM_TRIALS)*np.max(BANDIT_PROBABILITIES))
  plt.xscale('log')
  plt.show()

  # plot moving average ctr linear
  plt.plot(cumulative_average)
  plt.plot(np.ones(NUM_TRIALS)*np.max(BANDIT_PROBABILITIES))
  plt.show()

  for b in bandits:
    print(b.p_estimate)

  print("total reward earned:", rewards.sum())
  print("overall win rate:", rewards.sum() / NUM_TRIALS)
  print("num times selected each bandit:", [b.N for b in bandits])

  return cumulative_average


def run_exp () :

    bandits = [Bandit(p) for p in BANDIT_PROBABILITIES]

    data = np.empty(NUM_TRIALS) 
    for i in range(NUM_TRIALS): 
        # epsilon greedy 
        p = np.random.random() 
        if p < EPS: 
            j = np.random.choice(3) 
        else: 
            j = np.argmax([a.p_estimate for a in bandits]) 
        x = bandits[j].pull() 
        bandits[j].update(x) 
        
        # for the plot 
        data[i] = x 
    cumulative_average = np.cumsum(data) / (np.arange(NUM_TRIALS) + 1) 
    return cumulative_average


if __name__ == '__main__':
  ucb = run_experiment()
  c_1 = run_exp()


# log scale plot 
plt.plot(c_1, label ='eps = 0.01') 
plt.plot(ucb, label ='ucb') 
plt.legend() 
plt.xscale('log') 
plt.show()


  # linear plot 
plt.figure(figsize = (12, 8)) 
plt.plot(c_1, label ='eps = 0.01') 
plt.plot(ucb, label ='ucb') 
plt.legend() 
plt.show()   