"""Callbacks that can apply the MixUp (and variants) data augmentation to your training

Docs: https://docs.fast.ai/callback.mixup.html.md"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/19_callback.mixup.ipynb.

# %% auto #0
__all__ = ['reduce_loss', 'MixHandler', 'MixUp', 'CutMix']

# %% ../../nbs/19_callback.mixup.ipynb #abdd7b68
from ..basics import *
from torch.distributions.beta import Beta

# %% ../../nbs/19_callback.mixup.ipynb #49465863
def reduce_loss(
    loss:Tensor, 
    reduction:str='mean' # PyTorch loss reduction
)->Tensor:
    "Reduce the loss based on `reduction`"
    return loss.mean() if reduction == 'mean' else loss.sum() if reduction == 'sum' else loss

# %% ../../nbs/19_callback.mixup.ipynb #ec396808
class MixHandler(Callback):
    "A handler class for implementing `MixUp` style scheduling"
    run_valid = False
    def __init__(self, 
        alpha:float=0.5 # Determine `Beta` distribution in range (0.,inf]
    ):
        self.distrib = Beta(tensor(alpha), tensor(alpha))

    def before_train(self):
        "Determine whether to stack y"
        self.stack_y = getattr(self.learn.loss_func, 'y_int', False)
        if self.stack_y: self.old_lf,self.learn.loss_func = self.learn.loss_func,self.lf

    def after_train(self):
        "Set the loss function back to the previous loss"
        if self.stack_y: self.learn.loss_func = self.old_lf

    def after_cancel_train(self):
        "If training is canceled, still set the loss function back"
        self.after_train()

    def after_cancel_fit(self):
        "If fit is canceled, still set the loss function back"
        self.after_train()

    def lf(self, pred, *yb):
        "lf is a loss function that applies the original loss function on both outputs based on `self.lam`"
        if not self.training: 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'))

# %% ../../nbs/19_callback.mixup.ipynb #b668a5e3
class MixUp(MixHandler):
    "Implementation of https://arxiv.org/abs/1710.09412"
    def __init__(self, 
        alpha:float=.4 # Determine `Beta` distribution in range (0.,inf]
    ): 
        super().__init__(alpha)
        
    def before_batch(self):
        "Blend xb and yb with another random item in a second batch (xb1,yb1) with `lam` weights"
        lam = self.distrib.sample((self.y.size(0),)).squeeze().to(self.x.device)
        lam = torch.stack([lam, 1-lam], 1)
        self.lam = lam.max(1)[0]
        shuffle = torch.randperm(self.y.size(0)).to(self.x.device)
        xb1,self.yb1 = tuple(L(self.xb).itemgot(shuffle)),tuple(L(self.yb).itemgot(shuffle))
        nx_dims = len(self.x.size())
        self.learn.xb = tuple(L(xb1,self.xb).map_zip(torch.lerp,weight=unsqueeze(self.lam, n=nx_dims-1)))

        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)))

# %% ../../nbs/19_callback.mixup.ipynb #253ea6ae
class CutMix(MixHandler):
    "Implementation of https://arxiv.org/abs/1905.04899"
    def __init__(self,
        alpha:float=1. # Determine `Beta` distribution in range (0.,inf]
    ):
        super().__init__(alpha)

    def before_batch(self):
        "Add `rand_bbox` patches with size based on `lam` and location chosen randomly."
        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:int, # Input image width
        H:int, # Input image height
        lam:Tensor # lambda sample from Beta distribution i.e tensor([0.3647])
    ) -> tuple: # Represents the top-left pixel location and the bottom-right pixel location
        "Give a bounding box location based on the size of the im and a weight"
        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