import numpy as np

def calculate_confusion_matrix_value_result(outcome_pred, outcome_true):

    outcome_pred = 1 if outcome_pred > 0.5 else 0

    if outcome_pred == 1 and outcome_true == 1:

        return "tp"

    elif outcome_pred == 0 and outcome_true == 0:

        return "tn"

    elif outcome_pred == 1 and outcome_true == 0:

        return "fp"

    elif outcome_pred == 0 and outcome_true == 1:

        return "fn"

    else:

        raise ValueError("Unknown value occurred")

def calculate_es(los_true, threshold, case="tp"):

    metric = 0.0

    if case == "tp":

        if los_true >= threshold:  # predict correct in early stage

            metric = 1

        else:

            metric = los_true / threshold

    elif case == "fn":

        if los_true >= threshold:  # predict wrong in early stage

            metric = 0

        else:

            metric = los_true / threshold - 1

    elif case == "tn":

        metric = 0.0

    elif case == "fp":

        metric = -0.1 # penalty term

    return metric

def es_score(

    y_true_outcome,

    y_true_los,

    y_pred_outcome,

    threshold,

    verbose=0

):

    """

    Args:

        - threshold: 50%*mean_los (patient-wise) 

    Note:

        - y/predictions are already flattened here

        - so we don't need to consider visits_length

    """

    metric = []

    metric_optimal = []

    num_records = len(y_pred_outcome)

    for i in range(num_records):

        cur_outcome_pred = y_pred_outcome[i]

        cur_outcome_true = y_true_outcome[i]

        cur_los_true = y_true_los[i]

        prediction_result = calculate_confusion_matrix_value_result(cur_outcome_pred, cur_outcome_true)

        prediction_result_optimal = calculate_confusion_matrix_value_result(cur_outcome_true, cur_outcome_true)

        metric.append(

            calculate_es(

                cur_los_true,

                threshold,

                case=prediction_result,

            )

        )

        metric_optimal.append(

            calculate_es(

                cur_los_true,

                threshold,

                case=prediction_result_optimal,

            )

        )

    metric = np.array(metric)

    metric_optimal = np.array(metric_optimal)

    result = 0.0

    if metric_optimal.sum() > 0.0:

        result = metric.sum() / metric_optimal.sum()

    result = max(result, -1.0)

    if verbose:

        print("ES Score:", result)

    if isinstance(result, np.float64):

        result = result.item()

    return {"es": result}

if __name__ == "__main__":

    y_true_outcome = np.array([0,1])

    y_true_los = np.array([5,5])

    y_pred_outcome = np.array([0.7,0.7])

    y_pred_los = np.array([10,10])

    large_los = 110

    threshold = 10

    print(es_score(

        y_true_outcome,

        y_true_los,

        y_pred_outcome,

        threshold,

        verbose=0

    ))