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

# %% auto 0
__all__ = ['as_percentage', 'cut_model_by_name', 'MLP', 'Threshold', 'ExperimentalResults', 'RepeatedExperiment']

# %% ../nbs/utils.ipynb 3
from fastai.vision.all import *
from torchvision.models.feature_extraction import create_feature_extractor

from fastai_datasets.all import *

# %% ../nbs/utils.ipynb 4
def as_percentage(x, ndigits=2):
    return f'{round(x * 100, ndigits=ndigits)}%'

# %% ../nbs/utils.ipynb 5
def cut_model_by_name(model, cut):
    graph = create_feature_extractor(model, [cut])
    base_forward = graph.forward
    graph.forward = lambda *args, **kwargs: base_forward(*args, **kwargs)[cut]
    return graph

# %% ../nbs/utils.ipynb 7
class MLP(Module):
    """
    A Multilayer Perceptron comprized of linear layers with ReLU activations.
    """
    def __init__(self,
                 logits: Optional[int], # Number of logits. pass `None` to not include a logits layer after the hidden layers
                 hidden_depth=5, # Number of hidden layers
                 hidden_width=512, # Number of neurons in each hidden layer
                 features_dim=None  # Number of neurons in the last hidden layer. Pass `None` to use the same value as the other hidden layers
                 ):
        super().__init__()
        features_dim = features_dim or hidden_width

        def generate_hidden_layers():
            if hidden_depth == 0:
                return
            if hidden_depth >= 2:
                yield nn.LazyLinear(hidden_width)
                yield nn.ReLU()
                for _ in range(hidden_depth - 2):
                    yield nn.Linear(hidden_width, hidden_width)
                    yield nn.ReLU()
            yield nn.LazyLinear(features_dim)

        self.hidden_layers = nn.Sequential(*generate_hidden_layers())
        self.logits = nn.LazyLinear(logits) if logits else None

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = self.hidden_layers(x)
        if self.logits:
            x = self.logits(x)
        return x

# %% ../nbs/utils.ipynb 8
import torch
from torch import nn

from fastai.vision.all import *

class Threshold(nn.Module):
    """Classifies 1D inputs into 2 classes, based on whether they surpass a threshold or not"""
    def __init__(self) -> None:
        super().__init__()
        self.t = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        x = x - self.t
        return torch.stack([x, -x], dim=-1)

# %% ../nbs/utils.ipynb 11
@patch
def fit(self: Threshold, x, y):
    """Picks a threshold that maximizes the empirical accuracy"""
    with torch.no_grad():
        def accuracy_for_threshold(t):
            self.t[0] = t
            return accuracy(self(x), y)

        threshold_candidates = np.arange(0.0, 4.0, 0.01)
        self.t[0], accuracy_score = max(((t, accuracy_for_threshold(t)) for t in threshold_candidates),
                                                key=lambda p: p[1])

        return self.t.item(), accuracy_score.item()


# %% ../nbs/utils.ipynb 15
from abc import ABC, abstractmethod
from dataclasses import dataclass

import matplotlib.pyplot as plt
    

@dataclass
class ExperimentalResults(object):
    "Provides various ways of examining the results of a `RepeatedExperiment`"
    stats: List[Any]

    @property
    def collated_stats(self):
        return {k: v.numpy() for k, v in default_collate(self.stats).items()}

    def plot_stats(self):
        fig, axs = plt.subplots(len(self.collated_stats), sharex=True)
        fig.tight_layout()
        for ax, (stat_name, stat_values) in zip(axs, self.collated_stats.items()):
            ax.set_title(stat_name)

            min_val, max_val = min(stat_values), max(stat_values)
            val_range = max_val - min_val
            ax.set_ylim(min_val - .1*val_range, max_val + .1*val_range)
            
            ax.bar(range(len(stat_values)), stat_values)

    @property
    def stat_means(self):
        return {k: v.mean() for k, v in self.collated_stats.items()}

    @property
    def stat_stds(self):
        return {k: v.std() for k, v in self.collated_stats.items()}


class RepeatedExperiment(ABC):
    "Runs multiple independent iterations of the same procedure, and combines the results"

    def __init__(self,
                 model: nn.Module,  # The model to be used in each iteration. Parameter are reset to their initial values before each iteration
                 data: List[Datasets],  # A list of `Datasets`, each representing a different iteration. A `Dataloaders` of the current `Datasets` is available via `self.dls`
                 seed: int = 0  # Used for reproducibility of results. Use `None` to avoid reproducibility
                 ):
        super().__init__()
        store_attr('model, data, seed')
    
    def run(self) -> ExperimentalResults:
        "Runs the experiment, returning the results as an `ExperimentalResults`"
        return ExperimentalResults(self._run())

    @return_list
    def _run(self):
        if self.seed is not None:
            set_seed(self.seed, reproducible=True)

        initial_state_dict = deepcopy(self.model.state_dict())

        splits = master_bar(self.data)
        for i, split in enumerate(splits):
            self.dls = split.dls()
            yield self.iteration()
            self.model.load_state_dict(initial_state_dict)

    @abstractmethod
    def iteration(self) -> Dict[str, Any]:
        pass
