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
from sklearn.metrics import r2_score
from sklearn.metrics import mean_squared_error

def plot_regrets(results, oracle):
    [plt.plot(make_regret(payoff=x['r_payoff'], oracle=oracle), label="alpha: "+str(alpha)) for (alpha, x) in results.items()]

    
def plot_l2_norm_diff(results, true_theta, n_trial, n_arms, n_feature, FInv_Init, true_choice, K):
    [plt.plot(make_l2_norm_diff(x['theta_lst'], x['q_lst'], x['C_lst'], true_theta, n_trial, n_arms, n_feature, FInv_Init, true_choice, K), label="alpha: "+str(alpha)) for (alpha, x) in results.items()]    
# l2_norm_diff

def make_l2_norm_diff_linucb(theta, true_theta, n_trial, n_arms, true_choice): #, true_choice
    diff = []
    for t in range(n_trial):
        diff_t = 0
        for arm in range(n_arms):
            diff_vec = np.array([theta[t, arm] - (true_theta[arm])])
            diff_t += np.linalg.norm(diff_vec, ord=2)
        diff.append(diff_t)
    return np.array(diff)

def make_l2_norm_diff_goblin(theta, true_theta, n_trial, n_arms, n_feature, FInv_Init):
    diff = []
    for t in range(n_trial):
        theta_matrix = (FInv_Init.dot(theta[t])).reshape(n_arms, n_feature) 
        diff_t = 0
        for arm in range(n_arms):
            diff_vec = theta_matrix[arm] - (true_theta[arm])
            diff_t += np.linalg.norm(diff_vec, ord=2)
        diff.append(diff_t)
    return np.array(diff)

def make_l2_norm_diff_CLUCB(theta, q_lst, C_lst, true_theta, n_trial, n_arms, n_feature, FInv_Init, true_choice, K):
    diff = []
    for t in range(n_trial):
        theta_matrix = (q_lst[t].reshape(n_feature, K).T).dot((FInv_Init.dot(C_lst[t]).real).reshape(n_arms, K).T).T
        diff_t = 0
        for arm in range(n_arms):
            diff_vec = np.array([theta_matrix[arm] - (true_theta[arm])])
            diff_t += np.linalg.norm(diff_vec, ord=2) 
        diff.append(diff_t)
    return np.array(diff)

def reg_plot_all_c(A, B, C, D, n_trial, n, m):
    plt.figure()
    X = np.arange(n_trial)
    plt.plot(X, A, marker='.', markevery=1000, ls = ':', label='LinUCB')
    plt.plot(X, B, marker='s', markevery=1000, ls = '-.', label='GOB.Lin')
    plt.plot(X, C, marker='p', markevery=1000, ls = '-.', label='CLUCB (w/o sim. reg.)')
    plt.plot(X, D, marker='P', markevery=1000, ls = '--', label='CLUCB')
    plt.ylabel("Cumulative Regret",fontweight='bold',fontsize = 16)
    plt.xlabel("Iteration",fontweight='bold',fontsize = 16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.legend(prop={'size': 12}, loc='upper left')
    plt.show()
    name = "regret" + str(n) + str(m)
    plt.savefig(f'/results/{name}.png')

def reg_plot_all(A, B, C, D, n_trial):
    plt.figure()
    X = np.arange(n_trial)
    plt.plot(X, A, marker='.', markevery=1000, ls = ':', label='LinUCB')
    plt.plot(X, B, marker='s', markevery=1000, ls = '-.', label='GOB.Lin')
    plt.plot(X, C, marker='p', markevery=1000, ls = '-.', label='CLUCB (w/o sim. reg.)')
    plt.plot(X, D, marker='P', markevery=1000, ls = '--', label='CLUCB')
    plt.ylabel("Cumulative Regret",fontweight='bold',fontsize = 16)
    plt.xlabel("Iteration",fontweight='bold',fontsize = 16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.legend(prop={'size': 12}, loc='upper left')
    plt.show()
    plt.savefig('/results/regret.png')

def reg_plot_all_conv(A, B, n_trial):
    plt.figure()
    X = np.arange(n_trial)
    plt.plot(X, A, marker='.', markevery=1000, ls = ':', label='conv. max iters.')
    plt.plot(X, B, marker='s', markevery=1000, ls = '-.', label='conv. param.')
    plt.ylabel("Cumulative Regret",fontweight='bold',fontsize = 16)
    plt.xlabel("Iteration",fontweight='bold',fontsize = 16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.legend(prop={'size': 12}, loc='upper left')
    plt.show()
    plt.savefig('/results/conv.png')
    
def l2_plot_all(A, B, C, D, n_trial):
    plt.figure()
    X = np.arange(n_trial)
    plt.plot(X, A, marker='.', markevery=1000, ls = ':', label='LinUCB')
    plt.plot(X, B, marker='s', markevery=1000, ls = '-.', label='GOB.Lin')
    plt.plot(X, C, marker='p', markevery=1000, ls = '-.', label='CLUCB (w/o sim. reg.)')
    plt.plot(X, D, marker='P', markevery=1000, ls = '--', label='CLUCB')
    plt.ylabel("L2 Difference",fontweight='bold',fontsize = 16)
    plt.xlabel("Iteration",fontweight='bold',fontsize = 16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.legend(prop={'size': 12}, loc='upper left')
    plt.show()
    plt.savefig('/results/l2diff.png')

def l2_plot_all_c(A, B, C, D, n_trial, n, m):
    plt.figure()
    X = np.arange(n_trial)
    plt.plot(X, A, marker='.', markevery=1000, ls = ':', label='LinUCB')
    plt.plot(X, B, marker='s', markevery=1000, ls = '-.', label='GOB.Lin')
    plt.plot(X, C, marker='p', markevery=1000, ls = '-.', label='CLUCB (w/o sim. reg.)')
    plt.plot(X, D, marker='P', markevery=1000, ls = '--', label='CLUCB')
    plt.ylabel("L2 Difference",fontweight='bold',fontsize = 16)
    plt.xlabel("Iteration",fontweight='bold',fontsize = 16)
    plt.xticks(fontsize=14)
    plt.yticks(fontsize=14)
    plt.legend(prop={'size': 12}, loc='upper left')
    plt.show()
    name2 = "l2diff" + str(n) + str(m)
    plt.savefig(f'/results/{name2}.png')
    
def plot_y_all(A, B, C, D, y):
    left, width = .2, .75
    bottom, height = .2, .75
    right = left + width
    top = bottom + height
    #if applying on other datasets, please use A,B,C,D instead of -1*(A-30)
    r_linucb = r2_score(y, -1*(A-30))
    r_goblin = r2_score(y, -1*(B-30))
    r_clucbwo = r2_score(y, -1*(C-30))
    r_clucb = r2_score(y, -1*(D-30))
    mse_linucb = mean_squared_error(y, -1*(A-30))
    mse_goblin = mean_squared_error(y, -1*(B-30))
    mse_clucbwo = mean_squared_error(y, -1*(C-30))
    mse_clucb = mean_squared_error(y, -1*(D-30))
    ind = np.arange(30)
    fig, axs = plt.subplots(1, 4, figsize=(5, 5))
    axs[0].scatter(-1*(A-30), y, color="#1f77b4")
    axs[0].text(left, top, f"$R^2$ = {r_linucb:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[0].transAxes, size=12)
    axs[0].text(left, top-0.05, f"MSE = {mse_linucb:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[0].transAxes, size=12)
    axs[1].scatter(-1*(B-30), y, color="#ff7f0e")
    axs[1].text(left, top, f"$R^2$ = {r_goblin:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[1].transAxes, size=12)
    axs[1].text(left, top-0.05, f"MSE = {mse_goblin:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[1].transAxes, size=12)
    axs[2].scatter(-1*(C-30), y, color="#2ca02c")
    axs[2].text(left, top, f"$R^2$ = {r_clucbwo:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[2].transAxes, size=12)
    axs[2].text(left, top-0.05, f"MSE = {mse_clucbwo:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[2].transAxes, size=12)
    axs[3].scatter(-1*(D-30), y, color="#d62728")
    axs[3].text(left, top, f"$R^2$ = {r_clucb:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[3].transAxes, size=12)
    axs[3].text(left, top-0.05, f"MSE = {mse_clucb:.4f}",horizontalalignment='center',
        verticalalignment='center', transform=axs[3].transAxes, size=12)
    for ax in axs:
        ax.plot(ind, ind, color='black')
        ax.set_xlabel("Predicted Y",fontweight='bold',fontsize = 16)
        ax.set_ylabel("True Y",fontweight='bold',fontsize = 16)
        ax.set_xticks(range(0, 35,5))
        ax.set_yticks(range(0, 35,5))
    axs[0].set_title("LinUCB", fontweight='bold',fontsize = 16)
    axs[1].set_title("GOB.Lin", fontweight='bold',fontsize = 16)
    axs[2].set_title("CLUCBwo", fontweight='bold',fontsize = 16)
    axs[3].set_title("CLUCB", fontweight='bold',fontsize = 16)
    # plt.tight_layout()
    plt.show()
    plt.savefig('/results/predy.png')
    
def make_regret(payoff, oracle):
    return np.cumsum(oracle - payoff)

# X transformation from a sparse matrix
def X_reshape(X, X_tr, t, K, n_arms, n_feature):  #
  for arm in range(1, n_arms):
    X_tr = np.concatenate((X_tr,np.kron(np.identity(n = K),X[arm].reshape(-1,1))), axis = 1)
  return X_tr

# convert to a sparse matrix
def X_to_X_m(X, t, arm_choice, n_arms, n_feature):
  X_m = np.copy(X[t])
  for arm in np.arange(n_arms): # N x p
    if arm not in arm_choice:
      X_m[arm] = np.zeros(shape=n_feature)
  return X_m

# Create the F matrix
def constructAdjMatrix(W, n, threshold): #m
    Adj_mat = np.zeros(shape = (n, n))
    for ui in range(n):
        for uj in range(n):
            Adj_mat[ui][uj] = W[ui][uj]
        # trim the graph
            for i in range(n):
                if W[ui][i] <= threshold:
                    Adj_mat[ui][i] = 0;
#         Adj_mat[ui] /= sum(Adj_mat[ui])
    return Adj_mat
def constructLaplacianMatrix(W, n, Gepsilon):
    G = W.copy()
    #Convert adjacency matrix of weighted graph to adjacency matrix of unweighted graph
    for i in range(n):
        for j in range(n):
            if G[i][j] > 0:
                G[i][j] = 1
    L = csgraph.laplacian(G, normed = False)
    I = np.identity(n = G.shape[0])
    GW = I + Gepsilon*L  # W is a double stochastic matrix
    return GW.T

def construct_X(N_TRIAL, N_ARMS, inc):
    if inc == 0:
        incre = 0.001 + inc
        X_1_lst = []
        for T in np.arange(N_TRIAL):
          X_1t_lst = []
          for arm in np.arange(N_ARMS):
            temp = []
            temp.append(1)
            temp.append(incre*(T+1))
            temp.append((incre*(T+1))**2)
            X_1t_lst.append(np.array(temp))
          X_1_lst.append(np.array(X_1t_lst))
        X_1 = np.array(X_1_lst)
    else:
        X_1_lst = []
        for T in np.arange(N_TRIAL):
          X_1t_lst = []
          for arm in np.arange(N_ARMS):
            temp = []
            temp.append(1)
            temp.append(0.001 + inc*(T))
            temp.append((0.001 + inc*(T))**2)
            X_1t_lst.append(np.array(temp))
          X_1_lst.append(np.array(X_1t_lst))
        X_1 = np.array(X_1_lst)
    return X_1
    
def construct_opt(Y, N_TRIAL, N_ARMS, M):
    oracle_lst = []
    true_choice = []
    new_y = Y #
    for t in np.arange(N_TRIAL):
      # Find indices of M highest arms
      all_reward_t = [new_y.T[t, arm] for arm in np.arange(N_ARMS)]
      chosen_arms = np.array(all_reward_t).argsort()[-M:][::-1]
      # Sum of M highest rewards
      oracle_payoff_t = np.sum([new_y.T[t, choice] for choice in chosen_arms])
      # Append to the list
      oracle_lst.append(oracle_payoff_t)
      true_choice.append(chosen_arms)
    oracle = np.array(oracle_lst)
    return true_choice, oracle