# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/metrics.ipynb (unless otherwise specified).

__all__ = ['peace', 'ece', 'ece_v2', 'class_wise_error']

# Cell

from riverreliability import utils

import numpy as np
import scipy.stats
import scipy.integrate

import riverreliability.beta
import sklearn.metrics
import sklearn.datasets
import sklearn.model_selection
import sklearn.svm
import sklearn.preprocessing
import sklearn.utils

# Cell

def peace(y_probs, y_preds, y_true, samples=1000, bins="fd", **bin_args):
    """Compute the posterior expected balanced accuracy-based calibration error (PEACE).

    Parameters:
    y_probs (np.array): predicted class probabilities
    y_preds (np.array): predicted class labels
    y_true (np.array): true class labels
    samples (int): number of samples for numerical integration

    Returns:
    exp_ce (float): expected calibration error

    """

    sklearn.utils.check_consistent_length(y_preds, y_true)
    classes = sklearn.utils.multiclass.unique_labels(y_preds, y_true)

    # define the bin function
    def bin_func(y_probs_bin, y_preds_bin, y_true_bin):

        xs = np.linspace(0, 1, samples)
        conf = y_probs_bin.mean()

        if len(np.unique(y_preds_bin)) > 1:
            # estimate beta parameters
            confusion = sklearn.metrics.confusion_matrix(y_true_bin, y_preds_bin, labels=classes)
            params = riverreliability.beta.get_beta_parameters(confusion)
            ys = abs(xs - conf) * riverreliability.beta.beta_avg_pdf(xs, params, fft=True)
        else:
            params = sum(y_preds_bin == y_true_bin)+1, sum(y_preds_bin != y_true_bin)+1
            ys = abs(xs - conf) * scipy.stats.beta.pdf(xs, params[0], params[1])

        # approximate the integral using Simpson's rule
        return scipy.integrate.simps(ys, xs)

    # compute the full result
    bin_indices = utils.get_bin_indices(y_probs, bins=bins, lower=0, upper=1, **bin_args)
    return utils.binning(y_probs, y_preds, y_true, bin_indices, bin_func)

# Cell

def ece(y_probs, y_preds, y_true, balanced=False, bins="fd", **bin_args):
    """Compute the expected calibration error (ECE).

    Parameters:
    y_probs (np.array): predicted class probabilities
    y_preds (np.array): predicted class labels
    y_true (np.array): true class labels

    Returns:
    exp_ce (float): expected calibration error

    """

    sklearn.utils.check_consistent_length(y_preds, y_true)

    # define the bin function
    def bin_func(y_probs_bin, y_preds_bin, y_true_bin):
        acc = (y_preds_bin == y_true_bin).mean()
        conf = y_probs_bin.mean()
        return abs(acc - conf)

    # define the balanced bin function
    def balanced_bin_func(y_probs_bin, y_preds_bin, y_true_bin):
        balacc = sklearn.metrics.balanced_accuracy_score(y_true_bin, y_preds_bin)
        conf = y_probs_bin.mean()
        return abs(balacc - conf)

    # compute the full result
    bin_indices = utils.get_bin_indices(y_probs, bins=bins, lower=0, upper=1, **bin_args)
    func = balanced_bin_func if balanced else bin_func
    return utils.binning(y_probs, y_preds, y_true, bin_indices, func)

# Cell

def ece_v2(y_probs, y_preds, y_true, bins="fd", **bin_args):
    """Compute the expected calibration error based on the expected posterior balanced accuracy (ECEv2).

    Parameters:
    y_probs (np.array): predicted class probabilities
    y_preds (np.array): predicted class labels
    y_true (np.array): true class labels

    Returns:
    exp_ce (float): expected calibration error

    """

    sklearn.utils.check_consistent_length(y_preds, y_true)
    classes = sklearn.utils.multiclass.unique_labels(y_preds, y_true)

    # define the bin function
    def bin_func(y_probs_bin, y_preds_bin, y_true_bin):
        confusion = sklearn.metrics.confusion_matrix(y_true_bin, y_preds_bin, labels=classes)
        acc = riverreliability.beta.balanced_accuracy_expected(confusion, fft=True)
        conf = y_probs_bin.mean()
        return abs(acc - conf)

    # compute the full result
    bin_indices = utils.get_bin_indices(y_probs, bins=bins, lower=0, upper=1, **bin_args)
    return utils.binning(y_probs, y_preds, y_true, bin_indices, bin_func)

# Cell

def class_wise_error(y_probs, y_preds, y_true, base_error, *base_error_args, **base_error_kwargs):
    """Compute classwise-error as proposed in "Beyond temperature scaling: Obtaining well-calibrated
    multiclass probabilities with Dirichlet calibration" (Kull, 2019).

    Parameters:
    y_probs (np.array): predicted class probabilities
    y_preds (np.array): predicted class labels
    y_true (np.array): true class labels
    base_error (callable): function that returns ECE for given probabilities, label predictions and true labels
    base_error_[kw]args ([kw]args): [Keyword ]arguments that should be passed to the base_ece callable.

    Returns:
    exp_ce (float): class-wise expected calibration error

    """

    if hasattr(y_preds, "to_numpy"):
        y_preds = y_preds.to_numpy()
    if hasattr(y_preds, "to_numpy"):
        y_true = y_true.to_numpy()

    sklearn.utils.check_consistent_length(y_preds, y_true)
    classes = sklearn.utils.multiclass.unique_labels(y_preds, y_true)

    result = 0.
    for i, c in enumerate(classes):
        selector = y_preds == c
        if sum(selector) == 0:
            continue

        result += base_error(y_probs[selector, i], y_preds[selector], y_true[selector], *base_error_args, **base_error_kwargs)

    return result/len(classes)