# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_tune.ipynb.

# %% auto 0
__all__ = ['calculate_metrics_by_thresh_binary', 'calculate_metrics_by_thresh_multi', 'coverage', 'calculate_metric_ignoring_nan',
           'fbeta', 'confusion_matrix']

# %% ../nbs/03_tune.ipynb 2
from collections import defaultdict
from functools import partial
from typing import Callable, Optional, Sequence, Union

import numpy as np
import pandas as pd
from sklearn import metrics
from tqdm import tqdm

# %% ../nbs/03_tune.ipynb 4
def _calculate_metrics_by_thresh(
    y_true: np.array,
    y_prob: np.array,
    metrics: Union[Callable, Sequence[Callable]],
    prob_to_pred: Callable,
    thresh_picker: Callable,
) -> pd.DataFrame:
    if callable(metrics):
        metrics = [metrics]

    results = defaultdict(list)
    results["thresh"] = thresh_picker(y_prob)
    for thresh in tqdm(results["thresh"]):
        y_pred = prob_to_pred(y_prob, thresh)
        for metric in metrics:
            results[metric.__name__].append(metric(y_true, y_pred))

    return pd.DataFrame(results)

# %% ../nbs/03_tune.ipynb 5
def calculate_metrics_by_thresh_binary(
    y_true: np.array,
    y_prob: np.array,
    metrics: Union[Callable, Sequence[Callable]],
    thresholds: Optional[Sequence] = None,
) -> pd.DataFrame:
    """Calculate binary classification metrics as a function of
    threshold

    Takes prediction to be `1` when `y_prob` is greater than the
    threshold, `0` otherwise.

    Parameters:

    - `y_true`: Ground-truth values with shape (n_items,)
    - `y_prob`: Probability distributions with shape (n_items, 2)
    - `metrics`: Callables that take `y_true`, `y_pred` as positional
    arguments and return a number. Must have a `__name__` attribute.
    - `thresholds`: `Sequence` of `float` threshold values to use. By
    default uses `0` and the values that appear in `y_prob[:, 1]`, which
    is a minimal set that covers all of the relevant possibilities. One
    reason to override that default would be to save time with a large
    dataset.

    Returns: DataFrame with one column "thresh" indicating the
    thresholds used and an additional column for each input metric
    giving the value of that metric at that threshold.
    """
    thresh_picker = (
        (lambda y_prob: thresholds)
        if thresholds is not None
        else (lambda y_prob: sorted(np.hstack([0, np.unique(y_prob[:, 1])])))
    )

    return _calculate_metrics_by_thresh(
        y_true=y_true,
        y_prob=y_prob,
        metrics=metrics,
        prob_to_pred=lambda y_prob, thresh: np.where(y_prob[:, 1] > thresh, 1, 0),
        thresh_picker=thresh_picker,
    )

# %% ../nbs/03_tune.ipynb 10
def calculate_metrics_by_thresh_multi(
    y_true: np.array,
    y_prob: np.array,
    metrics: Union[Callable, Sequence[Callable]],
    thresholds: Optional[Sequence] = None,
) -> pd.DataFrame:
    """Calculate multiclass metrics as a function of threshold

    Takes prediction to be the position of the column in `y_prob` with
    the greatest value if that value is greater than the threshold,
    `np.nan` otherwise.

    Parameters:

    - `y_true`: Ground-truth values
    - `y_prob`: Probability distributions
    - `metrics`: Callables that take `y_true`, `y_pred` as positional arguments
    and return a number. Must have a `__name__` attribute.
    - `thresholds`: `Sequence` of `float` threshold values to use. By
    default uses 0 and all values that appear in `y_prob`, which is a
    minimal set that covers all of the relevant possibilities. One
    reason to override that default would be to save time with a large
    dataset.

    Returns: DataFrame with one column "thresh" indicating the
    thresholds used and an additional column for each input metric
    giving the value of that metric at that threshold.
    """
    thresh_picker = (
        (lambda y_prob: thresholds)
        if thresholds is not None
        else (lambda y_prob: sorted(np.hstack([0, np.unique(y_prob)])))
    )
    return _calculate_metrics_by_thresh(
        y_true=y_true,
        y_prob=y_prob,
        metrics=metrics,
        prob_to_pred=lambda y_prob, thresh: np.where(
            y_prob.max(axis=1) > thresh, y_prob.argmax(axis=1), np.nan
        ),
        thresh_picker=thresh_picker,
    )

# %% ../nbs/03_tune.ipynb 12
def coverage(y_true: np.array, y_pred: np.array):
    """How often the model makes a prediction, where `np.nan` indicates
    abstaining from predicting.

    Parameters:

    - `y_true`: Ground-truth values
    - `y_pred`: Predicted values, possibly including `np.nan` to
    indicate abstraining from predicting
    """
    return (~np.isnan(y_pred)).mean()

# %% ../nbs/03_tune.ipynb 13
def calculate_metric_ignoring_nan(
    y_true: np.array, y_pred: np.array, metric: Callable, *args, **kwargs
):
    """Calculate `metric` ignoring `np.nan` predictions

    Parameters:

    - `y_true`: Ground-truth values
    - `y_pred`: Predicted values, possibly including `np.nan` to
    indicate abstraining from predicting
    - `metric`: Function that takes `y_true`, `y_pred` as keyword
    arguments

    Any additional arguments will be passed to `metric`
    """
    return metric(
        y_true=y_true[~np.isnan(y_pred)],
        y_pred=y_pred[~np.isnan(y_pred)],
        *args,
        **kwargs,
    )

# %% ../nbs/03_tune.ipynb 14
def fbeta(precision: float, recall: float, beta: float = 1):
    weighted_mean_of_inverses = (
        1 / (1 + beta**2) * (1 / precision + beta**2 * 1 / recall)
    )
    weighted_harmonic_mean = 1 / weighted_mean_of_inverses
    return weighted_harmonic_mean

# %% ../nbs/03_tune.ipynb 18
def confusion_matrix(
    y_true: Union[np.array, pd.Series],
    y_pred: Union[np.array, pd.Series],
    shade_axis: Optional[Union[str, int]] = None,
    sample_weight: Optional[np.array] = None,
    normalize: Optional[str] = None,
) -> pd.DataFrame:
    """Get confusion matrix

    Parameters:

    - `y_true`: Ground-truth values
    - `y_pred`: Predicted values
    - `shade_axis`: `axis` argument to pass to
    `pd.DataFrame.style.background_gradient`

    The remaining parameters are passed to
    `sklearn.metrics.confusion_matrix`.
    """
    cm = metrics.confusion_matrix(
        y_true=y_true, y_pred=y_pred, sample_weight=sample_weight, normalize=normalize
    )
    vals = sorted(np.unique(y_true))
    cm = pd.DataFrame(
        cm,
        columns=[f"Predicted {val}" for val in vals],
        index=[f"Actual {val}" for val in vals],
    )
    cm.style.background_gradient("Blues", axis=shade_axis)
    cm.loc[:, "Totals"] = cm.sum(axis="columns")
    cm.loc["Totals"] = cm.sum(axis="rows")
    return cm.style.background_gradient(
        "Blues", subset=(cm.index[:-1], cm.columns[:-1]), axis=shade_axis
    )
