import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.linalg import block_diag
from numpy.linalg import multi_dot
from scipy.linalg import sqrtm
from scipy.sparse import csgraph
import os
import itertools
import util as ut
import LinUCB as linucb
import CLUCB as clucb
import GOBLin as goblin

# APPENDIX D: RUN EXPERIMENTS USING DIFFERENT W 

# For the full experiment, please change N_TRIAL to N_TRIAL = 30000
N_TRIAL = 5
N_ARMS = 100
K = 3
N_FEATURE = 3
M = 37

#Read files
Beta = np.genfromtxt(r"/data/new_W/Beta_hk.csv",delimiter=',')
q = np.genfromtxt(r"/data/new_W/q_init4.csv",delimiter=',')
c = np.genfromtxt(r"/data/new_W/c_init4.csv",delimiter=',')
Y = np.genfromtxt(r"/data/new_W/Y_hk.csv",delimiter=',')
W = np.genfromtxt(r"/data/new_W/A_hk.csv",delimiter=',')

# Create X_i = [1, t, t^2]
# Increment
# small (5000iter: 0.006)/ big(30000iter: 0)
inc = 0
X = ut.construct_X(N_TRIAL, N_ARMS, inc)

# Create true Y
true_choice, oracle = ut.construct_opt(-1 * Y + 30, N_TRIAL, N_ARMS, M)

# Initialize q and C
# vec_q: q (Kp x 1)
# vec_C: C (NK x 1)
vec_q = q[~np.isnan(q)]
vec_C = c


# Construct laplacian matrix
lda = 0.1
T = 0.5
test_adj = ut.constructAdjMatrix(W, N_ARMS, T)
test_F = ut.constructLaplacianMatrix(test_adj, N_ARMS, lda)
FInv_Init = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=N_FEATURE))))

# Construct laplacian matrix lda = 0
lda2 = 0
T = 0.5
test_adj2 = ut.constructAdjMatrix(W, N_ARMS, T)
test_F2 = ut.constructLaplacianMatrix(test_adj2, N_ARMS, lda2)
FInv_Init2 = sqrtm(np.linalg.inv(np.kron(test_F2, np.identity(n=N_FEATURE))))

# Construct laplacian matrix GOB.Lin
lda3 = 1.25
T = 0.5
test_adj3 = ut.constructAdjMatrix(W, N_ARMS, T)
test_F3 = ut.constructLaplacianMatrix(test_adj3, N_ARMS, lda3)
FInv_Init3 = sqrtm(np.linalg.inv(np.kron(test_F3, np.identity(n=N_FEATURE))))

print("N_TRIAL:", N_TRIAL, "N_ARMS:", N_ARMS, "K:", K, "N_FEATURE:", N_FEATURE, "M:", M)

# Run LinUCB
print('LinUCB')
alpha_to_test =[3.25]
results_dict_linucb = {alpha: linucb.linUCB(lda = 0.3, alpha = alpha, X=X, Y=(-1 * Y + 30).T, m=M, oracle=oracle, true_choice=true_choice)\
                for alpha in alpha_to_test}


# # Run GOB.Lin
print('GOB.Lin--','lda:', lda)
alpha_to_test = [1]
results_dict_gob = {alpha: goblin.GOBLin(alpha = alpha, eta = 0.3, X=X, Y=(-1 * Y + 30).T, m=M, FInv=FInv_Init3, oracle=oracle, true_choice=true_choice, bb = 5000)\
                for alpha in alpha_to_test}


# # Run CLUCB w/o
print('CLUCB w/o--', 'lda:', lda2)
alpha_to_test = [1]
results_dict_clucbwo = {alpha: clucb.CLUCB(eta_1 = 0.3, eta_2 = 0.3, alpha_q = 1, alpha_c = alpha, X=X, Y=(-1 * Y + 30).T, init_q=vec_q, init_c=vec_C,m=M, K = K, FInv=FInv_Init2, oracle=oracle, true_choice=true_choice, bb = 2000)\
                for alpha in alpha_to_test}

# Run CLUCB
print('CLUCB--', 'lda:', lda)
alpha_to_test = [0.5]
results_dict_clucb = {alpha: clucb.CLUCB(eta_1 = 0.3, eta_2 = 0.3, alpha_q = 1, alpha_c = alpha, X=X, Y=(-1 * Y + 30).T, init_q=vec_q, init_c=vec_C,m=M, K = K, FInv=FInv_Init, oracle=oracle, true_choice=true_choice, bb = 2000)\
                for alpha in alpha_to_test}

#plot results
ut.reg_plot_all_c(results_dict_linucb[3.25]['cum_regret'], results_dict_gob[1]['cum_regret'], results_dict_clucbwo[1]['cum_regret'], results_dict_clucb[0.5]['cum_regret'], N_TRIAL, 3, M)

