# -*- coding: utf-8 -*-
"""
Created on Fri May  6 13:22:05 2022

@author: turinici
"""

import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

k,T,M,epsilon=10,1000,2000,0.1

rewards= np.zeros((T,M))
qstar_vec=np.zeros((k,M))

for  m in tqdm(range(M)):

    qstar = np.random.randn(k)
    qstar_vec[:,m]=qstar
    count_vector=np.zeros(k)
    Q=np.zeros(k)

    for t in range(T):

        x = np.random.rand()
        if(x<epsilon):
            A = np.random.choice(k)
        else:
            A = np.argmax(Q)
        count_vector[A] +=1
        rewards[t,m]=qstar[A] +  np.random.randn() 
        Q[A] = Q[A] + (rewards[t,m]-Q[A])/count_vector[A]

mean_reward=np.mean(rewards,axis=1)

mean_maxq= np.mean(np.max(qstar_vec,axis=0))
plt.plot(range(T),mean_reward,"-r",[0,T],[mean_maxq,mean_maxq],"-b" )#plot mean of rewards
plt.legend([r'$\epsilon$=' + str(epsilon), r'mean max $q_*$'])