import torch
from torch import nn
from torch.nn import functional
import numpy as np
from cvxpy import *
import gurobipy as gp
from gurobipy import GRB
from gurobipy import quicksum
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import mean_squared_error
from sklearn.kernel_ridge import KernelRidge
import pytorch_forecasting


def opt_oracle(prob_vec, problem_params):
    dim_sol = problem_params['dim_sol']
    dim_sce = problem_params['dim_scenario']
    sce_list = problem_params['scenario_list']
    c_list = problem_params['c_list']
    problem_setting = problem_params['setting']
    flow = problem_params['flow']
    if problem_setting == 'network flow':
        w = Variable(dim_sol)
        constr = [ w >= 0]
        obj = 0
        for i in range(dim_sce):
            for j in range(dim_sol):
                obj += prob_vec[i]*c_list[j]*(w[j]-sce_list[i,j])**2
        constr += [w[0]+w[1] >= flow]
        constr += [w[2] + w[3] >= flow]
        constr += [w[0] + w[2] >= flow]
        constr += [w[1] + w[3] >= flow]
        # print("cvxpy objective", nv_obj)
        # print("input prob vec", prob_vec)
        prob = Problem(Minimize(obj), constr)
        prob.solve(solver=SCS)  # ECOS
        w = w.value
        if w is None:
            try:
                max_k = torch.argmax(prob_vec)
            except:
                max_k = np.argmax(prob_vec)
            w = sce_list[max_k, :]
        return np.reshape(np.array(w), dim_sol)
    else:
        print("ERROR! Invalid problem setting!")
        return None

### Gurobi version.
# def opt_oracle(prob_vec, problem_params):
#     dim_sol = problem_params['dim_sol']
#     dim_sce = problem_params['dim_scenario']
#     sce_list = problem_params['scenario_list']
#     b_list = problem_params['b_list']
#     h_list = problem_params['h_list']
#     problem_setting = problem_params['setting']
#     rho = problem_params['regularization_rho']
#     budget = problem_params['budget']
#     # if problem_setting == 'loigistic':
#     #     w = Variable(dim_sol)
#     #     constr = [sum(x) <= budget]
#     #     expr_obj = 0
#     #     for i in range(dim_sce):
#     #         z = sce_list[i]
#     #         expr_obj +=  prob_vec[i]*logistic(-z*w)
#     #     prob = Problem(Minimize(expr_obj), constr)
#     #     prob.solve(solver = ECOS)
#     #     w = w.value
#     #     return np.reshape(np.array(w), dim_sol)
#     if problem_setting == 'newsvendor':
#         w = Variable( dim_sol)
#         t1 = Variable((dim_sce, dim_sol))
#         t2 = Variable((dim_sce, dim_sol))
#         constr = [sum(w) <= budget, t1 >=0, t2>=0, w >=0]
#         nv_obj = 0
#        # prob_vec = [0,0,1,0]
#         for i in range(dim_sce):
#             for j in range(dim_sol):
#                 z = sce_list[i, j]
#                 nv_obj += prob_vec[i]*(h_list[j]*t1[i,j] + b_list[j]*t2[i,j])
#                 constr+= [t1[i,j] >= w[j]-z, t2[i,j] >= z-w[j]]
#         nv_obj += 0.5*rho*sum_squares(w)
#         #print("cvxpy objective", nv_obj)
#         #print("input prob vec", prob_vec)
#         prob = Problem(Minimize(nv_obj), constr)
#         prob.solve(solver=SCS) # ECOS
#         w = w.value
#         if w is None:
#             try:
#                 max_k = torch.argmax(prob_vec)
#             except:
#                 max_k = np.argmax(prob_vec)
#             w = sce_list[max_k,:]
#         # print("cvxpy_sol", w)
#         # print("input prob vec", prob_vec)
#         return np.reshape(np.array(w), dim_sol)
#     else:
#         print("ERROR! Invalid problem setting!")
#         return None


def fit_sol_map_KRR(p_train, w_train, p_test, w_test, problem_params):
    dim_x = problem_params['dim_scenario']
    dim_w = problem_params['dim_sol']
    oracle_kernel = problem_params['oracle_kernel']
    X = np.array(p_train).reshape(-1,dim_x)
    y = np.array(w_train).reshape(-1,dim_w)
    parameters = {'alpha': [1e-2, 1e-1, 1, 10],
                  'gamma': np.logspace(-3, 3, 10),
                  'coef0': [1e3, 1e2, 10, 0.1, 1e-2, 1e-2, 1e-3, 1e-4],
                  'degree': [3,4,5]}
    krr_cv = GridSearchCV(KernelRidge(kernel=oracle_kernel), parameters,  cv=5)
    krr_model1 = krr_cv.fit(X, y)
    oracle_params = krr_model1.best_params_
    X_test = np.array(p_test).reshape(-1,dim_x)
    w1_pred = krr_model1.predict(X_test)
    y_test = np.array(w_test).reshape(-1,dim_w)
    return oracle_params, mean_squared_error(y_test, w1_pred)


def dual_coef_KRR(p_train, w_train, p_test, w_test, oracle_params, problem_params):
    alpha = oracle_params["alpha"]
    gamma = oracle_params["gamma"]
    degree = oracle_params["degree"]
    coef0 = oracle_params["coef0"]
    dim_w = problem_params['dim_sol']
    oracle_kernel = problem_params['oracle_kernel']
    X = p_train
    y =np.array(w_train).reshape(-1, dim_w)
    krr = KernelRidge(kernel=oracle_kernel, alpha=alpha, gamma=gamma, coef0=coef0, degree=degree)
    model2 = krr.fit(X, y)
    w1_pred = model2.predict(p_test)
    y_test = np.array(w_test).reshape(-1, dim_w)
    mse = mean_squared_error(y_test, w1_pred)
    dual_coef = krr.dual_coef_
    # print("gamma", gamma)
    # print("coef0", coef0)
    # print("approx oracle mse", mse)
    # print("dual_coef size", np.shape(dual_coef))
    # print(w_train)
    return dual_coef, mse

def kernel_function(vector1, vector2, oracle_params, problem_params):
    if problem_params['oracle_kernel']=='polynomial':
        coef0 = oracle_params["coef0"]
        degree = oracle_params["degree"]
        gamma = oracle_params["gamma"]
        #print("oracle_params", oracle_params)
        return (gamma*torch.matmul(vector1, torch.t(vector2)) + coef0).pow(degree)
    else:
        print("ERROR!! Oracle kernel not valid!")
        return None

def kernel_function_grad(vector1, vector2, oracle_params, problem_params):
    if problem_params['oracle_kernel']=='polynomial':
        coef0 = oracle_params["coef0"]
        degree = oracle_params["degree"]
        gamma = oracle_params["gamma"]
        #print("vector1", vector1)
        #print("vector2", vector2)
        return degree*(gamma*torch.matmul(vector1, torch.t(vector2)) + coef0).pow(degree-1)*gamma*vector1
    else:
        print("ERROR!! Oracle kernel not valid!")
        return None

def approx_oracle(prob_vector, p_train, dual_coef, oracle_params, problem_params):
    # print(problem_params)
    # p_train = problem_params['p_train']
    # dual_coef = problem_params['dual_coef']
    #print(p_train.size())
    p_train = torch.as_tensor(p_train, dtype=torch.float64)
    n = list(p_train.size())[0]
    b_size = list(prob_vector.size())[0]
    K = problem_params['dim_sol']
    w_approx = torch.zeros((b_size, K), dtype=torch.float64)
    #print("prob_vec in approx oracle", prob_vector)
    for j in range(b_size):
        for i in range(n):
            kernel_out = kernel_function(p_train[i], prob_vector[j], oracle_params, problem_params)
            for k in range(K):
                # print("prob_vec size", prob_vector.size())
                #print("prob_vec[j]", prob_vector[j])
                #print("dual_coef", dual_coef[i,k])
                #print("kernel output", kernel_out)
                w_approx[j,k] += dual_coef[i,k]*kernel_out
    #print("w_approx in approx oracle", w_approx)
    return w_approx



def approx_oracle_back(grad, prob_vector,p_train, dual_coef, oracle_params, problem_params):
    # p_train = problem_params['p_train']
    # dual_coef = problem_params['dual_coef']
    p_train = torch.as_tensor(p_train, dtype=torch.float64)
    n = list(p_train.size())[0]
    b_size = list(prob_vector.size())[0]
    K = problem_params['dim_sol']
    dim_scenario = problem_params['dim_scenario']
    w_grad = torch.zeros((b_size, dim_scenario,K), dtype=torch.float64)
    for j in range(b_size):
        for i in range(n):
            #print("p_train[i]", p_train[i])
            #print("prob_vec[j]", prob_vector[j])
            kernel_grad_out = kernel_function_grad(p_train[i], prob_vector[j], oracle_params, problem_params)
            for k in range(K):
                # print("kernel_grad_out", kernel_grad_out)
                # print("dual_coef", dual_coef[i,k])
                w_grad[j,:,k] += dual_coef[i,k]*kernel_grad_out
    grad_l = grad['dw']
    #print("approx_oracle_back, w_grad size, grad l size", w_grad.size(), grad_l.size())
    grad_out = torch.einsum('bks, bs -> bk', w_grad, grad_l)
    # print("grad out size", grad_out.size())
    return grad_out

def fit_sol_map_NN(p_train, w_train, p_val, w_val, p_test, w_test, problem_params, learning_rate=0.01, n_ite= 100):
    dim_x = problem_params['dim_scenario']
    dim_w = problem_params['dim_sol']
    model = nn.Sequential(nn.Linear(dim_x, 128),
                          nn.Dropout(p=0.5), # 0.25 works
                          nn.ReLU(),
                          # nn.LeakyReLU(),
                          nn.Linear(128, 128),
                          nn.Dropout(p=0.5), # 0.25 works
                          nn.ReLU(),
                          # nn.LeakyReLU(),
                          nn.Linear(128, dim_w))
    loss_fn = pytorch_forecasting.metrics.MAPE()  #nn.MSELoss()
    #optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    optimizer =torch.optim.AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01, amsgrad=False)
    for t in range(n_ite):
        # load batch:
        # p_batch, w_batch = batch_loader(p_train, w_train, batch_size)
        optimizer.zero_grad()
        output = model(p_train.float())
        loss = loss_fn(output, w_train.float())
        loss.backward()
        optimizer.step()
        if t % 100 == 99:
            print(t,"training loss", loss.item())
            pred_val = model(p_val.float())
            loss_val = loss_fn(pred_val, w_val.float())
            print(t, "validation loss", loss_val.item())

    # test loss:

    pred = model(p_test.float())
    print("test set pred", pred)
    print("test set w", w_test)
    loss_test = loss_fn(pred, w_test.float())
    print("test mse", loss_test)
    return model, loss_test

def approx_oracle_nn(prob_vector, oracle_model, problem_params):
    w_approx = oracle_model(prob_vector.float())
    return w_approx

def approx_oracle_nn_back(grad, prob_vector, oracle_model, problem_params):
    # b_size = list(prob_vector.size())[0]
    # K = problem_params['dim_sol']
    # dim_scenario = problem_params['dim_scenario']
    prob_tensor = torch.tensor(prob_vector, requires_grad=True)
    out = oracle_model(prob_tensor.float())
    out.backward(gradient=grad['dw'])
    grad_out = prob_tensor.grad
    # print("gradient", grad_out)
    # print("gradient shape", grad_out.size())
    return grad_out






