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

__all__ = ['most_common_errors', 'n_groups', 'complete_sentence', 'LRPicker', 'ULMFineTuner', 'FastLabelEncoder']


# Cell
import numpy as np
import torch

from htools import auto_repr, valuecheck


# Cell
def most_common_errors(interp):
    """More concise version of `most_confused`. Find the single most common
    error for each true class.

    Parameters
    ----------
    interp: fastai ClassificationInterpretation

    Returns
    -------
    dict[str, str]: Map true class to its most common incorrectly predicted
        class.
    """
    res = {}
    for label, pred, _ in interp.most_confused():
        if label not in res: res[label] = pred
    return res


# Cell
def n_groups(learn):
    """Quickly get the number of layer groups in a fastai learner.

    Parameters
    ----------
    learn: fastai Learner
        It seems that the models themselves don't actually specify groups the
        way I've been doing in torch. Guessing this is because they've reworked
        the API so everything is Sequential and can be indexed into. This does
        make it easier to experiment with different groupings.

    Returns
    -------
    int: 4 for default awd_lstm.
    """
    return len(learn.opt.param_groups)


# Cell
def complete_sentence(learn, start_phrase, n_words=30, n_samples=3, temp=.75):
    """Generate text from a given input. This is a decent way to get a glimpse
    of how a language model is doing.

    Parameters
    ----------
    learn: fastai Learner
    start_phrase: str
        The prompt (start of a sentence) that you want the model to complete.
    n_words: int
        Number of words to generate after the given prompt.
    n_samples: int
        Number of sample sentences to generate.
    temp: float

    Returns
    -------
    list[str]: Each item is one completed sentence.
    """
    return [learn.predict(start_phrase, n_words, temperature=temp)
            for _ in range(n_samples)]


# Cell
@auto_repr
class LRPicker:
    """Take user-suggested LR into account as an "anchor" to ensure `lr_find`
    doesn't choose something too far out of the ordinary.
    """

    @valuecheck
    def __init__(self, sugg_type:('lr_min', 'lr_steep', 'avg')='lr_min',
                 resolve:('avg', 'auto', 'manual')='avg', tol_ratio=.1):
        """
        Parameters
        ----------
        sugg_type: str
            Determines type of auto LR to use. "avg" will take the mean of
            lr_min and lr_steep.
        resolve: str
            Determines method of resolving "disagreements" between default LR
            and suggested LR. You can default to the auto-LR, the
            user-suggested LR, or the mean of the two.
        tol_ratio: float
            Specifies minimum ratio of LR to user-provided LR to consider the
            auto-suggestion as acceptable. This should always be between 0 and
            1 (the ratio is calculated such that the smaller value is in the
            numerator).
        """
        self.sugg_type = sugg_type
        self.resolve = resolve
        self.tol_ratio = tol_ratio

    def lr_find(self, learn, manual_sugg=None, find=True):
        """Find learning rate, incorporating manual suggestion and results of
        `learn.lr_find()`.

        Parameters
        ----------
        learn: fastai Learner
        manual_sugg: float or None
            User-suggested learning rate.
        find: bool
            If True, run lr_find to get an auto-suggestion. Otherwise just use
            user-suggested LR.

        Returns
        -------
        float: Learning rate to use.
        """
        if find:
            auto_sugg = learn.lr_find()
            print('Auto suggestion:', auto_sugg)
            return self._approve_lr(auto_sugg, manual_sugg)
        return manual_sugg

    def _approve_lr(self, auto_sugg, manual_sugg=None):
        """Take auto LR and user-suggested into account to choose a final LR.

        Parameters
        ----------
        auto_sugg: SuggestedLRs
            Object returned by learn.lr_find(). Basically a namedtuple I think.
        manual_sugg: float or None
            User-suggested learning rate.

        Returns
        -------
        float: learning rate
        """
        if self.sugg_type == 'avg':
            sugg_lr = (auto_sugg.lr_min + auto_sugg.lr_steep) / 2
        else:
            sugg_lr = getattr(auto_sugg, self.sugg_type)

        if not manual_sugg or self.resolve == 'auto' or \
                np.divide(*sorted([sugg_lr, manual_sugg])) >= self.tol_ratio:
            return sugg_lr
        if self.resolve == 'avg':
            return (sugg_lr + manual_sugg) / 2
        if self.resolve == 'manual':
            return manual_sugg


# Cell
@auto_repr
class ULMFineTuner:
    """Fine tune language model using ULM Fit procedure. I noticed the built-in
    `fine_tune` method does not unfreeze 1 layer at a time as the paper
    describes - not sure if they found that to be a better practice or if it's
    just simpler for an automated method.

    Originally, part of the reason for building this was to also decrease the
    batch size at each stage since unfreezing eats up more memory with stored
    gradients. However, I decided I'd rather not have to account for changing
    batch size when selecting each stage's LR (we could run lr_find before each
    stage but I opted for the simpler approach).
    """
    moms = (0.8, 0.7, 0.8)
    lr_div = 2.6

    def __init__(self, learn, name_fmt='cls_stage_{}'):
        """
        Parameters
        ----------
        learn: fastai Learner
        name_fmt: str
            Determines name of model that will be saved at each stage.
        """
        self.learn = learn
        # Usually 4 layer groups in fastai text models, but check just in case.
        self.groups = len(learn.opt.param_groups)
        self.name_fmt = name_fmt

    def fine_tune(self, lr, unfrozen_epochs=15, **kwargs):
        """
        Parameters
        ----------
        lr: float
            The learning rate to use for the last layer in the first stage of
            training. The learning rate is halved after each stage. When
            earlier layers are unfrozen, they will use a smaller LR as
            described in the ULM Fit paper.
        unfrozen_epochs: int
            Number of epochs to fit after unfreezing everything. Couldn't find
            a good guideline on this but I think it's fine to be fairly high
            because we can use a callback to save the best model and stop early
            if appropriate.
        kwargs: any
            Passed on to `learn.fit_one_cycle()`.

        Returns
        -------
        None
        """
        for stage in range(1, self.groups+1):
            epochs = 1 if stage < self.groups else unfrozen_epochs
            self.learn.freeze_to(-stage)
            self.learn.fit_one_cycle(epochs,
                                     slice(lr/self.lr_div**self.groups, lr),
                                     moms=self.moms, **kwargs)
            self.learn.save(self.name_fmt.format(stage))
            lr /= 2


# Cell
class FastLabelEncoder:
    """Use a fastai learner to mimic a sklearn label encoder. This can be
    useful because our standard evaluation code is often built when we are
    trying out simple baseline models (e.g. logistic regression using sklearn).
    """

    def __init__(self, learn):
        """
        Parameters
        ----------
        learn: fastai Learner
            The dataloaders must have a `vocab` attribute (text dataloaders do,
            not sure about others).
        """
        self.i2o = self.classes_ = learn.dls.vocab[1]
        self.o2i = {o: i for i, o in enumerate(self.i2o)}

    def transform(self, label_names):
        """Convert strings to integer indices.

        Parameters
        ----------
        label_names: Iterable[str]

        Returns
        -------
        list[int]
        """
        return [self.o2i[o] for o in label_names]

    def inverse_transform(self, label_idx):
        """Convert integer indices to strings.

        Parameters
        ----------
        label_idx: Iterable[int]

        Returns
        -------
        list[str]
        """
        return [self.i2o[i] for i in label_idx]