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
from sklearn.metrics import r2_score
# APPENDIX D: RUN EXPERIMENTS for comparison between pred. y and true y

# 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 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.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))))

# 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))))

# Run LinUCB
print('LinUCB')
alpha_to_test =[2.5]
results_dict_linucb = {alpha: linucb.linUCB_conv_y(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 = [2.5]
results_dict_gob = {alpha: goblin.GOBLin_conv_y(alpha = alpha, eta = 0.3, X=X, Y=(-1 * Y + 30).T, m=M, FInv=FInv_Init, 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 = [2.5]
results_dict_clucbwo = {alpha: clucb.CLUCB_conv_y(eta_1 = 0.3, eta_2 = 0.3, alpha_q = 2.5, 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, max_iters=5)\
                for alpha in alpha_to_test}

# Run CLUCB
print('CLUCB--', 'lda:', lda)
alpha_to_test = [2.5]
results_dict_clucb = {alpha: clucb.CLUCB_conv_y(eta_1 = 0.3, eta_2 = 0.3, alpha_q = 2.5, 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, max_iters=5)\
                for alpha in alpha_to_test}

#plot pred. y vs true y
true_y = Y.T[N_TRIAL-1]
pd_y_linucb = results_dict_linucb[2.5]['pred_reward'][N_TRIAL-1]
pd_y_goblin = results_dict_gob[2.5]['pred_reward'][N_TRIAL-1]
pd_y_clucbwo = results_dict_clucbwo[2.5]['pred_reward'][N_TRIAL-1]
pd_y_clucb = results_dict_clucb[2.5]['pred_reward'][N_TRIAL-1]

ut.plot_y_all(pd_y_linucb, pd_y_goblin, pd_y_clucbwo, pd_y_clucb, true_y)

