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

__all__ = ['PredictionStd']


# Cell
import numpy as np
import torch
from torchmetrics import Metric

from htools import auto_repr, valuecheck
from .utils import pool_stds


# Cell
class PredictionStd(Metric):
    """Pytorch lightning metric that computes the standard deviation of soft
    predictions over the course of an epoch. This can be helpful when
    troubleshooting: I've often encountered models that output a very narrow
    range of values, either early in training or as the result of some sort
    of bug. This can be used for regression (raw outputs) or classification
    (logits or final soft predictions).
    """

    def __init__(self, dist_sync_on_step=False):
        super().__init__(dist_sync_on_step=dist_sync_on_step)
        self.add_state('batch_sizes', default=[], dist_reduce_fx='cat')
        self.add_state('stds', default=[], dist_reduce_fx='cat')
        self.add_state('means', default=[], dist_reduce_fx='cat')

    def update(self, preds, target):
        assert preds.shape == target.shape, 'Preds and target should be the '\
            ' same shape.'
        assert preds.ndim == 1 or (preds.ndim == 2 and preds.shape[1] == 1)

        n_unique = torch.unique(preds, sorted=True).shape[0]
        if n_unique <= 2:
            warnings.warn(f'Only {n_unique} unique values in preds. Are you '
                          'sure you\'re not passing in hard predictions from '
                          'a classifier?')

        self.batch_sizes.append(preds.shape[0])
        self.stds.append(preds.std())
        self.means.append(preds.mean())

    def compute(self):
        return pool_stds(self.batch_sizes, self.means, self.stds)