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 plot_all as ptall
import LinUCB as linucb
import CLUCB as clucb
import GOBLin as goblin

# SECTION 4.2: RUN CLUCB using different maximum iterations G

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

#Read files
Beta = np.genfromtxt(r"/data/Case1/Beta(noise0.1).csv",delimiter=',')
q = np.genfromtxt(r"/data/Case1/q_init4.csv",delimiter=',')
c = np.genfromtxt(r"/data/Case1/c_init4.csv",delimiter=',')
Y = np.genfromtxt(r"/data/Case1/Y(noise0.1).csv",delimiter=',')
W = np.genfromtxt(r"/data/Case1/W(noise0.1).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 the optimal strategy
# case1: -1 * Y + 30
true_choice, oracle = ut.construct_opt(-1 * Y + 30, N_TRIAL, N_ARMS, M)

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

# Construct laplacian matrix
lda = 0.01
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))))

# Run CLUCB
print('M:', M, 'lda:', lda, 'T:', T)
max_it = 5
print('max_iters', max_it)

alpha_to_test = [0.5]
results_dict_clucb = {alpha: clucb.CLUCB_conv2(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=3000, max_iters = max_it)\
                for alpha in alpha_to_test}

max_it = 10
print('max_iters', max_it)

alpha_to_test = [0.5]
results_dict_clucb = {alpha: clucb.CLUCB_conv2(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=3000, max_iters = max_it)\
                for alpha in alpha_to_test}

max_it = 14
print('max_iters', max_it)

alpha_to_test = [0.5]
results_dict_clucb = {alpha: clucb.CLUCB_conv2(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=3000, max_iters = max_it)\
                for alpha in alpha_to_test}
