# AUTOGENERATED! DO NOT EDIT! File to edit: 01_tune.ipynb (unless otherwise specified).

__all__ = ['calculate_metrics_by_thresh_binary', 'calculate_metrics_by_thresh_multi', 'coverage',
           'calculate_metric_ignoring_nan', 'fbeta', 'confusion_matrix']

# Cell
from collections import defaultdict
from functools import partial
from typing import Callable, Iterable, Optional, Sequence, Union

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

# Cell
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 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)

# Cell
def calculate_metrics_by_thresh_binary(
    y_true: np.array,
    y_prob: np.array,
    metrics: Union[Callable, Sequence[Callable]],
) -> 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.

    Returns: DataFrame with one column "thresh" indicating the
    thresholds used, which is 0 and the sorted set of values that occur in
    `y_prob`, and an additional column for each input metric giving
    the value of that metric at that threshold.
    """
    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=lambda y_prob: sorted(np.hstack([0, np.unique(y_prob[:, 1])])),
    )

# Cell
def calculate_metrics_by_thresh_multi(
    y_true: np.array,
    y_prob: np.array,
    metrics: Union[Callable, Sequence[Callable]],
) -> 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.

    Returns: DataFrame with one column "thresh" indicating the
    thresholds used, which is 0 and the sorted set of values that occur in
    `y_prob`, and an additional column for each input metric giving
    the value of that metric at that threshold.
    """
    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=lambda y_prob: sorted(np.hstack([0, np.unique(y_prob)])),
    )

# Cell
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()

# Cell
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
    )

# Cell
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

# Cell
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))
    result = pd.DataFrame(cm, columns=[f"Predicted {val}" for val in vals], index=[f"Actual {val}" for val in vals]).style.background_gradient('Blues', axis=shade_axis)
    return result