# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/02b_classification.baseline.baseline.ipynb.

# %% auto 0
__all__ = ['FastaiMetricWrapper', 'calculate_metrics', 'fixed_probs', 'random_probs']

# %% ../../../nbs/02b_classification.baseline.baseline.ipynb 1
from fastai.vision.all import *
from torch.utils.data import TensorDataset

# %% ../../../nbs/02b_classification.baseline.baseline.ipynb 3
class FastaiMetricWrapper(GetAttr):
    _default = 'metric'
    def __init__(self, metric):
        self.metric = metric
        
    def accumulate(self, pred, yb):
        fake_learner = namedtuple('FakeLearner', 'pred yb y to_detach')(pred, (yb,), yb, to_detach)
        self.metric.accumulate(fake_learner)

# %% ../../../nbs/02b_classification.baseline.baseline.ipynb 4
def calculate_metrics(dl, metrics, probs_func, n_classes=None):
    "Calculate metrics for a baseline method `probs_func`"
    n_classes = n_classes or get_c(dl)
    if n_classes == 0:
        raise ValueError("Couldn't automatically infer number of classes from dataloader. Please explicitly specify it with `n_classes`")
        
    metrics = L(metrics).map(compose(mk_metric, FastaiMetricWrapper))
    for metric in metrics: metric.reset()

    for x,y in dl:
        pred = probs_func(x, y, n_classes)
        for metric in metrics: metric.accumulate(pred, y)

    return {metric.name: metric.value for metric in metrics}

# %% ../../../nbs/02b_classification.baseline.baseline.ipynb 8
def fixed_probs(x, y, n_classes, class_id):
    pred = torch.zeros(find_bs(y), n_classes).to(y.device)
    pred[:, class_id] = 1
    return pred

# %% ../../../nbs/02b_classification.baseline.baseline.ipynb 11
def random_probs(x, y, n_classes):
    pred = torch.randn(find_bs(y), n_classes).to(y.device)
    pred = torch.nn.functional.softmax(pred, dim=-1)
    return pred
