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

# %% ../nbs/transform.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../nbs/transform.ipynb 3
from __future__ import annotations
from torch.distributions import Bernoulli

from fastcore.transform import DisplayedTransform, _is_tuple, retain_type
from fastcore.dispatch import typedispatch, explode_types

from fastai.data.core import TfmdDL

from .imports import *

# %% auto 0
__all__ = ['BatchRandTransform']

# %% ../nbs/transform.ipynb 5
@patch
def to(self:TfmdDL, device):
    self.device = device
    for tfm in self.after_batch.fs:
        for a in L(getattr(tfm, 'parameters', None)): setattr(tfm, a, getattr(tfm, a).to(device))
        if hasattr(tfm, 'to'): tfm.to(device)
    return self

# %% ../nbs/transform.ipynb 6
@patch
def _one_pass(self:TfmdDL):
    b = self.do_batch([self.do_item(None)])
    if self.device is not None:
        b = to_device(b, self.device)
        self.to(self.device)
    its = self.after_batch(b)
    self._n_inp = 1 if not isinstance(its, (list,tuple)) or len(its)==1 else len(its)-1
    self._types = explode_types(its)

# %% ../nbs/transform.ipynb 8
class BatchRandTransform(DisplayedTransform):
    "Randomly selects a subset of batch `b` to apply transform with per item probability `p` in `before_call`"
    do,supports,split_idx = True,[],0
    def __init__(self,
        p:Numeric=1., # Probability of applying Transform to each batch item
        before_call:Callable[[Tensor|tuple[Tensor,...],int],None]|None=None, # Batch preprocessing function
        **kwargs
    ):
        store_attr('p')
        super().__init__(**kwargs)
        self.before_call = ifnone(before_call,self.before_call)
        self.bernoulli = Bernoulli(p)

    def before_call(self,
        b:Tensor|tuple[Tensor,...], # Batch item(s)
        split_idx:int # Train (0) or valid (1) index
    ):
        "Randomly select `self.idxs` and set `self.do` based on `self.p` if not valid `split_idx`"
        self.idxs = self.bernoulli.sample((find_bs(b),)).bool() if not split_idx and self.p<1. else torch.ones(find_bs(b)).bool()
        self.do = self.p==1. or self.idxs.shape[-1] > 0

    def __call__(self,
        b:Tensor|tuple[Tensor,...], # Batch item(s)
        split_idx:int, # Train (0) or valid (1) index
        **kwargs
    ) -> Tensor|tuple[Tensor,...]:
        "Call `super().__call__` if `self.do`"
        self.before_call(b, split_idx=split_idx)
        return super().__call__(b, split_idx=split_idx, **kwargs) if self.do else b

    def _do_call(self,
        f, # Transform
        x:Tensor|tuple[Tensor,...], # Batch item(s)
        **kwargs
    ) -> Tensor|tuple[Tensor,...]:
        "Override `Transform._do_call` to apply transform `f` to `x[self.idxs]`"
        if not _is_tuple(x):
            if f is None: return x
            ret = f.returns(x) if hasattr(f,'returns') else None
            return retain_type(self._do_f(f, x, **kwargs), x, ret)
        res = tuple(self._do_call(f, x_, **kwargs) for x_ in x)
        return retain_type(res, x)

    def _do_f(self,
        f, # Transform
        x:Tensor, # Batch item
        **kwargs
    ) -> Tensor:
        "Apply transform `f` to `x[self.idxs]`"
        x[self.idxs] = f(x[self.idxs], **kwargs)
        return x
