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

__all__ = ['GamblersCallback', 'NoisyStudent', 'TransformScheduler', 'ShowGraphCallback2']

# Cell
from .imports import *
from .utils import *
from .data.transforms import *
from .models.layers import *
from fastai.callback.all import *

# Cell
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

# Cell
class GamblersCallback(Callback):
    "A callback to use metrics with gambler's loss"
    def after_loss(self): self.learn.pred = self.learn.pred[..., :-1]

# Cell
class NoisyStudent(Callback):
    """A callback to implement the Noisy Student approach. In the original paper this was used in combination with noise:
        - stochastic depth: .8
        - RandAugment: N=2, M=27
        - dropout: .5

    Steps:
        1. Build the dl you will use as a teacher
        2. Create dl2 with the unlabeled data plus pseudolabels (either soft or hard preds)
        3. Pass any required batch_tfms to the callback

    """

    def __init__(self, dl2:DataLoader, bs:Union[None, int]=None, unlabel_ratio:int=1, batch_tfms:Union[None, list]=None, do_setup:bool=True,
                 unlabeled_sample_weight:float=1., verbose=False):
        self.dl2, self.bs, self.unlabel_ratio, self.batch_tfms, self.do_setup, self.vebose = dl2, bs, unlabel_ratio, batch_tfms, do_setup, verbose
        self.unlabeled_sample_weight = unlabeled_sample_weight

    def before_fit(self):
        # Remove and store dl.train.batch_tfms
        self.old_train_dl_batch_tfms = self.dls.train.after_batch
        self.dls.train.after_batch = noop

        if self.do_setup and self.batch_tfms:
            for bt in self.batch_tfms: bt.setup(self.dls.train)

        unlabel_ratio = self.unlabel_ratio
        dsets_ratio = len(self.dl2.dataset) / len(self.dls.train.dataset)
        if dsets_ratio < self.unlabel_ratio:  self.unlabel_ratio = dsets_ratio
        if self.bs is None: self.bs = self.dls.train.bs
        self.dl2.bs = min(len(self.dl2.dataset), int(self.unlabel_ratio * self.bs / (1 + self.unlabel_ratio)))
        self.dls.train.bs = self.bs - self.dl2.bs
        if unlabel_ratio != self.dl2.bs/self.dls.train.bs:
            print(f'effective unlabel_ratio          : {unlabel_ratio:.3f} --> {self.dl2.bs/self.dls.train.bs:.3f} \
            dl.bs: {self.dls.train.bs}  dl2.bs: {self.dl2.bs}')
        print(f'relative unlabeled sample weight : {self.dl2.bs/self.dls.train.bs * len(self.dls.train.dataset) / len(self.dl2.dataset):.3f}')
        self.dl2iter = iter(self.dl2)

        self.old_loss_func = self.learn.loss_func
        self.learn.loss_func = self.loss

    def before_batch(self):
        if self.training:
            X, y = self.x, self.y
            try: X2, y2 = next(self.dl2iter)
            except StopIteration:
                self.dl2iter = iter(self.dl2)
                X2, y2 = next(self.dl2iter)
            if y.ndim == 1 and y2.ndim == 2: y = torch.eye(self.learn.dls.c)[y].to(device)
            pv(f'X: {X.shape}  X2: {X2.shape}  y: {y.shape} y2: {y2.shape}', self.vebose)
            X_comb = concat(X, X2) # concat and retain type for tfms
            y_comb = concat(y, y2)
#             perm = torch.randperm(y_comb.shape[0])
#             X_combv = X_comb[perm]
#             y_comb = y_comb[perm]
            if self.batch_tfms is not None:
                X_comb = compose_tfms(X_comb, self.batch_tfms, split_idx=0)
                y_comb = compose_tfms(y_comb, self.batch_tfms, split_idx=0)
            self.learn.xb = (X_comb,)
            self.learn.yb = (y_comb,)
            pv(f'X: {X.shape}  X2: {X2.shape}  X_comb: {X_comb.shape}', self.vebose)

    def loss(self, output, target):
        if target.ndim == 2: _, target = target.max(dim=1)
        if self.training and self.unlabeled_sample_weight != 1:
            loss = (1 - self.unlabeled_sample_weight) * self.old_loss_func(output[:self.dls.train.bs], target[:self.dls.train.bs])
            loss += self.unlabeled_sample_weight * self.old_loss_func(output[self.dls.train.bs:], target[self.dls.train.bs:])
            return loss
        else:
            return self.old_loss_func(output, target)

    def after_fit(self):
        self.dls.train.after_batch = self.old_train_dl_batch_tfms
        self.learn.loss_func = self.old_loss_func
        self.dls.train.bs = self.bs
        self.dls.bs = self.bs

# Cell
class TransformScheduler(Callback):
    "A callback to schedule batch transforms during training based on a function (sched_lin, sched_exp, sched_cos (default), etc)"
    def __init__(self, schedule_func:callable, show_plot:bool=False):
        self.schedule_func,self.show_plot = schedule_func,show_plot
        self.mult = []

    def before_fit(self):
        for pct in np.linspace(0, 1, len(self.dls.train) * self.n_epoch): self.mult.append(self.schedule_func(pct))
        # get initial magnitude values and update initial value
        self.mag = []
        self.mag_tfms = []
        for t in self.dls.after_batch:
            if hasattr(t, 'magnitude'):
                self.mag.append(t.magnitude)
                t.magnitude *= self.mult[0]
                self.mag_tfms.append(t)

    def after_batch(self):
        if self.training and len(self.mag_tfms)>0 and self.train_iter < len(self.mult):
            # set values for next batch
            for t,m in zip(self.mag_tfms, self.mag):
                t.magnitude = m * self.mult[self.train_iter]

    def after_fit(self):
        if self.show_plot and self.mult != [] and len(self.mag_tfms)>0:
            print()
            plt.plot(self.mult)
            plt.title('Scheduled tfms')
            plt.show()
            print()
            self.show_plot = False
        # set values to initial values
        for t,m in zip(self.mag_tfms, self.mag): t.magnitude = m

    def __repr__(self):
        return f'{self.__class__.__name__}({self.schedule_func})'

# Cell
class ShowGraphCallback2(Callback):
    "(Modified) Update a graph of training and validation loss"
    run_after,run_valid=ProgressCallback,False

    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, "gather_preds")
        if not(self.run): return
        self.nb_batches = []
        assert hasattr(self.learn, 'progress')

    def after_train(self): self.nb_batches.append(self.train_iter)

    def after_epoch(self):
        "Plot validation loss in the pbar graph"
        rec = self.learn.recorder
        iters = range_of(rec.losses)
        val_losses = [v[1] for v in rec.values]
        x_bounds = (0, (self.n_epoch - len(self.nb_batches)) * self.nb_batches[0] + len(rec.losses))
        y_min = min((min(rec.losses), min(val_losses)))
        y_max = max((max(rec.losses), max(val_losses)))
        margin = (y_max - y_min) * .05
        y_bounds = (y_min - margin, y_max + margin)
        self.progress.mbar.update_graph([(iters, rec.losses), (self.nb_batches, val_losses)], x_bounds, y_bounds)