from methods.wachter.gd import gd
from methods.reup import bayesian_utils

def generate_recourse(x0, model, random_state, params=dict()):
    dim = x0.shape[0]

    # General parameters
    cat_indices = params['cat_indices']

    # Wachter parameters
    lr = params['wachter_params']['lr']
    lmbda = params['wachter_params']['lmbda']
    
    #Recourse generation
    recourse, feasible = gd(model,
    x0.reshape(1, -1),
    cat_feature_indices=cat_indices,
    binary_cat_features=True,
    lr=lr,
    lambda_param=lmbda,
    n_iter=1000,
    t_max_min=1000,
    clamp=True)

    #define the cost based on the user preference
    A_0 = bayesian_utils.generate_A_0(dim)
    cost = bayesian_utils.evaluate_cost_diag(recourse.reshape(-1, 1), x0.reshape(-1, 1), A_0) 
    #cost = (recourse.reshape(-1, 1) - x0.reshape(-1, 1)).T @ A_0 @ (recourse.reshape(-1, 1) - x0.reshape(-1, 1))

    print(cost)
    print()

    #logging
    log_dict = {}
    log_dict['recourse'] = recourse
    log_dict['cost'] = cost
    log_dict['x_0'] = x0.reshape(-1, 1)
    log_dict['A_0'] = A_0

    return recourse, cost, feasible, log_dict