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 warnings
warnings.filterwarnings('ignore')
#eta_1, eta_2, lambda = tuning param.
#alpha_q, alpha_c = exploration parameters for q and C
#X = feature matrix
#Y = true reward matrix
#F = the sum of the identity matrix (N x N) and the laplacian matrix E
#FInv = The inverse of the sqrt of the kronecker product between the F and the identity matrix (K x K)
#m = number of selected arms
#K = number of latent structures

# Different versions of CLUCB
# CLUCB: in case of p = k + run ALS for 1 iteration
# CLUCB_alt: in case of p != k + run ALS for 1 iteration
# CLUCB_conv: run ALS until reach a param. threshold
# CLUCB_conv2: run ALS until finishing all maximum iterations
# CLUCB_conv_y: for collecting pred. y values in each trial
# CoCoUCB: prototype

# run ALS in CLUCB for 1 iteration to save comp. cost
def CLUCB(eta_1, eta_2, alpha_q, alpha_c, X, Y, init_q, init_c, m, K, FInv, oracle, true_choice, bb):
    print("CLUCB with eta_1:", eta_1, "eta_2:", eta_2, "alpha_q:", alpha_q, "alpha_c:", alpha_c)
    n_trial, n_arms, n_feature = X.shape
    #n_arms = N | n_trial = t | n_features = p | K

    # 1. Initialize objects
    # 1.1. Output objects                     
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    r_payoff = np.empty(n_trial)   
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    t_break = 0
    final_c_regret = 100000
    p = np.empty(shape=(n_trial, n_arms))
    
    ucb = np.empty(shape=(n_trial, n_arms))
    mse_sel = np.empty(n_trial)
    mse_sel_all = np.empty(n_trial)
    
    
    acc = np.empty(shape=(n_trial, n_arms))
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    opt_y = np.empty(n_trial)
    selected_pred = np.empty(n_trial)
    diff = np.empty(n_trial)
    
    
    opt_pred = np.empty(n_trial)  
    selected_err = np.empty(n_trial)
    opt_err = np.empty(n_trial)
    opt_payoff = np.empty(n_trial)
    
    theta_lst = np.empty(shape=(n_trial, n_arms * n_feature))
    C_lst = np.empty(shape=(n_trial, n_arms * K))
    q_lst = np.empty(shape=(n_trial, K * n_feature))

    
    # 1.2. Intialize q and C_tilde
    q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    q[0] = init_q
    C_tilde[0] = init_c
    
    te_q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    te_C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    
    # 1.4. A, b, D, d
    A = eta_1 * np.identity(n = K * n_feature) #Kp x Kp
    b = np.zeros(shape=K * n_feature)          #Kp x 1
    D = eta_2 * np.identity(n = n_arms * K)    #NK x NK
    d = np.zeros(shape=n_arms * K)             #NK x 1
    inv_A = np.linalg.inv(A)
    inv_D = np.linalg.inv(D)
    
    # 2. Algorithm
    for t in np.arange(n_trial):
        for a in np.arange(n_arms):
            
          #x_tr for q case 
          X_temp = ut.X_to_X_m(X, t, [a], n_arms, n_feature)    
          X_tr_init = np.kron(np.identity(n = K),X_temp[0].reshape(-1,1))
          X_tr = ut.X_reshape(X_temp, X_tr_init, t, K, n_arms, n_feature) #Kp x NK 
          
          #x_tr for c case
          X_tilde = FInv.dot(ut.X_to_X_m(X, t, [a], n_arms, n_feature).flatten()) #Np x 1
          # q_block = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T #Np x NK
          q_block = np.kron(np.eye(n_arms,dtype=int),q[t].reshape((K,n_feature)).T)
          #Calculate cb_q and cb_c
          #cb_q  
          X_q_a = X_tr.dot(FInv.dot(C_tilde[t]))
          cb_q = alpha_q * np.sqrt(X_q_a.T.dot(inv_A).dot(X_q_a))
        
          #cb_c
          X_c = q_block.T.dot(X_tilde)
          cb_c = alpha_c * np.sqrt((X_c).T.dot(inv_D).dot(X_c))
          
          #pred reward
          pred_reward[t, a] = (FInv.dot(C_tilde[t]).T).dot(X_tr.T).dot(q[t])
          
          #Predictions
          p[t, a] = pred_reward[t, a] + cb_q + cb_c
          acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        # Choose m best arms
        chosen_arms = p[t].argsort()[-m:][::-1]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
          
        selected_pred[t] = sum([pred_reward[t, a] for a in chosen_arms])
        opt_y[t] = sum([Y[t, a] for a in true_choice[t]])
        diff[t] = abs(selected_pred[t] - opt_y[t])
        
        opt_pred[t] = sum([pred_reward[t, a] for a in true_choice[t]])
        
        #Update A, b, D, d for each selected arm
        te_q[t] = q[t]
        te_C_tilde[t] = C_tilde[t]
        #for j in range(5):
        for chosen_arm in arm_choice[t]:
          X_tr_chosen_temp = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)
          X_tr_init_cs = np.kron(np.identity(n = K),X_tr_chosen_temp[0].reshape(-1,1))
          X_1_tr_chosen =  ut.X_reshape(X_tr_chosen_temp, X_tr_init_cs, t, K, n_arms, n_feature)
            
                #x_tr for c case
          X_tilde_chosen = FInv.dot(ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature).flatten())
                # q_block_chosen = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
                #q_block_chosen = (block_diag(*[te_q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
          q_block_chosen = np.kron(np.eye(n_arms,dtype=int),te_q[t].reshape((K,n_feature)).T)
            
              #Update  
          X_q = FInv.dot(te_C_tilde[t]).dot(X_1_tr_chosen.T)
          X_C_Tilde = q_block_chosen.T.dot(X_tilde_chosen)
          A = A + np.outer(X_q, X_q)
          b = b + Y[t, chosen_arm] * X_q
          D = D + np.outer(X_C_Tilde, X_C_Tilde)           
          d = d + Y[t, chosen_arm] * X_C_Tilde
        
            #inverse calculation
        inv_A = np.linalg.inv(A)
        inv_D = np.linalg.inv(D)
            # q[t + 1] = inv_A.dot(b)
            # C_tilde[t + 1] = inv_D.dot(d)
        te_q[t] = inv_A.dot(b)
        te_C_tilde[t] = inv_D.dot(d)
        q[t + 1] = te_q[t]
        C_tilde[t + 1] = te_C_tilde[t]
        
        
        q_lst[t] = q[t + 1]
        C_lst[t] = C_tilde[t + 1]
        theta_lst[t] = ((q[t + 1].reshape(n_feature, K)).dot(C_tilde[t + 1].reshape(K,n_arms)).flatten())
                
        # Cumulative rewards
        if t == 0:
            c_payoff[t] = np.sum([Y[t, choice] for choice in arm_choice[t]])
        else:
            c_payoff[t]   = c_payoff[t-1] + np.sum([Y[t, choice] for choice in arm_choice[t]])
        r_payoff[t]   = np.sum([Y[t, choice] for choice in arm_choice[t]])
        
        cum_regret[t] = np.sum(oracle[0:t+1] - r_payoff[0:t+1])
        
        opt_payoff[t] = np.sum([Y[t, choice] for choice in true_choice[t]])
        selected_err[t] = selected_pred[t] - r_payoff[t]
        opt_err[t] = opt_pred[t] - opt_payoff[t]

        mse_sel_all[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in np.arange(n_arms))/m
        mse_sel[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in chosen_arms)/m
        # print('TRIAL:', t, 'cum_regret:', cum_regret[t])
        
        if (t+1) % 250 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = t
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = n_trial
            final_c_regret = cum_regret[t]
        if cum_regret[t] > bb:
            print('break at t:', t)
            break
            # print('arm selected:', chosen_arms)
    return dict(mse_sel_all=mse_sel_all, mse_sel=mse_sel, alpha_q=alpha_q, alpha_c=alpha_c,theta_lst=theta_lst, q_lst=q_lst, C_lst=C_lst, q = q, C_tilde = C_tilde, p = p, arm_choice = arm_choice, r_payoff = r_payoff, c_payoff = c_payoff, cum_regret=cum_regret, t_break=t_break, opt_err=opt_err, final_c_regret=final_c_regret, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff)

# run ALS in CLUCB until each param. norm is less than thres.
def CLUCB_conv(eta_1, eta_2, alpha_q, alpha_c, X, Y, init_q, init_c, m, K, FInv, oracle, true_choice, bb, thres):
    print("CLUCB with eta_1:", eta_1, "eta_2:", eta_2, "alpha_q:", alpha_q, "alpha_c:", alpha_c)
    n_trial, n_arms, n_feature = X.shape
    #n_arms = N | n_trial = t | n_features = p | K

    # 1. Initialize objects
    # 1.1. Output objects                     
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    r_payoff = np.empty(n_trial)   
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    t_break = 0
    final_c_regret = 100000
    p = np.empty(shape=(n_trial, n_arms)) 
    acc = np.empty(shape=(n_trial, n_arms))
    
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    ucb = np.empty(shape=(n_trial, n_arms))
    mse_sel = np.empty(n_trial)
    mse_sel_all = np.empty(n_trial)


    selected_pred = np.empty(n_trial)
    opt_pred = np.empty(n_trial)  
    selected_err = np.empty(n_trial)
    opt_err = np.empty(n_trial)
    opt_payoff = np.empty(n_trial)
    
    theta_lst = np.empty(shape=(n_trial, n_arms * n_feature))
    C_lst = np.empty(shape=(n_trial, n_arms * K))
    q_lst = np.empty(shape=(n_trial, K * n_feature))

    
    # 1.2. Intialize q and C_tilde
    q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    q[0] = init_q
    C_tilde[0] = init_c
    
    te_q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    te_C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    
    # 1.4. A, b, D, d
    A = eta_1 * np.identity(n = K * n_feature) #Kp x Kp
    b = np.zeros(shape=K * n_feature)          #Kp x 1
    D = eta_2 * np.identity(n = n_arms * K)    #NK x NK
    d = np.zeros(shape=n_arms * K)             #NK x 1
    inv_A = np.linalg.inv(A)
    inv_D = np.linalg.inv(D)
    
    # 2. Algorithm
    for t in np.arange(n_trial):
        for a in np.arange(n_arms):
            
          #x_tr for q case 
          X_temp = ut.X_to_X_m(X, t, [a], n_arms, n_feature)    
          X_tr_init = np.kron(np.identity(n = K),X_temp[0].reshape(-1,1))
          X_tr = ut.X_reshape(X_temp, X_tr_init, t, K, n_arms, n_feature) #Kp x NK 
          
          #x_tr for c case
          X_tilde = FInv.dot(ut.X_to_X_m(X, t, [a], n_arms, n_feature).flatten()) #Np x 1
          # q_block = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T #Np x NK
          q_block = np.kron(np.eye(n_arms,dtype=int),q[t].reshape((K,n_feature)).T)
          #Calculate cb_q and cb_c
          #cb_q  
          X_q_a = X_tr.dot(FInv.dot(C_tilde[t]))
          cb_q = alpha_q * np.sqrt(X_q_a.T.dot(inv_A).dot(X_q_a))
        
          #cb_c
          X_c = q_block.T.dot(X_tilde)
          cb_c = alpha_c * np.sqrt((X_c).T.dot(inv_D).dot(X_c))
          
          #pred reward
          pred_reward[t, a] = (FInv.dot(C_tilde[t]).T).dot(X_tr.T).dot(q[t])
          
          #Predictions
          p[t, a] = pred_reward[t, a] + cb_q + cb_c
          acc[t, a]   = p[t, a] - Y[t, a]
        # Choose m best arms
        chosen_arms = p[t].argsort()[-m:][::-1]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
          
        selected_pred[t] = sum([pred_reward[t, a] for a in chosen_arms])
        opt_pred[t] = sum([pred_reward[t, a] for a in true_choice[t]])
        
        # Update A, b, D, d for each selected arm
        te_q[t] = q[t]
        te_C_tilde[t] = C_tilde[t]
        
        # Assign temp t + 1
        temp_q = 100*np.ones(shape= K * n_feature)
        temp_c = 100*np.ones(shape= K * n_arms)
        # te_q[t + 1] = 2 * te_q[t]
        # te_C_tilde[t + 1] = 2 * te_C_tilde[t]
        #for j in range(5):
        v = 0
        while np.linalg.norm(te_q[t] - temp_q, ord=2) > thres and np.linalg.norm(te_C_tilde[t] - temp_c, ord=2) > thres:
        # while np.linalg.norm(te_q[t + 1] - te_q[t], ord=2) > 10**-2 or np.linalg.norm(te_C_tilde[t + 1] - te_C_tilde[t], ord=2) > 10**-2:
            v += 1
            
            for chosen_arm in arm_choice[t]:
              X_tr_chosen_temp = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)
              X_tr_init_cs = np.kron(np.identity(n = K),X_tr_chosen_temp[0].reshape(-1,1))
              X_1_tr_chosen =  ut.X_reshape(X_tr_chosen_temp, X_tr_init_cs, t, K, n_arms, n_feature)
              X_tilde_chosen = FInv.dot(ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature).flatten())
              q_block_chosen = np.kron(np.eye(n_arms,dtype=int),te_q[t].reshape((K,n_feature)).T)
              # q_block_chosen = np.kron(np.eye(n_arms,dtype=int),temp_q.reshape((K,n_feature)).T)
              X_q = FInv.dot(te_C_tilde[t]).dot(X_1_tr_chosen.T)
              # X_q = FInv.dot(temp_c).dot(X_1_tr_chosen.T)
              X_C_Tilde = q_block_chosen.T.dot(X_tilde_chosen)
                #Update 
              A = A + np.outer(X_q, X_q)
              b = b + Y[t, chosen_arm] * X_q
              D = D + np.outer(X_C_Tilde, X_C_Tilde)           
              d = d + Y[t, chosen_arm] * X_C_Tilde

                #inverse calculation
            inv_A = np.linalg.inv(A)
            inv_D = np.linalg.inv(D)
                # q[t + 1] = inv_A.dot(b)
                # C_tilde[t + 1] = inv_D.dot(d)
            #save to temp before updating
            temp_q = np.copy(te_q[t])
            temp_c = np.copy(te_C_tilde[t])
            te_q[t] = inv_A.dot(b)
            te_C_tilde[t] = inv_D.dot(d)
        q[t + 1] = te_q[t]
        C_tilde[t + 1] = te_C_tilde[t]
        q_lst[t] = q[t + 1]
        C_lst[t] = C_tilde[t + 1]
        theta_lst[t] = ((q[t + 1].reshape(n_feature, K)).dot(C_tilde[t + 1].reshape(K,n_arms)).flatten())
                
        # Cumulative rewards
        if t == 0:
            c_payoff[t] = np.sum([Y[t, choice] for choice in arm_choice[t]])
        else:
            c_payoff[t]   = c_payoff[t-1] + np.sum([Y[t, choice] for choice in arm_choice[t]])
        r_payoff[t]   = np.sum([Y[t, choice] for choice in arm_choice[t]])
        
        cum_regret[t] = np.sum(oracle[0:t+1] - r_payoff[0:t+1])
        
        opt_payoff[t] = np.sum([Y[t, choice] for choice in true_choice[t]])
        selected_err[t] = selected_pred[t] - r_payoff[t]
        opt_err[t] = opt_pred[t] - opt_payoff[t]

        mse_sel_all[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in np.arange(n_arms))/m
        mse_sel[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in chosen_arms)/m
        # print('TRIAL:', t, 'cum_regret:', cum_regret[t])
        
        if (t+1) % 250 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = t
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = n_trial
            final_c_regret = cum_regret[t]
        if cum_regret[t] > bb:
            print('break at t:', t)
            break
            # print('arm selected:', chosen_arms)
    return dict(mse_sel_all=mse_sel_all, mse_sel=mse_sel, alpha_q=alpha_q, alpha_c=alpha_c,theta_lst=theta_lst, q_lst=q_lst, C_lst=C_lst, q = q, C_tilde = C_tilde, p = p, arm_choice = arm_choice, r_payoff = r_payoff, c_payoff = c_payoff, acc=acc, cum_regret=cum_regret, t_break=t_break,selected_err=selected_err, opt_err=opt_err, final_c_regret=final_c_regret)

# run ALS in CLUCB for max_iters iterations
def CLUCB_conv2(eta_1, eta_2, alpha_q, alpha_c, X, Y, init_q, init_c, m, K, FInv, oracle, true_choice, bb, max_iters):
    print("CLUCB with eta_1:", eta_1, "eta_2:", eta_2, "alpha_q:", alpha_q, "alpha_c:", alpha_c)
    n_trial, n_arms, n_feature = X.shape
    #n_arms = N | n_trial = t | n_features = p | K

    # 1. Initialize objects
    # 1.1. Output objects                     
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    r_payoff = np.empty(n_trial)   
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    t_break = 0
    final_c_regret = 100000
    p = np.empty(shape=(n_trial, n_arms)) 
    # acc = np.empty(shape=(n_trial, n_arms))
    
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    ucb = np.empty(shape=(n_trial, n_arms))
    mse_sel = np.empty(n_trial)
    mse_sel_all = np.empty(n_trial)

    acc = np.empty(shape=(n_trial, n_arms))
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    opt_y = np.empty(n_trial)
    selected_pred = np.empty(n_trial)
    diff = np.empty(n_trial)
    
    
    opt_pred = np.empty(n_trial)  
    selected_err = np.empty(n_trial)
    opt_err = np.empty(n_trial)
    opt_payoff = np.empty(n_trial)
    
    theta_lst = np.empty(shape=(n_trial, n_arms * n_feature))
    C_lst = np.empty(shape=(n_trial, n_arms * K))
    q_lst = np.empty(shape=(n_trial, K * n_feature))

    
    # 1.2. Intialize q and C_tilde
    q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    q[0] = init_q
    C_tilde[0] = init_c
    
    te_q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    te_C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    
    # 1.4. A, b, D, d
    A = eta_1 * np.identity(n = K * n_feature) #Kp x Kp
    b = np.zeros(shape=K * n_feature)          #Kp x 1
    D = eta_2 * np.identity(n = n_arms * K)    #NK x NK
    d = np.zeros(shape=n_arms * K)             #NK x 1
    inv_A = np.linalg.inv(A)
    inv_D = np.linalg.inv(D)
    
    # 2. Algorithm
    for t in np.arange(n_trial):
        for a in np.arange(n_arms):
            
          #x_tr for q case 
          X_temp = ut.X_to_X_m(X, t, [a], n_arms, n_feature)    
          X_tr_init = np.kron(np.identity(n = K),X_temp[0].reshape(-1,1))
          X_tr = ut.X_reshape(X_temp, X_tr_init, t, K, n_arms, n_feature) #Kp x NK 
          
          #x_tr for c case
          X_tilde = FInv.dot(ut.X_to_X_m(X, t, [a], n_arms, n_feature).flatten()) #Np x 1
          # q_block = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T #Np x NK
          q_block = np.kron(np.eye(n_arms,dtype=int),q[t].reshape((K,n_feature)).T)
          #Calculate cb_q and cb_c
          #cb_q  
          X_q_a = X_tr.dot(FInv.dot(C_tilde[t]))
          cb_q = alpha_q * np.sqrt(X_q_a.T.dot(inv_A).dot(X_q_a))
        
          #cb_c
          X_c = q_block.T.dot(X_tilde)
          cb_c = alpha_c * np.sqrt((X_c).T.dot(inv_D).dot(X_c))
          
          #pred reward
          pred_reward[t, a] = (FInv.dot(C_tilde[t]).T).dot(X_tr.T).dot(q[t])
          
          #Predictions
          p[t, a] = pred_reward[t, a] + cb_q + cb_c
          acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        # Choose m best arms
        chosen_arms = p[t].argsort()[-m:][::-1]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
          
        selected_pred[t] = sum([pred_reward[t, a] for a in chosen_arms])
        opt_y[t] = sum([Y[t, a] for a in true_choice[t]])
        diff[t] = abs(selected_pred[t] - opt_y[t])
        opt_pred[t] = sum([pred_reward[t, a] for a in true_choice[t]])
        
        #Update A, b, D, d for each selected arm
        te_q[t] = q[t]
        te_C_tilde[t] = C_tilde[t]
        for j in range(max_iters):
            for chosen_arm in arm_choice[t]:
              X_tr_chosen_temp = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)
              X_tr_init_cs = np.kron(np.identity(n = K),X_tr_chosen_temp[0].reshape(-1,1))
              X_1_tr_chosen =  ut.X_reshape(X_tr_chosen_temp, X_tr_init_cs, t, K, n_arms, n_feature)

                    #x_tr for c case
              X_tilde_chosen = FInv.dot(ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature).flatten())
                    # q_block_chosen = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
                    #q_block_chosen = (block_diag(*[te_q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
              q_block_chosen = np.kron(np.eye(n_arms,dtype=int),te_q[t].reshape((K,n_feature)).T)

                  #Update  
              X_q = FInv.dot(te_C_tilde[t]).dot(X_1_tr_chosen.T)
              X_C_Tilde = q_block_chosen.T.dot(X_tilde_chosen)
              A = A + np.outer(X_q, X_q)
              b = b + Y[t, chosen_arm] * X_q
              D = D + np.outer(X_C_Tilde, X_C_Tilde)           
              d = d + Y[t, chosen_arm] * X_C_Tilde
        
            #inverse calculation
            inv_A = np.linalg.inv(A)
            inv_D = np.linalg.inv(D)
            # q[t + 1] = inv_A.dot(b)
            # C_tilde[t + 1] = inv_D.dot(d)
            te_q[t] = inv_A.dot(b)
            te_C_tilde[t] = inv_D.dot(d)
        q[t + 1] = te_q[t]
        C_tilde[t + 1] = te_C_tilde[t]
        
        
        q_lst[t] = q[t + 1]
        C_lst[t] = C_tilde[t + 1]
        theta_lst[t] = ((q[t + 1].reshape(n_feature, K)).dot(C_tilde[t + 1].reshape(K,n_arms)).flatten())
                
        # Cumulative rewards
        if t == 0:
            c_payoff[t] = np.sum([Y[t, choice] for choice in arm_choice[t]])
        else:
            c_payoff[t]   = c_payoff[t-1] + np.sum([Y[t, choice] for choice in arm_choice[t]])
        r_payoff[t]   = np.sum([Y[t, choice] for choice in arm_choice[t]])
        
        cum_regret[t] = np.sum(oracle[0:t+1] - r_payoff[0:t+1])
        
        opt_payoff[t] = np.sum([Y[t, choice] for choice in true_choice[t]])
        selected_err[t] = selected_pred[t] - r_payoff[t]
        opt_err[t] = opt_pred[t] - opt_payoff[t]

        mse_sel_all[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in np.arange(n_arms))/m
        mse_sel[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in chosen_arms)/m
        # print('TRIAL:', t, 'cum_regret:', cum_regret[t])
        
        if (t+1) % 250 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = t
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = n_trial
            final_c_regret = cum_regret[t]
        if cum_regret[t] > bb:
            print('break at t:', t)
            break
            # print('arm selected:', chosen_arms)
    return dict(mse_sel_all=mse_sel_all, mse_sel=mse_sel, alpha_q=alpha_q, alpha_c=alpha_c,theta_lst=theta_lst, q_lst=q_lst, C_lst=C_lst, q = q, C_tilde = C_tilde, p = p, arm_choice = arm_choice, r_payoff = r_payoff, c_payoff = c_payoff, cum_regret=cum_regret, t_break=t_break,selected_err=selected_err, opt_err=opt_err, final_c_regret=final_c_regret, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff)

# run ALS in CLUCB for max_iters iterations + for testing predicted y values
def CLUCB_conv_y(eta_1, eta_2, alpha_q, alpha_c, X, Y, init_q, init_c, m, K, FInv, oracle, true_choice, bb, max_iters):
    print("This run is for testing predicted y.")
    print("CLUCB with eta_1:", eta_1, "eta_2:", eta_2, "alpha_q:", alpha_q, "alpha_c:", alpha_c)
    n_trial, n_arms, n_feature = X.shape
    #n_arms = N | n_trial = t | n_features = p | K

    # 1. Initialize objects
    # 1.1. Output objects                     
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    r_payoff = np.empty(n_trial)   
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    t_break = 0
    final_c_regret = 100000
    p = np.empty(shape=(n_trial, n_arms)) 
    # acc = np.empty(shape=(n_trial, n_arms))
    
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    ucb = np.empty(shape=(n_trial, n_arms))
    mse_sel = np.empty(n_trial)
    mse_sel_all = np.empty(n_trial)

    acc = np.empty(shape=(n_trial, n_arms))
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    opt_y = np.empty(n_trial)
    selected_pred = np.empty(n_trial)
    diff = np.empty(n_trial)
    
    
    opt_pred = np.empty(n_trial)  
    selected_err = np.empty(n_trial)
    opt_err = np.empty(n_trial)
    opt_payoff = np.empty(n_trial)
    
    theta_lst = np.empty(shape=(n_trial, n_arms * n_feature))
    C_lst = np.empty(shape=(n_trial, n_arms * K))
    q_lst = np.empty(shape=(n_trial, K * n_feature))

    
    # 1.2. Intialize q and C_tilde
    q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    q[0] = init_q
    C_tilde[0] = init_c
    
    te_q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    te_C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    
    # 1.4. A, b, D, d
    A = eta_1 * np.identity(n = K * n_feature) #Kp x Kp
    b = np.zeros(shape=K * n_feature)          #Kp x 1
    D = eta_2 * np.identity(n = n_arms * K)    #NK x NK
    d = np.zeros(shape=n_arms * K)             #NK x 1
    inv_A = np.linalg.inv(A)
    inv_D = np.linalg.inv(D)
    
    # 2. Algorithm
    for t in np.arange(n_trial):
        for a in np.arange(n_arms):
            
          #x_tr for q case 
          X_temp = ut.X_to_X_m(X, t, [a], n_arms, n_feature)    
          X_tr_init = np.kron(np.identity(n = K),X_temp[0].reshape(-1,1))
          X_tr = ut.X_reshape(X_temp, X_tr_init, t, K, n_arms, n_feature) #Kp x NK 
          
          #x_tr for c case
          X_tilde = FInv.dot(ut.X_to_X_m(X, t, [a], n_arms, n_feature).flatten()) #Np x 1
          # q_block = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T #Np x NK
          q_block = np.kron(np.eye(n_arms,dtype=int),q[t].reshape((K,n_feature)).T)
          #Calculate cb_q and cb_c
          #cb_q  
          X_q_a = X_tr.dot(FInv.dot(C_tilde[t]))
          cb_q = alpha_q * np.sqrt(X_q_a.T.dot(inv_A).dot(X_q_a))
        
          #cb_c
          X_c = q_block.T.dot(X_tilde)
          cb_c = alpha_c * np.sqrt((X_c).T.dot(inv_D).dot(X_c))
          
          #pred reward
          pred_reward[t, a] = (FInv.dot(C_tilde[t]).T).dot(X_tr.T).dot(q[t])
          
          #Predictions
          p[t, a] = pred_reward[t, a] + cb_q + cb_c
          acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        if t+1 == n_trial:
            # theta_matrix = (q[t].reshape(n_feature, K).T).dot((FInv.dot(C_tilde[t]).real).reshape(n_arms, K).T).T
            for a in np.arange(n_arms):
                print("arm", a, ":", pred_reward[t, a]) 
        # Choose m best arms
        chosen_arms = p[t].argsort()[-m:][::-1]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
          
        selected_pred[t] = sum([pred_reward[t, a] for a in chosen_arms])
        opt_y[t] = sum([Y[t, a] for a in true_choice[t]])
        diff[t] = abs(selected_pred[t] - opt_y[t])
        opt_pred[t] = sum([pred_reward[t, a] for a in true_choice[t]])
        
        #Update A, b, D, d for each selected arm
        te_q[t] = q[t]
        te_C_tilde[t] = C_tilde[t]
        for j in range(max_iters):
            for chosen_arm in arm_choice[t]:
              X_tr_chosen_temp = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)
              X_tr_init_cs = np.kron(np.identity(n = K),X_tr_chosen_temp[0].reshape(-1,1))
              X_1_tr_chosen =  ut.X_reshape(X_tr_chosen_temp, X_tr_init_cs, t, K, n_arms, n_feature)

                    #x_tr for c case
              X_tilde_chosen = FInv.dot(ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature).flatten())
                    # q_block_chosen = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
                    #q_block_chosen = (block_diag(*[te_q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
              q_block_chosen = np.kron(np.eye(n_arms,dtype=int),te_q[t].reshape((K,n_feature)).T)

                  #Update  
              X_q = FInv.dot(te_C_tilde[t]).dot(X_1_tr_chosen.T)
              X_C_Tilde = q_block_chosen.T.dot(X_tilde_chosen)
              A = A + np.outer(X_q, X_q)
              b = b + Y[t, chosen_arm] * X_q
              D = D + np.outer(X_C_Tilde, X_C_Tilde)           
              d = d + Y[t, chosen_arm] * X_C_Tilde
        
            #inverse calculation
            inv_A = np.linalg.inv(A)
            inv_D = np.linalg.inv(D)
            # q[t + 1] = inv_A.dot(b)
            # C_tilde[t + 1] = inv_D.dot(d)
            te_q[t] = inv_A.dot(b)
            te_C_tilde[t] = inv_D.dot(d)
        q[t + 1] = te_q[t]
        C_tilde[t + 1] = te_C_tilde[t]
        
        
        q_lst[t] = q[t + 1]
        C_lst[t] = C_tilde[t + 1]
        theta_lst[t] = ((q[t + 1].reshape(n_feature, K)).dot(C_tilde[t + 1].reshape(K,n_arms)).flatten())
                
        # Cumulative rewards
        if t == 0:
            c_payoff[t] = np.sum([Y[t, choice] for choice in arm_choice[t]])
        else:
            c_payoff[t]   = c_payoff[t-1] + np.sum([Y[t, choice] for choice in arm_choice[t]])
        r_payoff[t]   = np.sum([Y[t, choice] for choice in arm_choice[t]])
        
        cum_regret[t] = np.sum(oracle[0:t+1] - r_payoff[0:t+1])
        
        opt_payoff[t] = np.sum([Y[t, choice] for choice in true_choice[t]])
        selected_err[t] = selected_pred[t] - r_payoff[t]
        opt_err[t] = opt_pred[t] - opt_payoff[t]

        mse_sel_all[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in np.arange(n_arms))/m
        mse_sel[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in chosen_arms)/m
        # print('TRIAL:', t, 'cum_regret:', cum_regret[t])
        
        if (t+1) % 1000 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = t
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = n_trial
            final_c_regret = cum_regret[t]
            
        if cum_regret[t] > bb:
            print('break at t:', t)
            break
            # print('arm selected:', chosen_arms)
    return dict(mse_sel_all=mse_sel_all, mse_sel=mse_sel, alpha_q=alpha_q, alpha_c=alpha_c,theta_lst=theta_lst, q_lst=q_lst, C_lst=C_lst, q = q, C_tilde = C_tilde, p = p, arm_choice = arm_choice, r_payoff = r_payoff, c_payoff = c_payoff, cum_regret=cum_regret, t_break=t_break,selected_err=selected_err, opt_err=opt_err, final_c_regret=final_c_regret, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff)

def CLUCB_alt(eta_1, eta_2, alpha_q, alpha_c, X, Y, init_q, init_c, m, K, FInv_NK, FInv_Np, oracle, min_cregret, true_choice):
    print("CoCoUCB with eta_1:", eta_1, "eta_2:", eta_2, "alpha_q:", alpha_q, "alpha_c:", alpha_c)
    n_trial, n_arms, n_feature = X.shape
    #n_arms = N | n_trial = t | n_features = p | K

    # 1. Initialize objects
    # 1.1. Output objects                     
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    r_payoff = np.empty(n_trial)   
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    t_break = 0
    final_c_regret = 100000
    p = np.empty(shape=(n_trial, n_arms)) 
    acc = np.empty(shape=(n_trial, n_arms))
    
    pred_reward = np.empty(shape=(n_trial, n_arms)) 
    ucb = np.empty(shape=(n_trial, n_arms))
    mse_sel = np.empty(n_trial)
    mse_sel_all = np.empty(n_trial)


    selected_pred = np.empty(n_trial)
    opt_pred = np.empty(n_trial)  
    selected_err = np.empty(n_trial)
    opt_err = np.empty(n_trial)
    opt_payoff = np.empty(n_trial)
    
    theta_lst = np.empty(shape=(n_trial, n_arms * n_feature))
    C_lst = np.empty(shape=(n_trial, n_arms * K))
    q_lst = np.empty(shape=(n_trial, K * n_feature))

    
    # 1.2. Intialize q and C_tilde
    q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    q[0] = init_q
    C_tilde[0] = init_c
    
    te_q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    te_C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    
    # 1.4. A, b, D, d
    A = eta_1 * np.identity(n = K * n_feature) #Kp x Kp
    b = np.zeros(shape=K * n_feature)          #Kp x 1
    D = eta_2 * np.identity(n = n_arms * K)    #NK x NK
    d = np.zeros(shape=n_arms * K)             #NK x 1
    inv_A = np.linalg.inv(A)
    inv_D = np.linalg.inv(D)
    
    # 2. Algorithm
    for t in np.arange(n_trial):
        for a in np.arange(n_arms):
            
          #x_tr for q case 
          X_temp = ut.X_to_X_m(X, t, [a], n_arms, n_feature)    
          X_tr_init = np.kron(np.identity(n = K),X_temp[0].reshape(-1,1))
          X_tr = ut.X_reshape(X_temp, X_tr_init, t, K, n_arms, n_feature) #Kp x NK 
          
          #x_tr for c case
          # X_tilde = FInv_Np.dot(X_to_X_m(X, t, [a], n_arms, n_feature).flatten()) #Np x 1
        
        
        
          X_tilde = ut.X_to_X_m(X, t, [a], n_arms, n_feature).flatten() #Np x 1
          # q_block = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T #Np x NK
          q_block = np.kron(np.eye(n_arms,dtype=int),q[t].reshape((K,n_feature)).T).dot(FInv_NK)
          #Calculate cb_q and cb_c
          #cb_q  
          X_q_a = X_tr.dot(FInv_NK.dot(C_tilde[t]))
          cb_q = alpha_q * np.sqrt(X_q_a.T.dot(inv_A).dot(X_q_a))
        
          #cb_c
          X_c = q_block.T.dot(X_tilde)
          cb_c = alpha_c * np.sqrt((X_c).T.dot(inv_D).dot(X_c))
          
          #pred reward
          pred_reward[t, a] = (FInv_NK.dot(C_tilde[t]).T).dot(X_tr.T).dot(q[t])
          
          #Predictions
          p[t, a] = pred_reward[t, a] + cb_q + cb_c
          acc[t, a]   = p[t, a] - Y[t, a]
        # Choose m best arms
        chosen_arms = p[t].argsort()[-m:][::-1]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
          
        selected_pred[t] = sum([pred_reward[t, a] for a in chosen_arms])
        opt_pred[t] = sum([pred_reward[t, a] for a in true_choice[t]])
        
        #Update A, b, D, d for each selected arm
        te_q[t] = q[t]
        te_C_tilde[t] = C_tilde[t]
        #for j in range(5):
        for chosen_arm in arm_choice[t]:
          X_tr_chosen_temp = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)
          X_tr_init_cs = np.kron(np.identity(n = K),X_tr_chosen_temp[0].reshape(-1,1))
          X_1_tr_chosen =  ut.X_reshape(X_tr_chosen_temp, X_tr_init_cs, t, K, n_arms, n_feature)
            
                #x_tr for c case
          X_tilde_chosen = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature).flatten()
                # q_block_chosen = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
                #q_block_chosen = (block_diag(*[te_q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
          q_block_chosen = np.kron(np.eye(n_arms,dtype=int),te_q[t].reshape((K,n_feature)).T).dot(FInv_NK)
            
              #Update  
          X_q = FInv_NK.dot(te_C_tilde[t]).dot(X_1_tr_chosen.T)
          X_C_Tilde = q_block_chosen.T.dot(X_tilde_chosen)
          A = A + np.outer(X_q, X_q)
          b = b + Y[t, chosen_arm] * X_q
          D = D + np.outer(X_C_Tilde, X_C_Tilde)           
          d = d + Y[t, chosen_arm] * X_C_Tilde
        
            #inverse calculation
        inv_A = np.linalg.inv(A)
        inv_D = np.linalg.inv(D)
            # q[t + 1] = inv_A.dot(b)
            # C_tilde[t + 1] = inv_D.dot(d)
        te_q[t] = inv_A.dot(b)
        te_C_tilde[t] = inv_D.dot(d)
        q[t + 1] = te_q[t]
        C_tilde[t + 1] = te_C_tilde[t]
        
        
        q_lst[t] = q[t + 1]
        C_lst[t] = C_tilde[t + 1]
        theta_lst[t] = ((q[t + 1].reshape(n_feature, K)).dot(C_tilde[t + 1].reshape(K,n_arms)).flatten())
                
        # Cumulative rewards
        if t == 0:
            c_payoff[t] = np.sum([Y[t, choice] for choice in arm_choice[t]])
        else:
            c_payoff[t]   = c_payoff[t-1] + np.sum([Y[t, choice] for choice in arm_choice[t]])
        r_payoff[t]   = np.sum([Y[t, choice] for choice in arm_choice[t]])
        
        cum_regret[t] = np.sum(oracle[0:t+1] - r_payoff[0:t+1])
        
        opt_payoff[t] = np.sum([Y[t, choice] for choice in true_choice[t]])
        selected_err[t] = selected_pred[t] - r_payoff[t]
        opt_err[t] = opt_pred[t] - opt_payoff[t]

        mse_sel_all[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in np.arange(n_arms))/m
        mse_sel[t] = sum(np.square(pred_reward[t, v] - Y[t, v]) for v in chosen_arms)/m
        # print('TRIAL:', t, 'cum_regret:', cum_regret[t])
        if cum_regret[t] > min_cregret:
            print('break at trial:', t, 'with cum_regret:', cum_regret[t])
            break
        if (t+1) % 3000 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = t
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
            t_break = n_trial
            final_c_regret = cum_regret[t]
    return dict(mse_sel_all=mse_sel_all, mse_sel=mse_sel, alpha_q=alpha_q, alpha_c=alpha_c,theta_lst=theta_lst, q_lst=q_lst, C_lst=C_lst, q = q, C_tilde = C_tilde, p = p, arm_choice = arm_choice, r_payoff = r_payoff, c_payoff = c_payoff, acc=acc, cum_regret=cum_regret, t_break=t_break,selected_err=selected_err, opt_err=opt_err, final_c_regret=final_c_regret)

# prototype
def CoCoUCB(eta_1, eta_2, alpha_q, alpha_c, X, Y, init_q, init_c, m, K, FInv, oracle):
    print("CoCoUCB with eta_1:", eta_1, "eta_2:", eta_2, "alpha_q:", alpha_q, "alpha_c:", alpha_c)
    n_trial, n_arms, n_feature = X.shape
    #n_arms = N | n_trial = t | n_features = p | K

    # 1. Initialize objects
    # 1.1. Output objects                     
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    r_payoff = np.empty(n_trial)   
    c_payoff = np.empty(n_trial) 
    p = np.empty(shape=(n_trial, n_arms)) 
    acc = np.empty(shape=(n_trial, n_arms))
    cum_regret = np.empty(n_trial)
    theta_lst = np.empty(shape=(n_trial, n_arms * n_feature))
    C_lst = np.empty(shape=(n_trial, n_arms * K))
    q_lst = np.empty(shape=(n_trial, K * n_feature))

    
    # 1.2. Intialize q and C_tilde
    q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    q[0] = init_q
    C_tilde[0] = init_c
    
    te_q = np.empty(shape = (n_trial + 1, K * n_feature)) #Kp x 1
    te_C_tilde = np.empty(shape = (n_trial + 1, n_arms * K)) #NK x 1
    
    # 1.4. A, b, D, d
    A = eta_1 * np.identity(n = K * n_feature) #Kp x Kp
    b = np.zeros(shape=K * n_feature)          #Kp x 1
    D = eta_2 * np.identity(n = n_arms * K)    #NK x NK
    d = np.zeros(shape=n_arms * K)             #NK x 1
    inv_A = np.linalg.inv(A)
    inv_D = np.linalg.inv(D)
    
    # 2. Algorithm
    for t in np.arange(n_trial):
        for a in np.arange(n_arms):
            
          #x_tr for q case 
          X_temp = ut.X_to_X_m(X, t, [a], n_arms, n_feature)    
          X_tr_init = np.kron(np.identity(n = K),X_temp[0].reshape(-1,1))
          X_tr = ut.X_reshape(X_temp, X_tr_init, t, K, n_arms, n_feature) #Kp x NK 
          
          #x_tr for c case
          X_tilde = FInv.dot(ut.X_to_X_m(X, t, [a], n_arms, n_feature).flatten()) #Np x 1
          q_block = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T #Np x NK  
        
          #Calculate cb_q and cb_c
          #cb_q  
          X_q_a = X_tr.dot(FInv.dot(C_tilde[t]))
          cb_q = alpha_q * np.sqrt(X_q_a.T.dot(inv_A).dot(X_q_a))
        
          #cb_c
          X_c = q_block.T.dot(X_tilde)
          cb_c = alpha_c * np.sqrt((X_c).T.dot(inv_D).dot(X_c))
        
          #Predictions
          p[t, a] = (FInv.dot(C_tilde[t]).T).dot(X_tr.T).dot(q[t]) + cb_q + cb_c
          acc[t, a]   = p[t, a] - Y[t, a]
        # Choose m best arms
        chosen_arms = p[t].argsort()[-m:][::-1]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
        
        #Update A, b, D, d for each selected arm
        te_q[t] = q[t]
        te_C_tilde[t] = C_tilde[t]
        for chosen_arm in arm_choice[t]:
            X_tr_chosen_temp = ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)
            X_tr_init_cs = np.kron(np.identity(n = K),X_tr_chosen_temp[0].reshape(-1,1))
            X_1_tr_chosen =  ut.X_reshape(X_tr_chosen_temp, X_tr_init_cs, t, K, n_arms, n_feature)
            
            #x_tr for c case
            X_tilde_chosen = FInv.dot(ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature).flatten())
            # q_block_chosen = (block_diag(*[q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
            q_block_chosen = (block_diag(*[te_q[t].reshape((K, n_feature)) for _ in np.arange(n_arms)])).T
            
            #Update  
            X_q = FInv.dot(te_C_tilde[t]).dot(X_1_tr_chosen.T)
            X_C_Tilde = q_block_chosen.T.dot(X_tilde_chosen)
            A = A + np.outer(X_q, X_q)
            b = b + Y[t, chosen_arm] * X_q
            D = D + np.outer(X_C_Tilde, X_C_Tilde)           
            d = d + Y[t, chosen_arm] * X_C_Tilde
        
        #inverse calculation
        inv_A = np.linalg.inv(A)
        inv_D = np.linalg.inv(D)
        # q[t + 1] = inv_A.dot(b)
        # C_tilde[t + 1] = inv_D.dot(d)
        te_q[t] = inv_A.dot(b)
        te_C_tilde[t] = inv_D.dot(d)
        q[t + 1] = te_q[t]
        C_tilde[t + 1] = te_C_tilde[t]
        #just for recording data
        q_lst[t] = q[t + 1]
        C_lst[t] = C_tilde[t + 1]
        theta_lst[t] = ((q[t + 1].reshape(n_feature, K)).dot(C_tilde[t + 1].reshape(K,n_arms)).flatten())
                
        # Cumulative rewards
        if t == 0:
            c_payoff[t] = np.sum([Y[t, choice] for choice in arm_choice[t]])
        else:
            c_payoff[t]   = c_payoff[t-1] + np.sum([Y[t, choice] for choice in arm_choice[t]])
        r_payoff[t]   = np.sum([Y[t, choice] for choice in arm_choice[t]])
        cum_regret[t] = np.sum(oracle[0:t+1] - r_payoff[0:t+1])
        print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
        if (t+1) % 1000 == 0:
            # print('TRIAL:',t,'DONE', '| arm selected:', chosen_arms)
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
    return dict(theta_lst=theta_lst, q_lst=q_lst, C_lst=C_lst, q = q, C_tilde = C_tilde, p = p, arm_choice = arm_choice, r_payoff = r_payoff, c_payoff = c_payoff, acc=acc)
