# Plot regret rate of EI-TLR under Example 3
# Reproduce the right panel of Figure 3 in the appendix

import os
import sys
import numpy as np
import matplotlib.pyplot as plt

sys.path.append('/Your Path/Data_Code/Methods')
from Pre_Functions import lin_reg
from math import log
from math import ceil


# read data, i.e., the exploration, exploitation and overall regrets for multiple horizon lengths
l = 20
reg_exp3 = np.zeros((3,l))
os.chdir('/Your Path/Data_Code/Simulation/Results/Ex3_RegretRate')
for i in range(l):
    s = "EI-TLR_Ex3_reg_exp2_mean_" + str(i + 1) + ".npy"
    reg_exp2 = np.load(s)
    reg_exp3[0,i] = reg_exp2[0]
    reg_exp3[1,i] = reg_exp2[1]
    reg_exp3[2,i] = np.sum(reg_exp2)


# construct log regret and log horizon lengths
regret = reg_exp3[2,]
range_T = np.zeros(20)
for i in range(l):
    range_T[i] = ceil(200000*10**(i/(l-1)))

log_regret = np.zeros(len(regret))
log_range_T = np.zeros(len(range_T))
for i in range(len(regret)):
    log_range_T[i] = log(range_T[i])
    log_regret[i] = log(regret[i])


# fit simple linear regression
alpha,beta = lin_reg(log_range_T,log_regret)
print(alpha,beta)


# plot
plt.figure()
log_l = plt.scatter(log_range_T,log_regret,s=5,label='_nolegend_')
plt.axis([min(log_range_T),max(log_range_T),min(log_regret),max(log_regret)])
plt.plot([0,100],[alpha,alpha+100*beta])
plt.legend(labels=[str(round(alpha,4))+'+'+str(round(beta,4))+'x'],loc='best',fontsize = 12)
plt.xlabel("Log Horizon Length",fontsize = 20)
plt.ylabel("Log Regret",fontsize = 20)
plt.title("Example 3",fontsize = 24)
plt.show()