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 GOBLin(alpha, eta, X, Y, m, FInv, oracle, true_choice, bb):
    print("GOB.Lin with alpha: ", alpha)
    n_trial, n_arms, n_feature = X.shape
    
    r_payoff = np.empty(n_trial)
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    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)
    # 1.2. A and b

    A = eta * np.identity(n = n_feature * n_arms)   #dn x dn
    b = np.zeros(shape=n_feature * n_arms)        #dn x 1
    inv_A = np.linalg.inv(A)
    theta_lst = np.empty(shape=(n_trial, n_arms *  n_feature))
    theta = np.dot(inv_A, b)
    
    for t in np.arange(n_trial):
        
        #Compute estimates (theta) and prediction (p) for all arms
        for a in np.arange(n_arms):
            X_1_tr = FInv.dot((ut.X_to_X_m(X, t, [a], n_arms, n_feature)).flatten())
            p[t, a] = theta.dot(X_1_tr) + alpha * np.sqrt(np.dot(np.dot(X_1_tr, inv_A), X_1_tr)) #.clip(0))
            pred_reward[t, a] = theta.dot(X_1_tr)
            acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        # Choose m best arms
        idx = np.argpartition(p[t], -m)[-m:]
        chosen_arms = idx[np.argsort(-(p[t])[idx])]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
        
        # if t+1 == n_trial:
        #     for a in np.arange(n_arms):
        #         print("arm", a, ":", pred_reward[t, a])
        
        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])
        
        #update A and b
        for chosen_arm in arm_choice[t]:
          X_1_tr_chosen = FInv.dot((ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)).flatten())
          A = A + np.outer(X_1_tr_chosen, X_1_tr_chosen)
          b = b + Y[t, chosen_arm] * X_1_tr_chosen
        inv_A = np.linalg.inv(A)
        theta = inv_A.dot(b)
        theta_lst[t] = theta
        
        # 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])
        if cum_regret[t] > bb:
            break
        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(theta=theta, theta_lst=theta_lst, p=p, arm_choice = arm_choice, r_payoff=r_payoff, A=A, c_payoff=c_payoff, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff,cum_regret =cum_regret)

def GOBLin_conv_y(alpha, eta, X, Y, m, FInv, oracle, true_choice, bb):
    print("GOB.Lin with alpha: ", alpha)
    n_trial, n_arms, n_feature = X.shape
    
    r_payoff = np.empty(n_trial)
    c_payoff = np.empty(n_trial)
    cum_regret = np.empty(n_trial)
    arm_choice = np.empty(shape=(n_trial, m), dtype=int)
    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)
    # 1.2. A and b

    A = eta * np.identity(n = n_feature * n_arms)   #dn x dn
    b = np.zeros(shape=n_feature * n_arms)        #dn x 1
    inv_A = np.linalg.inv(A)
    theta_lst = np.empty(shape=(n_trial, n_arms *  n_feature))
    theta = np.dot(inv_A, b)
    
    for t in np.arange(n_trial):
        
        #Compute estimates (theta) and prediction (p) for all arms
        for a in np.arange(n_arms):
            X_1_tr = FInv.dot((ut.X_to_X_m(X, t, [a], n_arms, n_feature)).flatten())
            p[t, a] = theta.dot(X_1_tr) + alpha * np.sqrt(np.dot(np.dot(X_1_tr, inv_A), X_1_tr)) #.clip(0))
            pred_reward[t, a] = theta.dot(X_1_tr)
            acc[t, a]   = abs(pred_reward[t, a] - Y[t, a])
        # Choose m best arms
        idx = np.argpartition(p[t], -m)[-m:]
        chosen_arms = idx[np.argsort(-(p[t])[idx])]
        for i in np.arange(m):
          arm_choice[t][i] = chosen_arms[i]
        
        if t+1 == n_trial:
            for a in np.arange(n_arms):
                print("arm", a, ":", pred_reward[t, a])
        
        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])
        
        #update A and b
        for chosen_arm in arm_choice[t]:
          X_1_tr_chosen = FInv.dot((ut.X_to_X_m(X, t, [chosen_arm], n_arms, n_feature)).flatten())
          A = A + np.outer(X_1_tr_chosen, X_1_tr_chosen)
          b = b + Y[t, chosen_arm] * X_1_tr_chosen
        inv_A = np.linalg.inv(A)
        theta = inv_A.dot(b)
        theta_lst[t] = theta
        
        # 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])
        if cum_regret[t] > bb:
            break
        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(theta=theta, theta_lst=theta_lst, p=p, arm_choice = arm_choice, r_payoff=r_payoff, A=A, c_payoff=c_payoff, acc=acc,selected_pred=selected_pred, pred_reward=pred_reward, opt_y=opt_y, diff=diff,cum_regret =cum_regret)

