# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/11_DSCM_cost.ipynb.

# %% auto 0
__all__ = ['get_total_cost', 'get_exp_cost', 'calculate_costs_per_q', 'empirical_cdf', 'get_marginal_cost',
           'get_exp_marginal_cost', 'calculate_marginal_costs_per_q']

# %% ../../nbs/11_DSCM_cost.ipynb 3
import numpy as np

def get_total_cost(y_diff, cu, co, return_total_cu_co = False, round_to_int = False):
### TODO: adapt the cost function in ddop along these lines
    
    if not isinstance(y_diff, np.ndarray):
        y_diff = np.array(y_diff)
    
    total_underage = np.sum(y_diff[y_diff<0]*cu*-1)
    total_overage = np.sum(y_diff[y_diff>0]*co)

    if round_to_int:
        total_underage = int(round(total_underage), 0)
        total_overage = int(round(total_overage), 0)

    total = total_overage + total_underage
    
    if return_total_cu_co:  
    
        return total, total_underage, total_overage
    else:
        return total
    
def get_exp_cost(y_true, y_pred, cu, co, return_total_cu_co = False, round_to_int = False):
    
    y_diff = y_pred - y_true
        
    total, total_underage, total_overage = get_total_cost(y_diff, cu, co, return_total_cu_co = True)
    
    total_avg = total/len(y_true)
    
    if return_total_cu_co:
        
        
        underage_avg = total_underage / len(y_true)
        overage_avg = total_overage / len(y_true)
        
        return total_avg, underage_avg, overage_avg
    
    else:
        return total_avg

def calculate_costs_per_q(y_true, cu, co, q_tried=None, upper_bound=100):
    
    if q_tried == None:
        q_values = np.arange(upper_bound + 1)  # Generate an array of integer values from 0 to the specified upper bound (inclusive)
    else: 
        q_values = q_tried
    
    total_costs = list()
    cu_costs = list()
    co_costs = list()

    for i, q in enumerate(q_values):
        total_avg, cu_avg, co_avg = get_exp_cost(y_true=y_true, y_pred=[q]*len(y_true), cu=cu, co=co, return_total_cu_co=True, round_to_int=True)
        total_costs.append(total_avg)
        cu_costs.append(cu_avg)
        co_costs.append(co_avg)

    return total_costs, cu_costs, co_costs


def empirical_cdf(y_true, q):

    count = np.sum(y_true < q)
    cdf = count / len(y_true)
    
    return cdf


def get_marginal_cost(y_true, q, cu, co, return_total_cu_co = False, round_to_int = False):
    ### TODO: adapt the cost function in ddop along these lines
    
    if not isinstance(y_true, np.ndarray):
        y_true = np.array(y_true)

    quantile = empirical_cdf(y_true, q)
    
    marginal_underage = cu*(1-quantile)
    marginal_overage = co*quantile

    if round_to_int:  
        marginal_underage = int(round(marginal_underage, 0))
        marginal_overage = int(round(marginal_overage, 0))
    
    return marginal_underage, marginal_overage


def get_exp_marginal_cost(y_true, q, cu, co, return_total_cu_co = False, round_to_int = False):
        
    total_underage, total_overage = get_marginal_cost(y_true, q, cu, co, return_total_cu_co = True, round_to_int = round_to_int)
    


def calculate_marginal_costs_per_q(y_true, cu, co, q_tried=None, upper_bound=100):
    
    if q_tried == None:
        q_values = np.arange(upper_bound + 1)  # Generate an array of integer values from 0 to the specified upper bound (inclusive)
    else: 
        q_values = q_tried
    
    total_marginal_costs = list()
    marginal_cu_costs = list()
    marginal_co_costs = list()

    for i, q in enumerate(q_values):
        cu_avg, co_avg = get_exp_marginal_cost(y_true=y_true, q=q, cu=cu, co=co, return_total_cu_co=True, round_to_int=False)
        marginal_cu_costs.append(cu_avg)
        marginal_co_costs.append(co_avg)

    return marginal_cu_costs, marginal_co_costs




# %% ../../nbs/11_DSCM_cost.ipynb 4
#| export
