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

__all__ = ['CutMixUp', 'CutMixUpAugment']

# Cell
from fastai.basics import *
from fastai.data.transforms import IntToFloatTensor, Normalize
from fastai.callback.mixup import MixHandler, CutMix, MixUp, reduce_loss
from torch.distributions.beta import Beta

# Internal Cell
class CutMix(MixHandler):
    "Implementation of https://arxiv.org/abs/1905.04899"
    def __init__(self, alpha=1.): super().__init__(alpha)
    def before_batch(self):
        bs, _, H, W = self.x.size()
        self.lam = self.distrib.sample((1,)).to(self.x.device)
        shuffle = torch.randperm(bs).to(self.x.device)
        xb1,self.yb1 = self.x[shuffle], tuple((self.y[shuffle],))
        x1, y1, x2, y2 = self.rand_bbox(W, H, self.lam)
        self.learn.xb[0][..., y1:y2, x1:x2] = xb1[..., y1:y2, x1:x2]
        self.lam = (1 - ((x2-x1)*(y2-y1))/float(W*H))
        if not self.stack_y:
            ny_dims = len(self.y.size())
            self.learn.yb = tuple(L(self.yb1,self.yb).map_zip(torch.lerp,weight=unsqueeze(self.lam, n=ny_dims-1)))

    def rand_bbox(self, W, H, lam):
        cut_rat = torch.sqrt(1. - lam).to(self.x.device)
        cut_w = torch.round(W * cut_rat).type(torch.long).to(self.x.device)
        cut_h = torch.round(H * cut_rat).type(torch.long).to(self.x.device)
        # uniform
        cx = torch.randint(0, W, (1,)).to(self.x.device)
        cy = torch.randint(0, H, (1,)).to(self.x.device)
        x1 = torch.clamp(cx - torch.div(cut_w, 2, rounding_mode='floor'), 0, W)
        y1 = torch.clamp(cy - torch.div(cut_h, 2, rounding_mode='floor'), 0, H)
        x2 = torch.clamp(cx + torch.div(cut_w, 2, rounding_mode='floor'), 0, W)
        y2 = torch.clamp(cy + torch.div(cut_h, 2, rounding_mode='floor'), 0, H)
        return x1, y1, x2, y2

# Cell
class CutMixUp(MixUp, CutMix):
    "Combo implementation of https://arxiv.org/abs/1710.09412 and https://arxiv.org/abs/1905.04899"
    run_valid = False
    def __init__(self, mix_alpha=.4, cut_alpha=1., cutmix_ratio=1, mixup_ratio=1):
        MixUp.__init__(self, mix_alpha)
        CutMix.__init__(self, cut_alpha)
        self.mix_distrib = Beta(tensor(mix_alpha), tensor(mix_alpha))
        self.cut_distrib = Beta(tensor(cut_alpha), tensor(cut_alpha))
        self.ratio = mixup_ratio / (cutmix_ratio + mixup_ratio)

    def before_batch(self):
        if torch.rand(1) <= self.ratio: #mixup
            self.distrib = self.mix_distrib
            MixUp.before_batch(self)
        else:
            self.distrib = self.cut_distrib
            CutMix.before_batch(self)

# Cell
class CutMixUpAugment(MixUp, CutMix):
    "Combo implementation of https://arxiv.org/abs/1710.09412 and https://arxiv.org/abs/1905.04899 plus Augmentation"
    run_valid = False
    def __init__(self, mix_alpha=.4, cut_alpha=1., augment_ratio=1, cutmix_ratio=1, mixup_ratio=1, augs_only=None):
        MixUp.__init__(self, mix_alpha)
        CutMix.__init__(self, cut_alpha)
        self.mix_distrib = Beta(tensor(mix_alpha), tensor(mix_alpha))
        self.cut_distrib = Beta(tensor(cut_alpha), tensor(cut_alpha))
        self.aug_cutmix_ratio = augment_ratio / (augment_ratio + cutmix_ratio + mixup_ratio)
        if self.aug_cutmix_ratio == 1: self.cut_mix_ratio = 0
        else: self.cut_mix_ratio = mixup_ratio / (cutmix_ratio + mixup_ratio)
        self.augs_only = augs_only

    def before_fit(self):
        if self.augs_only is None: self.augs_only = (self.learn.n_epoch + 1)/self.learn.n_epoch
        elif self.augs_only >=1: self.augs_only = self.augs_only/self.learn.n_epoch
        else: self.augs_only = self.augs_only

        self._inttofloat_pipe = Pipeline([])
        self._norm_pipe = Pipeline([])

        # first copy transforms
        self._orig_pipe = self.dls.train.after_batch
        self._orig_pipe.split_idx = 0 # need to manually set split_idx for training augmentations to run

        # loop through existing transforms looking for IntToFloatTensor and Normalize
        for i in range(len(self.dls.train.after_batch.fs)):
            if isinstance(self.dls.train.after_batch[i], IntToFloatTensor):
                self._inttofloat_pipe = Pipeline([self.dls.train.after_batch[i]])
            elif isinstance(self.dls.train.after_batch[i], Normalize):
                self._norm_pipe = Pipeline([self.dls.train.after_batch[i]])

        # set existing transforms to an empty Pipeline
        self.dls.train.after_batch = Pipeline([])

    def before_batch(self):
        if self.augs_only >= self.learn.pct_train and torch.rand(1) >= self.aug_cutmix_ratio: # augs or mixup/cutmix
            self._aug = False
            self.learn.xb = self._inttofloat_pipe(self.xb) # apply inttofloat first
            if self.cut_mix_ratio > 0 and torch.rand(1) <= self.cut_mix_ratio: # mixup or cutmix
                self.distrib = self.mix_distrib
                MixUp.before_batch(self)
            else:
                self.distrib = self.cut_distrib
                CutMix.before_batch(self)
            self.learn.xb = self._norm_pipe(self.xb) # now normalize
        else:
            self._aug = True
            self.learn.xb = self._orig_pipe(self.xb) # original transforms

    def after_fit(self):
        self.dls.train.after_batch = self._orig_pipe

    def after_cancel_fit(self):
        self.after_fit()
        MixUp.after_cancel_fit(self)

    def lf(self, pred, *yb):
        if not self.training or self._aug: return self.old_lf(pred, *yb)
        with NoneReduce(self.old_lf) as lf:
            loss = torch.lerp(lf(pred,*self.yb1), lf(pred,*yb), self.lam)
        return reduce_loss(loss, getattr(self.old_lf, 'reduction', 'mean'))