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

def linUCB(lda, alpha, X, Y, m, oracle, true_choice):
    print("linUCB with alpha:", alpha)
    n_trial, n_arms, n_feature = X.shape
    # 1.Initialization
    c_payoff = np.empty(n_trial)
    c_payoff[0] = 0
    r_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    r_payoff[0] = 0
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    theta = np.empty(shape=(n_trial, n_arms, n_feature)) 
    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)) 
    opt_y = np.empty(n_trial)
    selected_pred = np.empty(n_trial)
    diff = np.empty(n_trial)
    
    ucb = np.empty(shape=(n_trial, n_arms))
    # 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)
    
    mse_sel = np.empty(n_trial)
    mse_opt = np.empty(n_trial)
    
    A = np.array([lda * np.diag(np.ones(shape=n_feature)) for _ in np.arange(n_arms)])
    b = np.array([np.zeros(shape=n_feature) for _ in np.arange(n_arms)])

    # 2. Algorithm
    for t in np.arange(n_trial): 
        for a in np.arange(n_arms):
            inv_A       = np.linalg.inv(A[a])
            theta[t, a] = inv_A.dot(b[a])
            pred_reward[t, a] = theta[t, a].dot(X[t, a])
            ucb[t, a] = alpha * np.sqrt(X[t, a].dot(inv_A).dot(X[t, a]))
            p[t, a] = pred_reward[t, a] + ucb[t, a]
            acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        
        # if t+1 == n_trial:
        #     for a in np.arange(n_arms):
        #         print("arm", a, ":", pred_reward[t, a])
        # Choose m best arms
        idx = np.argpartition(p[t], -m)[-m:]
        chosen_arms = idx[np.argsort(-(p[t])[idx])]

        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] = np.abs(selected_pred[t] - opt_y[t])

        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
        
        # Update A and b
        for chosen_arm in arm_choice[t]:
          A[chosen_arm] += np.outer(X[t, chosen_arm], X[t, chosen_arm].T)
          b[chosen_arm] += Y[t, chosen_arm] * X[t, chosen_arm]
        
        # 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[t] = sum(np.square(pred_reward[t, a] - Y[t, a]) for a in chosen_arms)/m
        mse_opt[t] = sum(np.square(pred_reward[t, a] - Y[t, a]) for a in true_choice[t])/m

        if (t+1) % 1000 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])

    return dict(mse_sel=mse_sel, mse_opt=mse_opt, theta=theta, p=p, arm_choice = arm_choice, r_payoff=r_payoff, c_payoff=c_payoff, A=A, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff,cum_regret=cum_regret)

def linUCB_conv_y(lda, alpha, X, Y, m, oracle, true_choice):
    print("linUCB with alpha:", alpha)
    n_trial, n_arms, n_feature = X.shape
    # 1.Initialization
    c_payoff = np.empty(n_trial)
    c_payoff[0] = 0
    r_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    r_payoff[0] = 0
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    theta = np.empty(shape=(n_trial, n_arms, n_feature)) 
    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)) 
    opt_y = np.empty(n_trial)
    selected_pred = np.empty(n_trial)
    diff = np.empty(n_trial)
    
    ucb = np.empty(shape=(n_trial, n_arms))
    # 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)
    
    mse_sel = np.empty(n_trial)
    mse_opt = np.empty(n_trial)
    
    A = np.array([lda * np.diag(np.ones(shape=n_feature)) for _ in np.arange(n_arms)])
    b = np.array([np.zeros(shape=n_feature) for _ in np.arange(n_arms)])

    # 2. Algorithm
    for t in np.arange(n_trial): 
        for a in np.arange(n_arms):
            inv_A       = np.linalg.inv(A[a])
            theta[t, a] = inv_A.dot(b[a])
            pred_reward[t, a] = theta[t, a].dot(X[t, a])
            ucb[t, a] = alpha * np.sqrt(X[t, a].dot(inv_A).dot(X[t, a]))
            p[t, a] = pred_reward[t, a] + ucb[t, a]
            acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        
        if t+1 == n_trial:
            for a in np.arange(n_arms):
                print("arm", a, ":", pred_reward[t, a])
        # Choose m best arms
        idx = np.argpartition(p[t], -m)[-m:]
        chosen_arms = idx[np.argsort(-(p[t])[idx])]

        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] = np.abs(selected_pred[t] - opt_y[t])

        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
        
        # Update A and b
        for chosen_arm in arm_choice[t]:
          A[chosen_arm] += np.outer(X[t, chosen_arm], X[t, chosen_arm].T)
          b[chosen_arm] += Y[t, chosen_arm] * X[t, chosen_arm]
        
        # 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[t] = sum(np.square(pred_reward[t, a] - Y[t, a]) for a in chosen_arms)/m
        mse_opt[t] = sum(np.square(pred_reward[t, a] - Y[t, a]) for a in true_choice[t])/m

        if (t+1) % 1000 == 0:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])
        if t+1 == n_trial:
            print('TRIAL:',t,'DONE', '| cum_regret:', cum_regret[t])

    return dict(mse_sel=mse_sel, mse_opt=mse_opt, theta=theta, p=p, arm_choice = arm_choice, r_payoff=r_payoff, c_payoff=c_payoff, A=A, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff,cum_regret=cum_regret)