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 misspecified K
# results vary: al_q and al_c need to be tuned

# For the full experiment, please change N_TRIAL to N_TRIAL = 30000
N_TRIAL = 5
N_ARMS = 100
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

min_target_cregret= 100000

# Construct laplacian matrix
K = 2
lda = 0.01
T = 0.5
test_adj = ut.constructAdjMatrix(W, N_ARMS, T)
test_F = ut.constructLaplacianMatrix(test_adj, N_ARMS, lda)
F_kron = np.kron(test_F, np.identity(n=N_FEATURE))
FInv_NK = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=K))))
FInv_Np = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=N_FEATURE))))

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 * K)])
print("K:", K)
# Run CLUCB
alpha_to_test = [2]
results_dict_clucb = {alpha: clucb.CLUCB_alt(min_cregret = min_target_cregret, eta_1 = 0.3, eta_2 = 0.3, alpha_q = 2, alpha_c = alpha, X=X, Y=(-1 * Y + 30).T, init_q=vec_q, init_c=vec_C,m=M, K = K, FInv_NK=FInv_NK, FInv_Np=FInv_Np, oracle=oracle, true_choice=true_choice)\
                for alpha in alpha_to_test}

# Construct laplacian matrix
K = 4
lda = 0.01
T = 0.5
test_adj = ut.constructAdjMatrix(W, N_ARMS, T)
test_F = ut.constructLaplacianMatrix(test_adj, N_ARMS, lda)
F_kron = np.kron(test_F, np.identity(n=N_FEATURE))
FInv_NK = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=K))))
FInv_Np = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=N_FEATURE))))

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 * K)])
print("K:", K)
# Run CLUCB
alpha_to_test = [2]
results_dict_clucb = {alpha: clucb.CLUCB_alt(min_cregret = min_target_cregret, eta_1 = 0.3, eta_2 = 0.3, alpha_q = 2, alpha_c = alpha, X=X, Y=(-1 * Y + 30).T, init_q=vec_q, init_c=vec_C,m=M, K = K, FInv_NK=FInv_NK, FInv_Np=FInv_Np, oracle=oracle, true_choice=true_choice)\
                for alpha in alpha_to_test}

# Construct laplacian matrix
K = 5
lda = 0.01
T = 0.5
test_adj = ut.constructAdjMatrix(W, N_ARMS, T)
test_F = ut.constructLaplacianMatrix(test_adj, N_ARMS, lda)
F_kron = np.kron(test_F, np.identity(n=N_FEATURE))
FInv_NK = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=K))))
FInv_Np = sqrtm(np.linalg.inv(np.kron(test_F, np.identity(n=N_FEATURE))))

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 * K)])
print("K:", K)
# Run CLUCB
alpha_to_test = [2]
results_dict_clucb = {alpha: clucb.CLUCB_alt(min_cregret = min_target_cregret, eta_1 = 0.3, eta_2 = 0.3, alpha_q = 2, alpha_c = alpha, X=X, Y=(-1 * Y + 30).T, init_q=vec_q, init_c=vec_C,m=M, K = K, FInv_NK=FInv_NK, FInv_Np=FInv_Np, oracle=oracle, true_choice=true_choice)\
                for alpha in alpha_to_test}

plt.figure()
K_plot = np.array([2,3,4,5])
Y_CLUCB_plot = np.array([1582.26,1236.98,1499.99,2768.88])
Y_CLUCBWO_plot = np.array([2282.20,1434.41,1795.32,2862.11])
# plt.figure()
plt.plot(K_plot, Y_CLUCBWO_plot, marker='p', markevery=1, ls = '-.', label='CLUCB (w/o sim. reg.)')
plt.plot(K_plot, Y_CLUCB_plot, marker='P', markevery=1, ls = '--', label='CLUCB')
plt.ylabel("Cumulative Regret",fontweight='bold',fontsize = 16)
plt.xlabel("K",fontweight='bold',fontsize = 16)
plt.xticks(range(2, 6,1),fontsize=14)
plt.yticks(fontsize=14)
plt.legend(prop={'size': 12}, loc='upper left')
plt.show()
plt.savefig('/results/miss_K.png')
