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

__all__ = ['LogMetric', 'MetricType', 'ActivationType', 'MetricX', 'AvgMetricX', 'AccumMetricX', 'AvgSmoothMetricX',
           'AvgLossX', 'AvgSmoothLossX', 'ValueMetricX', 'func_to_metric', 'skm_to_fastxtend', 'Accuracy', 'ErrorRate',
           'TopKAccuracy', 'APScoreBinary', 'BalancedAccuracy', 'BrierScore', 'CohenKappa', 'F1Score', 'FBeta',
           'HammingLoss', 'Jaccard', 'Precision', 'Recall', 'RocAuc', 'RocAucBinary', 'MatthewsCorrCoef',
           'AccuracyMulti', 'APScoreMulti', 'BrierScoreMulti', 'F1ScoreMulti', 'FBetaMulti', 'HammingLossMulti',
           'JaccardMulti', 'MatthewsCorrCoefMulti', 'PrecisionMulti', 'RecallMulti', 'RocAucMulti', 'MSE', 'RMSE',
           'MAE', 'MSLE', 'ExpRMSE', 'ExplainedVariance', 'R2Score', 'PearsonCorrCoef', 'SpearmanCorrCoef',
           'ForegroundAcc', 'Dice', 'DiceMulti', 'JaccardCoeff', 'CorpusBLEUMetric', 'Perplexity', 'perplexity',
           'LossMetric', 'LossMetrics']

# Cell
import inspect, time, math

import numpy as np
import sklearn.metrics as skm
import scipy.stats as scs

from fastcore.basics import class2attr
from fastcore.dispatch import cast

from fastprogress.core import format_time

from fastai.callback.core import Callback
from fastai.learner import AvgMetric, Learner, Metric, Recorder, _maybe_item
from fastai.metrics import AccumMetric
from fastai.torch_core import flatten_check

from .imports import *

# Cell
class LogMetric(Enum):
    "All possible logging types for `MetricX`",
    Train = 1
    Valid = 2
    Both = 3

# Cell
class MetricType(Enum):
    "All possible types of `MetricX`",
    Avg = 1
    Accum = 2
    Smooth = 3

# Cell
class ActivationType(Enum):
    "All possible activation classes for `MetricX",
    No = 1
    Sigmoid = 2
    Softmax = 3
    BinarySoftmax = 4

# Cell
class MetricX(Metric):
    "Blueprint for defining an extended metric with accumulate"
    log_metric=LogMetric.Valid
    def __init__(self, dim_argmax=None, activation=ActivationType.No, thresh=None, log_metric=None, name=None):
        store_attr(but='log_metric, name')
        self.log_metric = ifnone(log_metric, self.log_metric)
        self._name = ifnone(name, class2attr(self, 'MetricX'))

    def reset(self):
        "Reset inner state to prepare for new computation"
        pass

    def accumulate(self, learn):
        "Store targs and preds from `learn`, using activation function and argmax as appropriate"
        self.pred, self.targ = learn.pred, *learn.yb
        if self.activation in [ActivationType.Softmax, ActivationType.BinarySoftmax]:
            self.pred = F.softmax(self.pred, dim=self.dim_argmax)
            if self.activation == ActivationType.BinarySoftmax: self.pred = self.pred[:, -1]
        elif self.activation == ActivationType.Sigmoid: self.pred = torch.sigmoid(self.pred)
        elif self.dim_argmax: self.pred = self.pred.argmax(dim=self.dim_argmax)
        if self.thresh: self.pred = (self.pred >= self.thresh)

    @property
    def value(self):
        "The value of the metric"
        raise NotImplementedError

    @property
    def name(self):
        "Name of the `Metric`, camel-cased and with Metric removed. Or custom name if provided"
        return self._name

    @name.setter
    def name(self, value): self._name = value

    def _split_kwargs(self, method, **kwargs):
        args = [k for k,v in inspect.signature(method).parameters.items() if v.default != inspect.Parameter.empty]
        return {k: kwargs[k] for k in kwargs.keys() if k in args}

# Cell
@delegates(MetricX)
class AvgMetricX(MetricX):
    "Average the values of `func` taking into account potential different batch sizes"
    def __init__(self, func, to_np=False, invert_arg=False, **kwargs):
        super().__init__(**self._split_kwargs(MetricX.__init__, **kwargs))
        self.func, self.fkwargs = func, self._split_kwargs(func, **kwargs)
        self.to_np, self.invert_arg = to_np, invert_arg
        self._name = ifnone(kwargs.get('name', None), self.func.func.__name__ if hasattr(self.func, 'func') else self.func.__name__)

    def reset(self): self.total,self.count = 0.,0

    def accumulate(self, learn):
        super().accumulate(learn)
        bs = find_bs(learn.yb)
        if self.to_np: self.pred,self.targ = learn.to_detach(self.pred).numpy(),learn.to_detach(self.targ).numpy()
        self.total += (self.func(self.targ, self.pred, **self.fkwargs) if self.invert_arg else self.func(self.pred, self.targ, **self.fkwargs))*bs
        self.count += bs

    @property
    def value(self): return self.total/self.count if self.count != 0 else None

# Cell
@delegates(MetricX)
class AccumMetricX(MetricX):
    "Stores predictions and targets on CPU in accumulate to perform final calculations with `func`."
    def __init__(self, func, to_np=False, invert_arg=False, flatten=True, **kwargs):
        super().__init__(**self._split_kwargs(MetricX.__init__, **kwargs))
        self.flatten, self.func, self.fkwargs = flatten, func, self._split_kwargs(func, **kwargs)
        self.to_np, self.invert_arg = to_np, invert_arg
        self._name = ifnone(kwargs.get('name', None), self.func.func.__name__ if hasattr(self.func, 'func') else self.func.__name__)

    def reset(self):
        "Clear all targs and preds"
        self.targs,self.preds = [],[]

    def accumulate(self, learn):
        "Store targs and preds from `learn`, using activation function and argmax as appropriate"
        super().accumulate(learn)
        self.pred,self.targ = learn.to_detach(self.pred),learn.to_detach(self.targ)
        self.accum_values(self.pred, self.targ)

    def accum_values(self, preds, targs):
        "Store targs and preds"
        if self.flatten: preds,targs = flatten_check(preds,targs)
        self.preds.append(preds)
        self.targs.append(targs)

    def __call__(self, preds, targs):
        "Calculate metric on one batch of data"
        self.reset()
        self.accum_values(preds,targs)
        return self.value

    @property
    def value(self):
        "Value of the metric using accumulated preds and targs"
        if len(self.preds) == 0: return
        preds,targs = torch.cat(self.preds),torch.cat(self.targs)
        if self.to_np: preds,targs = preds.numpy(),targs.numpy()
        return self.func(targs, preds, **self.fkwargs) if self.invert_arg else self.func(preds, targs, **self.fkwargs)

# Cell
@delegates(MetricX, but='log_metric')
class AvgSmoothMetricX(MetricX):
    "Smooth average the values of `func` (exponentially weighted with `beta`). Only computed on training set."
    log_metric = LogMetric.Train
    def __init__(self, func, beta=0.98, to_np=False, invert_arg=False, **kwargs):
        super().__init__(**self._split_kwargs(MetricX.__init__, **kwargs))
        self.func, self.fkwargs = func, self._split_kwargs(func, **kwargs)
        self.beta, self.to_np, self.invert_arg = beta, to_np, invert_arg
        self._name = ifnone(kwargs.get('name', None), self.func.func.__name__ if hasattr(self.func, 'func') else self.func.__name__)

    def reset(self): self.count,self.val = 0,tensor(0.)

    def accumulate(self, learn):
        super().accumulate(learn)
        if self.to_np: self.pred,self.targ = learn.to_detach(self.pred).numpy(),learn.to_detach(self.targ).numpy()
        val = self.func(self.targ, self.pred, **self.fkwargs) if self.invert_arg else self.func(self.pred, self.targ, **self.fkwargs)
        if self.to_np: self.val = val*self.beta + val*(1-self.beta)
        else: self.val = torch.lerp(to_detach(val, gather=False), self.val, self.beta)
        self.count += 1

    @property
    def value(self): return self.val/(1-self.beta**self.count) if self.count != 0 else None

# Cell
class AvgLossX(MetricX):
    "Average the losses taking into account potential different batch sizes"
    def reset(self): self.total,self.count = 0.,0
    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        self.total += learn.to_detach(learn.loss.mean())*bs
        self.count += bs
    @property
    def value(self): return self.total/self.count if self.count != 0 else None

# Cell
class AvgSmoothLossX(MetricX):
    "Smooth average of the losses (exponentially weighted with `beta`)"
    log_metric = LogMetric.Train
    def __init__(self, beta=0.98):
        self.beta = beta
    def reset(self):
        self.count,self.val = 0,tensor(0.)
    def accumulate(self, learn):
        self.count += 1
        self.val = torch.lerp(to_detach(learn.loss.mean(), gather=False), self.val, self.beta)
    @property
    def value(self): return self.val/(1-self.beta**self.count)

# Cell
class ValueMetricX(MetricX):
    "Use to include a pre-calculated metric value (for instance calculated in a `Callback`) and returned by `func`"
    def __init__(self, func, name=None, log_metric=None):
        super().__init__(log_metric=log_metric)
        self.func = func
        self._name = ifnone(name, self.func.func.__name__ if hasattr(self.func, 'func') else self.func.__name__)

    @property
    def value(self): return self.func()

# Internal Cell
def _dedup_metric_names(metrics, names):
    dup = set()
    log = metrics.map(lambda o: (LogMetric.Valid if o.log_metric==LogMetric.Both else o.log_metric) if hasattr(o, 'log_metric') else LogMetric.Valid)
    dups = L(set([o[1] for o in zip(log, names) if o in dup or dup.add(o)]))
    indices = names.argwhere(lambda o: o in dups)
    for i in indices:
        if hasattr(metrics[i], 'func'):
            name = metrics[i].func.func.__name__ if hasattr(metrics[i].func, 'func') else metrics[i].func.__name__
        else:
            if isinstance(metrics[i], MetricX):
                name = class2attr(metrics[i], 'MetricX')
            else:
                name = class2attr(metrics[i], 'Metric')
        if not hasattr(metrics[i], 'name') or metrics[i].name == name: # only deduplicate default metric names
            if isinstance(metrics[i], (AvgMetricX, AvgMetric)):
                names[i] = f'avg_{names[i]}'
            elif isinstance(metrics[i], (AccumMetricX, AccumMetric)):
                names[i] = f'accm_{names[i]}'
            elif isinstance(metrics[i], AvgSmoothMetricX):
                names[i] = f'smth_{names[i]}'
    return names

# Internal Cell
@patch
def __init__(self:Recorder, add_time=True, train_metrics=False, valid_metrics=True, beta=0.98):
    store_attr('add_time')
    self.loss,self.smooth_loss = AvgLossX(),AvgSmoothLossX(beta=beta)

@patch
def before_fit(self:Recorder):
    "Prepare state for training"
    self.lrs,self.iters,self.losses,self.values = [],[],[],[]
    names = self.metrics.attrgot('name')
    if len(names.unique()) != len(names):
        names = _dedup_metric_names(self.metrics, names)
    train = self.metrics.argwhere(lambda o: hasattr(o, 'log_metric') and o.log_metric != LogMetric.Valid)
    valid = self.metrics.argwhere(lambda o: not hasattr(o, 'log_metric') or o.log_metric != LogMetric.Train)
    self._train_metsX = self.metrics[train]
    self._valid_metsX = self.metrics[valid]
    train_names = names[train]
    valid_names = names[valid]
    if len(self._train_metsX) > 0:
        train_names = train_names.map('train_{}')
        valid_names = valid_names.map('valid_{}')
    smooth = self._train_metsX.argwhere(lambda o: isinstance(o, (AvgSmoothLossX, AvgSmoothMetricX)))
    self.smooth_mets  = self._train_metsX[smooth]
    self.smooth_names = train_names[smooth]
    self.train_names = L('train_loss') + train_names
    self.valid_names = L('valid_loss') + valid_names
    self.metric_names = L('epoch') + self.train_names + self.valid_names
    if self.add_time: self.metric_names.append('time')
    self.smooth_loss.reset()
    self.loss.reset()
    self.smooth_mets.map(Self.reset())

@patch
def after_batch(self:Recorder):
    "Update all metrics and records lr and smooth loss in training"
    if len(self.yb) == 0: return
    mets = self.train_mets() if self.training else self.valid_mets()
    for met in mets:
        met.accumulate(self.learn)
    if not self.training: return
    self.lrs.append(self.opt.hypers[-1]['lr'])
    self.losses.append(self.smooth_loss.value)
    self.learn.smooth_loss = self.smooth_loss.value

@patch
def before_epoch(self:Recorder):
    "Set timer if `self.add_time=True`"
    self.cancel_train,self.cancel_valid = False,False
    if self.add_time: self.start_epoch = time.time()
    self.log = L(getattr(self, 'epoch', 0))

@patch
def before_train(self:Recorder):
    self._train_metsX.filter(lambda o: not isinstance(o, (AvgSmoothLossX, AvgSmoothMetricX))).map(Self.reset())

@patch
def before_validate(self:Recorder):
    self.valid_mets().map(Self.reset())

@patch
def after_train(self:Recorder):
    self.log += self.train_mets().map(_maybe_item)

@patch
def after_validate(self:Recorder):
    self.log += self.valid_mets().map(_maybe_item)

@patch
def after_epoch(self:Recorder):
    "Store and log the loss/metric values"
    self.learn.final_record = self.log[1:].copy()
    self.values.append(self.learn.final_record)
    if self.add_time: self.log.append(format_time(time.time() - self.start_epoch))
    self.logger(self.log)
    self.iters.append(self.smooth_loss.count)

@patch
def train_mets(self:Recorder):
    if getattr(self, 'cancel_train', False): return L()
    return L(self.smooth_loss) + self._train_metsX

@patch
def valid_mets(self:Recorder):
    if getattr(self, 'cancel_valid', False): return L()
    return L(self.loss) + self._valid_metsX

# Cell
@delegates(MetricX)
def func_to_metric(func, metric_type, is_class, thresh=None, axis=-1, activation=None, log_metric=LogMetric.Valid, **kwargs):
    "Convert `func` metric to a fastai metric"

    dim_argmax = axis if is_class and thresh is None else None
    if activation is None:
        activation = ActivationType.Sigmoid if (is_class and thresh is not None) else ActivationType.No

    if metric_type==MetricType.Accum:
        return AccumMetricX(func, dim_argmax=dim_argmax, activation=activation,
                            thresh=thresh, log_metric=log_metric, **kwargs)
    elif metric_type==MetricType.Avg:
        return AvgMetricX(func, dim_argmax=dim_argmax, activation=activation,
                          thresh=thresh, log_metric=log_metric, **kwargs)
    elif metric_type==MetricType.Smooth:
        if log_metric!=LogMetric.Train:
            name = func.func.__name__ if hasattr(func, 'func') else  func.__name__
            raise ValueError(f'Error with {name}: AvgSmoothMetricX can only run on train. Set `log_metric` to LogMetric.Train.')
        return AvgSmoothMetricX(func, dim_argmax=dim_argmax, activation=activation, thresh=thresh, **kwargs)
    else:
        name = func.func.__name__ if hasattr(func, 'func') else  func.__name__
        raise ValueError(f"Unsupported `metric_type` {metric_type} for metric {name}.")

# Cell
@delegates(MetricX)
def skm_to_fastxtend(func, is_class=True, thresh=None, axis=-1, activation=None, log_metric=LogMetric.Valid, **kwargs):
    "Convert `func` from sklearn.metrics to a fastai metric"
    return func_to_metric(func, MetricType.Accum, is_class, thresh, axis, activation,
                          log_metric, to_np=True, invert_arg=True, **kwargs)

# Internal Cell
def accuracy(inp, targ):
    "Compute accuracy with `targ` when `pred` is bs * n_classes"
    pred,targ = flatten_check(inp, targ)
    return (pred == targ).float().mean()

# Cell
def Accuracy(axis=-1, metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Compute accuracy with `targ` when `pred` is bs * n_classes"
    return func_to_metric(accuracy, metric_type, True, axis=axis, log_metric=log_metric, **kwargs)

# Internal Cell
def error_rate(inp, targ):
    "1 - `accuracy`"
    return 1 - accuracy(inp, targ)

# Cell
def ErrorRate(axis=-1, metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Compute 1 - accuracy with `targ` when `pred` is bs * n_classes"
    return func_to_metric(error_rate, metric_type, True, axis=axis, log_metric=log_metric, **kwargs)

# Internal Cell
def top_k_accuracy(inp, targ, k=5, axis=-1):
    "Computes the Top-k accuracy (`targ` is in the top `k` predictions of `inp`)"
    inp = inp.topk(k=k, dim=axis)[1]
    targ = targ.unsqueeze(dim=axis).expand_as(inp)
    return (inp == targ).sum(dim=-1).float().mean()

# Cell
def TopKAccuracy(k=5, axis=-1, metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Computes the Top-k accuracy (`targ` is in the top `k` predictions of `inp`)"
    return func_to_metric(partial(top_k_accuracy, k=k, axis=axis), metric_type, False,
                          log_metric=log_metric, **kwargs)

# Cell
def APScoreBinary(axis=-1, average='macro', pos_label=1, sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "Average Precision for single-label binary classification problems"
    return skm_to_fastxtend(skm.average_precision_score, axis=axis, activation=ActivationType.BinarySoftmax,
                         average=average, pos_label=pos_label, sample_weight=sample_weight, log_metric=log_metric,
                         **kwargs)

# Cell
def BalancedAccuracy(axis=-1, sample_weight=None, adjusted=False, log_metric=LogMetric.Valid, **kwargs):
    "Balanced Accuracy for single-label binary classification problems"
    return skm_to_fastxtend(skm.balanced_accuracy_score, axis=axis,
                         sample_weight=sample_weight, adjusted=adjusted, log_metric=log_metric, **kwargs)

# Cell
def BrierScore(axis=-1, sample_weight=None, pos_label=None, log_metric=LogMetric.Valid, **kwargs):
    "Brier score for single-label classification problems"
    return skm_to_fastxtend(skm.brier_score_loss, axis=axis,
                         sample_weight=sample_weight, pos_label=pos_label, log_metric=log_metric, **kwargs)

# Cell
def CohenKappa(axis=-1, labels=None, weights=None, sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "Cohen kappa for single-label classification problems"
    return skm_to_fastxtend(skm.cohen_kappa_score, axis=axis, labels=labels, weights=weights,
                         sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def F1Score(axis=-1, labels=None, pos_label=1, average='binary', sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "F1 score for single-label classification problems"
    return skm_to_fastxtend(skm.f1_score, axis=axis, labels=labels, pos_label=pos_label,
                         average=average, sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def FBeta(beta, axis=-1, labels=None, pos_label=1, average='binary', sample_weight=None,
          log_metric=LogMetric.Valid, **kwargs):
    "FBeta score with `beta` for single-label classification problems"
    return skm_to_fastxtend(skm.fbeta_score, axis=axis, beta=beta, labels=labels, pos_label=pos_label,
                         average=average, sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def HammingLoss(axis=-1, sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "Hamming loss for single-label classification problems"
    return skm_to_fastxtend(skm.hamming_loss, axis=axis,
                         sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def Jaccard(axis=-1, labels=None, pos_label=1, average='binary', sample_weight=None,
            log_metric=LogMetric.Valid, **kwargs):
    "Jaccard score for single-label classification problems"
    return skm_to_fastxtend(skm.jaccard_score, axis=axis, labels=labels, pos_label=pos_label,
                         average=average, sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def Precision(axis=-1, labels=None, pos_label=1, average='binary', sample_weight=None,
              log_metric=LogMetric.Valid, **kwargs):
    "Precision for single-label classification problems"
    return skm_to_fastxtend(skm.precision_score, axis=axis, labels=labels, pos_label=pos_label,
                         average=average, sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def Recall(axis=-1, labels=None, pos_label=1, average='binary', sample_weight=None,
           log_metric=LogMetric.Valid, **kwargs):
    "Recall for single-label classification problems"
    return skm_to_fastxtend(skm.recall_score, axis=axis, labels=labels, pos_label=pos_label,
                         average=average, sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def RocAuc(axis=-1, average='macro', sample_weight=None, max_fpr=None, multi_class='ovr',
           log_metric=LogMetric.Valid, **kwargs):
    "Area Under the Receiver Operating Characteristic Curve for single-label multiclass classification problems"
    assert multi_class in ['ovr', 'ovo']
    return skm_to_fastxtend(skm.roc_auc_score, axis=axis, activation=ActivationType.Softmax, flatten=False,
                         average=average, sample_weight=sample_weight, max_fpr=max_fpr, multi_class=multi_class,
                         log_metric=log_metric, **kwargs)

# Cell
def RocAucBinary(axis=-1, average='macro', sample_weight=None, max_fpr=None, multi_class='raise',
                 log_metric=LogMetric.Valid, **kwargs):
    "Area Under the Receiver Operating Characteristic Curve for single-label binary classification problems"
    return skm_to_fastxtend(skm.roc_auc_score, axis=axis, activation=ActivationType.BinarySoftmax,
                         average=average, sample_weight=sample_weight, max_fpr=max_fpr, multi_class=multi_class,
                         log_metric=log_metric, **kwargs)

# Cell
def MatthewsCorrCoef(sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "Matthews correlation coefficient for single-label classification problems"
    return skm_to_fastxtend(skm.matthews_corrcoef, sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Internal Cell
def accuracy_multi(inp, targ):
    "Compute accuracy when `inp` and `targ` are the same size."
    inp,targ = flatten_check(inp,targ)
    return (inp==targ.bool()).float().mean()

# Cell
def AccuracyMulti(thresh=0.5, sigmoid=True, metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Compute accuracy when `inp` and `targ` are the same size."
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return func_to_metric(accuracy_multi, metric_type, False, thresh=thresh, activation=activation,
                          log_metric=log_metric, **kwargs)

# Cell
def APScoreMulti(sigmoid=True, average='macro', pos_label=1, sample_weight=None,
                 log_metric=LogMetric.Valid, **kwargs):
    "Average Precision for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.average_precision_score, activation=activation, flatten=False,
                         average=average, pos_label=pos_label, sample_weight=sample_weight,
                         log_metric=log_metric, **kwargs)

# Cell
def BrierScoreMulti(thresh=0.5, sigmoid=True, sample_weight=None, pos_label=None,
                    log_metric=LogMetric.Valid, **kwargs):
    "Brier score for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.brier_score_loss, thresh=thresh, activation=activation, flatten=False,
                         sample_weight=sample_weight, pos_label=pos_label, log_metric=log_metric, **kwargs)

# Cell
def F1ScoreMulti(thresh=0.5, sigmoid=True, labels=None, pos_label=1, average='macro', sample_weight=None,
                 log_metric=LogMetric.Valid, **kwargs):
    "F1 score for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.f1_score, thresh=thresh, activation=activation, flatten=False,
                         labels=labels, pos_label=pos_label, average=average, sample_weight=sample_weight,
                         log_metric=log_metric, **kwargs)

# Cell
def FBetaMulti(beta, thresh=0.5, sigmoid=True, labels=None, pos_label=1, average='macro', sample_weight=None,
               log_metric=LogMetric.Valid, **kwargs):
    "FBeta score with `beta` for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.fbeta_score, thresh=thresh, activation=activation, flatten=False,
                beta=beta, labels=labels, pos_label=pos_label, average=average, sample_weight=sample_weight,
                log_metric=log_metric, **kwargs)

# Cell
def HammingLossMulti(thresh=0.5, sigmoid=True, labels=None, sample_weight=None,
                     log_metric=LogMetric.Valid, **kwargs):
    "Hamming loss for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.hamming_loss, thresh=thresh, activation=activation, flatten=False,
                         sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def JaccardMulti(thresh=0.5, sigmoid=True, labels=None, pos_label=1, average='macro', sample_weight=None,
                 log_metric=LogMetric.Valid, **kwargs):
    "Jaccard score for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.jaccard_score, thresh=thresh, activation=activation, flatten=False,
                         labels=labels, pos_label=pos_label, average=average, sample_weight=sample_weight,
                         log_metric=log_metric, **kwargs)

# Cell
def MatthewsCorrCoefMulti(thresh=0.5, sigmoid=True, sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "Matthews correlation coefficient for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.matthews_corrcoef, thresh=thresh, activation=activation, flatten=False,
                         sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def PrecisionMulti(thresh=0.5, sigmoid=True, labels=None, pos_label=1, average='macro', sample_weight=None,
                   log_metric=LogMetric.Valid, **kwargs):
    "Precision for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.precision_score, thresh=thresh, activation=activation, flatten=False,
                         labels=labels, pos_label=pos_label, average=average, sample_weight=sample_weight,
                         log_metric=log_metric, **kwargs)

# Cell
def RecallMulti(thresh=0.5, sigmoid=True, labels=None, pos_label=1, average='macro', sample_weight=None,
                log_metric=LogMetric.Valid, **kwargs):
    "Recall for multi-label classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.recall_score, thresh=thresh, activation=activation, flatten=False,
                         labels=labels, pos_label=pos_label, average=average, sample_weight=sample_weight,
                         log_metric=log_metric, **kwargs)

# Cell
def RocAucMulti(sigmoid=True, average='macro', sample_weight=None, max_fpr=None, log_metric=LogMetric.Valid, **kwargs):
    "Area Under the Receiver Operating Characteristic Curve for multi-label binary classification problems"
    activation = ActivationType.Sigmoid if sigmoid else ActivationType.No
    return skm_to_fastxtend(skm.roc_auc_score, activation=activation, flatten=False,
                         average=average, sample_weight=sample_weight, max_fpr=max_fpr, log_metric=log_metric, **kwargs)

# Internal Cell
def mse(inp,targ):
    "Mean squared error between `inp` and `targ`."
    return F.mse_loss(*flatten_check(inp,targ))

# Cell
def MSE(metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Mean squared error between `inp` and `targ`."
    return func_to_metric(mse, metric_type, False, log_metric=log_metric, **kwargs)

# Internal Cell
def rmse(inp, targ):
    return torch.sqrt(F.mse_loss(inp, targ))

# Cell
def RMSE(log_metric=LogMetric.Valid, **kwargs):
    "Mean squared error between `inp` and `targ`."
    return func_to_metric(rmse, MetricType.Accum, False, log_metric=log_metric, **kwargs)

# Internal Cell
def mae(inp,targ):
    "Mean absolute error between `inp` and `targ`."
    inp,targ = flatten_check(inp,targ)
    return torch.abs(inp - targ).mean()

# Cell
def MAE(metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Mean absolute error between `inp` and `targ`."
    return func_to_metric(mae, metric_type, False, log_metric=log_metric, **kwargs)

# Internal Cell
def msle(inp, targ):
    "Mean squared logarithmic error between `inp` and `targ`."
    inp,targ = flatten_check(inp,targ)
    return F.mse_loss(torch.log(1 + inp), torch.log(1 + targ))

# Cell
def MSLE(metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Mean squared logarithmic error between `inp` and `targ`."
    return func_to_metric(msle, metric_type, False, log_metric=log_metric, **kwargs)

# Internal Cell
def exp_rmspe(inp,targ):
    inp,targ = torch.exp(inp),torch.exp(targ)
    return torch.sqrt(((targ - inp)/targ).pow(2).mean())

# Cell
def ExpRMSE(log_metric=LogMetric.Valid, **kwargs):
    "Root mean square percentage error of the exponential of  predictions and targets"
    return func_to_metric(exp_rmspe, MetricType.Accum, False, log_metric=log_metric, **kwargs)

# Cell
def ExplainedVariance(sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "Explained variance between predictions and targets"
    return skm_to_fastxtend(skm.explained_variance_score, is_class=False,
                         sample_weight=sample_weight, log_metric=log_metric, **kwargs)

# Cell
def R2Score(sample_weight=None, log_metric=LogMetric.Valid, **kwargs):
    "R2 score between predictions and targets"
    return skm_to_fastxtend(skm.r2_score, is_class=False, sample_weight=sample_weight,
                         log_metric=log_metric, **kwargs)

# Cell
def PearsonCorrCoef(dim_argmax=None, log_metric=LogMetric.Valid, **kwargs):
    "Pearson correlation coefficient for regression problem"
    def pearsonr(x,y): return scs.pearsonr(x,y)[0]
    return AccumMetricX(pearsonr, invert_arg=False, dim_argmax=dim_argmax,
                        log_metric=log_metric, **kwargs)

# Cell
def SpearmanCorrCoef(dim_argmax=None, axis=0, nan_policy='propagate', log_metric=LogMetric.Valid, **kwargs):
    "Spearman correlation coefficient for regression problem"
    def spearmanr(a,b=None,**kwargs): return scs.spearmanr(a,b,**kwargs)[0]
    return AccumMetricX(partial(spearmanr, axis=axis, nan_policy=nan_policy),
                        invert_arg=False, dim_argmax=dim_argmax, log_metric=log_metric, **kwargs)

# Internal Cell
def foreground_acc(inp, targ, bkg_idx=0, axis=1):
    "Computes non-background accuracy for multiclass segmentation"
    targ = cast(targ.squeeze(1), TensorBase)
    mask = targ != bkg_idx
    return (inp[mask]==targ[mask]).float().mean()

# Cell
def ForegroundAcc(bkg_idx=0, axis=1, metric_type=MetricType.Avg, log_metric=LogMetric.Valid, **kwargs):
    "Computes non-background accuracy for multiclass segmentation"
    return func_to_metric(foreground_acc, metric_type, True, bkg_idx=bkg_idx, axis=axis,
                          log_metric=log_metric, **kwargs)

# Cell
class Dice(MetricX):
    "Dice coefficient metric for binary target in segmentation"
    def __init__(self, axis=1, log_metric=LogMetric.Valid, **kwargs):
        super().__init__(dim_argmax=axis, log_metric=log_metric, **kwargs)
    def reset(self): self.inter,self.union = 0,0
    def accumulate(self, learn):
        super().accumulate(learn)
        self.pred,self.targ = flatten_check(self.pred, self.targ)
        self.inter += (self.pred*self.targ).float().sum().item()
        self.union += (self.pred+self.targ).float().sum().item()

    @property
    def value(self): return 2. * self.inter/self.union if self.union > 0 else None

# Cell
class DiceMulti(MetricX):
    "Averaged Dice metric (Macro F1) for multiclass target in segmentation"
    def __init__(self, axis=1, log_metric=LogMetric.Valid, **kwargs):
        super().__init__(dim_argmax=axis, log_metric=log_metric, **kwargs)
    def reset(self): self.inter,self.union = {},{}
    def accumulate(self, learn):
        super().accumulate(learn)
        self.pred,self.targ = flatten_check(self.pred, self.targ)
        for c in range(learn.pred.shape[self.dim_argmax]):
            p = torch.where(self.pred == c, 1, 0)
            t = torch.where(self.targ == c, 1, 0)
            c_inter = (p*t).float().sum().item()
            c_union = (p+t).float().sum().item()
            if c in self.inter:
                self.inter[c] += c_inter
                self.union[c] += c_union
            else:
                self.inter[c] = c_inter
                self.union[c] = c_union

    @property
    def value(self):
        binary_dice_scores = np.array([])
        for c in self.inter:
            binary_dice_scores = np.append(binary_dice_scores, 2.*self.inter[c]/self.union[c] if self.union[c] > 0 else np.nan)
        return np.nanmean(binary_dice_scores)

# Cell
class JaccardCoeff(Dice):
    "Implementation of the Jaccard coefficient that is lighter in RAM"
    @property
    def value(self): return self.inter/(self.union-self.inter) if self.union > 0 else None

# Cell
class CorpusBLEUMetric(MetricX):
    "BLEU Metric calculated over the validation corpus"
    def __init__(self, vocab_sz=5000, axis=-1, log_metric=LogMetric.Valid, name='CorpusBLEU', **kwargs):
        super().__init__(log_metric=log_metric, name=name, **kwargs)
        self.axis, self.vocab_sz = axis, vocab_sz
        self.pred_len,self.targ_len,self.samp_idx,self.corrects,self.counts, = 0,0,0,[0]*4,[0]*4

    def reset(self):
        self.pred_len,self.targ_len,self.corrects,self.counts = 0,0,[0]*4,[0]*4

    class NGram():
        def __init__(self, ngram, max_n=5000): self.ngram,self.max_n = ngram,max_n
        def __eq__(self, other):
            if len(self.ngram) != len(other.ngram): return False
            return np.all(np.array(self.ngram) == np.array(other.ngram))
        def __hash__(self): return int(sum([o * self.max_n**i for i,o in enumerate(self.ngram)]))

    def get_grams(self, x, n, max_n=5000):
        return x if n==1 else [self.NGram(x[i:i+n], max_n=max_n) for i in range(len(x)-n+1)]

    def get_correct_ngrams(self, pred, targ, n, max_n=5000):
        pred_grams,targ_grams = self.get_grams(pred, n, max_n=max_n),self.get_grams(targ, n, max_n=max_n)
        pred_cnt,targ_cnt = Counter(pred_grams),Counter(targ_grams)
        return sum([min(c, targ_cnt[g]) for g,c in pred_cnt.items()]),len(pred_grams)

    def accumulate(self, learn):
        if learn.training: return None
        else:
            last_output = learn.pred.argmax(dim=self.axis)
            last_target = learn.y
            for pred,targ in zip(last_output.cpu().numpy(),last_target.cpu().numpy()):
                self.pred_len += len(pred)
                self.targ_len += len(targ)
                smooth_mteval = 1
                for i in range(4):
                    c,t = self.get_correct_ngrams(pred, targ, i+1, max_n=self.vocab_sz)
                    if c == 0:
                        smooth_mteval *= 2
                        c = 1 / smooth_mteval    # exp smoothing, method 3 from http://acl2014.org/acl2014/W14-33/pdf/W14-3346.pdf
                    self.corrects[i] += c
                    self.counts[i]   += t

    @property
    def value(self):
        if self.counts == 0: return None
        elif max(self.corrects) == 0: return 0.0
        else:
            precs = [c/t for c,t in zip(self.corrects,self.counts)]
            len_penalty = math.exp(1 - self.targ_len/self.pred_len) if self.pred_len < self.targ_len else 1
            return len_penalty * ((precs[0]*precs[1]*precs[2]*precs[3]) ** 0.25)

# Cell
class Perplexity(AvgLossX):
    "Perplexity (exponential of cross-entropy loss) for Language Models"
    @property
    def value(self): return torch.exp(self.total/self.count) if self.count != 0 else None
    @property
    def name(self):  return "perplexity"

perplexity = Perplexity()

# Cell
class LossMetric(AvgMetricX):
    "Create a metric from `loss_func.attr` named `nm`"
    def __init__(self, attr, nm=None, log_metric=LogMetric.Valid, **kwargs):
        super().__init__(noop, log_metric=log_metric, **kwargs)
        store_attr('attr,nm')
    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        self.total += learn.to_detach(getattr(learn.loss_func, self.attr, 0))*bs
        self.count += bs

    @property
    def name(self): return self.attr if self.nm is None else self.nm

# Cell
def LossMetrics(attrs, nms=None):
    "List of `LossMetric` for each of `attrs` and `nms`"
    if isinstance(attrs, str): attrs = attrs.split(',')
    nms = attrs if nms is None else nms.split(',') if isinstance(nms, str) else nms
    return [LossMetric(a, n) for a,n in zip(attrs,nms)]

# Internal Cell
try:
    import wandb
    from fastai.callback.wandb import WandbCallback
    from fastai.callback.training import FetchPredsCallback

    if not hasattr(WandbCallback,'_metrics_before_fit'): WandbCallback._metrics_before_fit = WandbCallback.before_fit

    @patch
    def before_fit(self:WandbCallback):
        self._metrics_before_fit()
        self.log_smooth = len(self.recorder.smooth_names) > 0

    @patch
    def after_batch(self:WandbCallback):
        "Log hyper-parameters and training loss"
        if self.training:
            self._wandb_step += 1
            self._wandb_epoch += 1/self.n_iter
            hypers = {f'{k}_{i}':v for i,h in enumerate(self.opt.hypers) for k,v in h.items()}
            if self.log_smooth:
                for n,m in zip(self.recorder.smooth_names, self.recorder.smooth_mets): hypers[n]=m.value
            wandb.log({'epoch': self._wandb_epoch, 'train_loss': to_detach(self.smooth_loss.clone()), 'raw_loss': to_detach(self.loss.clone()), **hypers}, step=self._wandb_step)

    @patch
    def after_epoch(self:WandbCallback):
        "Log validation loss and custom metrics & log prediction samples"
        # Correct any epoch rounding error and overwrite value
        self._wandb_epoch = round(self._wandb_epoch)
        wandb.log({'epoch': self._wandb_epoch}, step=self._wandb_step)
        # Log sample predictions
        if self.log_preds:
            try:
                self.log_predictions(self.learn.fetch_preds.preds)
            except Exception as e:
                self.log_preds = False
                self.remove_cb(FetchPredsCallback)
                print(f'WandbCallback was not able to get prediction samples -> {e}')
        wandb.log({n:s for n,s in zip(self.recorder.metric_names, self.recorder.log) if n not in ['train_loss', 'epoch', 'time']+self.recorder.smooth_names}, step=self._wandb_step)
except:
    pass

# Internal Cell
try:
    import tensorboard
    from fastai.callback.tensorboard import TensorBoardCallback, tensorboard_log

    if not hasattr(TensorBoardCallback,'_metrics_before_fit'): TensorBoardCallback._metrics_before_fit = TensorBoardCallback.before_fit

    @patch
    def before_fit(self:TensorBoardCallback):
        self._metrics_before_fit()
        self.log_smooth = len(self.recorder.smooth_names) > 0

    @patch
    def after_batch(self:TensorBoardCallback):
        self.writer.add_scalar('train_loss', self.smooth_loss, self.train_iter)
        for i,h in enumerate(self.opt.hypers):
            for k,v in h.items(): self.writer.add_scalar(f'{k}_{i}', v, self.train_iter)
        if self.log_smooth:
            for k,v in zip(self.recorder.smooth_names, self.recorder.smooth_mets):
                self.writer.add_scalar(f'{k}', v, self.train_iter)

    @patch
    def after_epoch(self:TensorBoardCallback):
        for n,v in zip(self.recorder.metric_names[2:-1], self.recorder.log[2:-1]):
            if n not in ['train_loss', 'time']+self.recorder.smooth_names:
                self.writer.add_scalar(n, v, self.train_iter)
        if self.log_preds:
            b = self.dls.valid.one_batch()
            self.learn.one_batch(0, b)
            preds = getattr(self.loss_func, 'activation', noop)(self.pred)
            out = getattr(self.loss_func, 'decodes', noop)(preds)
            x,y,its,outs = self.dls.valid.show_results(b, out, show=False, max_n=self.n_preds)
            tensorboard_log(x, y, its, outs, self.writer, self.train_iter)
except:
    pass