import numpy as np
from sklearn.metrics.pairwise import pairwise_kernels
from collections import Counter
from tools import optimization_oracle_tools
from tools import loss_func_tools
from tools import sto_forest_tools
from tools import prescriptive_rf_tools
from tools import cvxpy_layer_tools
from tools import e2e_tools

def get_knn_index(x, x_train,knn_num):
    dist_list = np.array([np.linalg.norm(x_i - x) for x_i in x_train])
    idx = np.argpartition(dist_list,knn_num)
   # print(idx)
    return idx[:knn_num]

def prespective_knn_batch(params, x_train, z_train, x_test, knn_num=10):
    K = params['dim_scenario']
    dim_sol = params['dim_sol']
    z_tilde = params['scenario_list']
    n_test = np.shape(x_test)[0]
    w_knn_batch = np.zeros((n_test, dim_sol))
    for i in range(n_test):
        count = np.zeros(K)
        x = x_test[i]
        idx = get_knn_index(x, x_train, knn_num)
        #x_train_temp = x_train[idx]
        z_train_temp = z_train[idx,:]
        #  print("in knn idx", idx)
        # print("in knn z_train_temp", z_train_temp)
        if np.shape(z_train_temp)[0] != knn_num:
            raise ValueError("knn shape not consistent!")
        for j in range(knn_num):
            # print('i',i)
            for k in range(K):
                # print('k',k)
                # print(z_tilde[k], z_train_temp[j])
                if np.array_equiv(z_tilde[k], z_train_temp[j]):
                    count[k] += 1
        if np.sum(count) != knn_num:
            print(count)
            print(knn_num)
            raise Exception('SAA count error!')
        prob_vec = count / np.sum(count)
        # print(count)
        w_knn_batch[i] = optimization_oracle_tools.opt_oracle(prob_vec, params)
    #print("knn batch", w_knn_batch)
    return w_knn_batch




# Gaussian kernel:
def get_kernel_weight(x, x_train, h):
    weight = np.array([np.exp(-(np.linalg.norm(x_i - x)/h)**2) for x_i in x_train])
    return weight
def prespective_kernel_batch(params, x_train, z_train, x_test, h=10):
    K = params['dim_scenario']
    dim_sol = params['dim_sol']
    z_tilde = params['scenario_list']
    n_test = np.shape(x_test)[0]
    n_train = np.shape(z_train)[0]
    w_kernel_batch = np.zeros((n_test, dim_sol))
    for i in range(n_test):
        count = np.zeros(K)
        x = x_test[i]
        
        weight = get_kernel_weight(x, x_train, h)
        # print('weight', weight)
        for j in range(n_train):
            for k in range(K):
                if np.array_equiv(z_tilde[k], z_train[j]):
                    # print('x',x)
                    # print('x_train',  x_train[j])
                    count[k] += weight[j]
                    # print('weight', weight[j])
        if not np.allclose([np.sum(count)], [np.sum(weight)]):
            print(np.sum(count))
            print(np.sum(weight))
            raise Exception('kernel count error!')
        prob_vec = count / np.sum(count)
        if np.sum(count) == 0:
            prob_vec = (count+1)/np.sum(count+1)
        # print(count)
        w_kernel_batch[i] = optimization_oracle_tools.opt_oracle(prob_vec, params)
    #print("knn batch", w_knn_batch)
    return w_kernel_batch


def saa_batch(params, x_train, z_train, x_test):
    K = params['dim_scenario']
    dim_sol = params['dim_sol']
    count = np.zeros(K)
    z_tilde = params['scenario_list']
    n_train = np.shape(z_train)[0]
    n_test = np.shape(x_test)[0]
    for i in range(n_train):
        for k in range(K):
            # if (z_tilde[k]==z_train[i]).all():
            if np.array_equiv(z_tilde[k], z_train[i]):
                #print("k", k)
                #print(z_tilde[k], z_train[i])
                count[k] +=1
    if np.sum(count) != n_train:
        print(count)
        print(n_train)
        raise Exception(
            'SAA count error!')
    prob_vec = count/np.sum(count)
    w_saa_batch = np.zeros((n_test, dim_sol))
    for i in range(n_test):
        w_saa_batch[i] = optimization_oracle_tools.opt_oracle(prob_vec, params)
    #print("saa batch", w_saa_batch)
    return w_saa_batch



def get_baseline_loss(params, x_train, z_train, x_test, z_test, x_valid, z_valid):
    baseline_list = params['baseline']
    baseline_test_loss = []
    for method in baseline_list:
        if method == 'saa':
            w_batch = saa_batch(params, x_train, z_train, x_test)
            loss = loss_func_tools.nf_loss_func_np(params, w_batch, z_test)
            print(method, loss)
            baseline_test_loss.append(loss)
        elif method == 'knn':
            knn_list = params['knn_grid']
            knn_loss = 10000
            knn_num_best = knn_list[0]
            for knn_num in knn_list:
                w_valid_batch = prespective_knn_batch(params, x_train, z_train, x_valid, knn_num)
                loss_valid = loss_func_tools.nf_loss_func_np(params, w_valid_batch, z_valid)
                print(method, knn_num, loss_valid)
                if loss_valid <= knn_loss:
                    knn_loss = loss_valid
                    knn_num_best = knn_num
            w_batch = prespective_knn_batch(params, x_train, z_train, x_test, knn_num_best)
            test_loss = loss_func_tools.nf_loss_func_np(params, w_batch, z_test)
            baseline_test_loss.append(test_loss)
        elif method == 'kernel':
            kernel_list = params['kernel_grid']
            kernel_loss = 10000
            h_best = kernel_list[0]
            for h in kernel_list:
                w_valid_batch = prespective_kernel_batch(params, x_train, z_train, x_valid, h)
                loss_valid = loss_func_tools.nf_loss_func_np(params, w_valid_batch, z_valid)
                print(method, h, loss_valid)
                if loss_valid <= kernel_loss:
                    kernel_loss = loss_valid
                    h_best = h
            w_batch = prespective_kernel_batch(params, x_train, z_train, x_test, h_best)
            test_loss = loss_func_tools.nf_loss_func_np(params, w_batch, z_test)
            baseline_test_loss.append(test_loss)
        elif method == 'forest':
            min_leaf_list = params['rf_min_leaf_grid']
            rf_loss = 10000
            for min_leaf in min_leaf_list:
                w_batch = prescriptive_rf_tools.prescriptive_rf_batch(params, x_train, z_train, x_test, z_test, min_leaf)
                loss = loss_func_tools.nf_loss_func_np(params, w_batch, z_test)
                print('forest min leaf size', min_leaf, loss)
                if loss <= rf_loss:
                    rf_loss = loss
            baseline_test_loss.append(rf_loss)
        elif method == 'cvxpy':
            cvxpy_train_loss, net = cvxpy_layer_tools.train(params, 300, x_train, z_train)
            cvxpy_loss = cvxpy_layer_tools.cvxpy_layer_test(params, net, x_test, z_test)
            baseline_test_loss.append(cvxpy_loss.item())
        elif method == 'e2e':
            _,net, layers = e2e_tools.train(params, 300, x_train, z_train)
            print(net)
            print(layers)
            e2e_loss = e2e_tools.cvxpy_layer_test(params, net, layers, x_test, z_test)
            baseline_test_loss.append(e2e_loss.item())

    return baseline_test_loss



