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

# SECTION 4.4: RUN EXPERIMENTS for Case 2

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

#Read files
Y = np.genfromtxt(r"/data/Case2/poly2_Y_all(mon).csv",delimiter=',', skip_header=1)
W = np.genfromtxt(r"/data/Case2/poly2_W.csv",delimiter=',', skip_header=1)

# 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(Y, N_TRIAL, N_ARMS, M)

# Initialize q and C
# vec_q: q (Kp x 1)
np.random.seed(3) #3 #59
vec_q = np.array([np.random.rand() for _ in range(K * N_FEATURE)])
# vec_C: C (NK x 1)
np.random.seed(42)
vec_C = np.array([np.random.rand() for _ in range(N_ARMS * N_FEATURE)])


# Construct laplacian matrix
lda = 1.25
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
lda3 = 5
T = 0.5
test_adj3 = ut.constructAdjMatrix(W, N_ARMS, T)
test_F3 = ut.constructLaplacianMatrix(test_adj3, N_ARMS, lda)
FInv_Init3 = sqrtm(np.linalg.inv(np.kron(test_F3, np.identity(n=N_FEATURE))))

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


# # Run GOB.Lin
print('GOB.Lin--','lda:', lda3)
alpha_to_test = [1.25]
results_dict_gob = {alpha: goblin.GOBLin(alpha = alpha, eta = 0.3, X=X, Y=Y.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 = [3]
results_dict_clucbwo = {alpha: clucb.CLUCB_conv2(eta_1 = 0.3, eta_2 = 0.3, alpha_q = 1.5, alpha_c = alpha, X=X, Y=Y.T, init_q=vec_q, init_c=vec_C,m=M, K = K, FInv=FInv_Init2, oracle=oracle, true_choice=true_choice, bb = 2000,max_iters = 5)\
                for alpha in alpha_to_test}

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

#plot regrets
ut.reg_plot_all_c(results_dict_linucb[3.5]['cum_regret'], results_dict_gob[1.25]['cum_regret'], results_dict_clucbwo[3]['cum_regret'], results_dict_clucb[3]['cum_regret'], N_TRIAL, 2, M)