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

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

# %% ../../nbs/data.loader.ipynb 4
from __future__ import annotations

from typing import Tuple

from fastcore.basics import GetAttr, detuplify, Inf
from fastcore.dispatch import retain_types, explode_types
from fastcore.meta import funcs_kwargs
from fastcore.transform import Pipeline

from fastai.data.core import show_batch, show_results

from ..imports import *

# %% auto 0
__all__ = []

# %% ../../nbs/data.loader.ipynb 6
@funcs_kwargs
class BaseLoader(GetAttr):
    "Provides callbacks for DataLoaders which inherit from `BaseLoader`"
    _methods = 'before_iter after_batch after_iter'.split()
    def __init__(self, **kwargs):
        pass

    def before_iter(self, x=None, *args, **kwargs):
        "Called before `BaseLoader` starts to read/iterate over the dataset."
        return x

    def after_batch(self, x=None, *args, **kwargs):
        "After collating mini-batch of items, the mini-batch is passed through this function."
        return x

    def after_iter(self, x=None, *args, **kwargs):
        "Called after `BaseLoader` has fully read/iterated over the dataset."
        return x

# %% ../../nbs/data.loader.ipynb 7
class DataLoaderMixin(BaseLoader):
    "A Mixin class for adding fastai DataLoader functionality to other DataLoader"
    _methods = 'before_iter after_batch after_iter'.split()
    def __init__(self,
        device:str|int|torch.device|None=None, # Device to place batch. Defaults to fastai's `default_device`
        n_inp:int|None=None, # Number of inputs to the model. Defaults to pipelines length minus 1
        split_idx:int|None=None, # Apply batch transform(s) to training (0) or validation (1) set. Defaults to valid if order is `SEQEUNTIAL`
        do_setup:bool=True, # Run `setup()` for batch transform(s)
        **kwargs
    ):
        if 'batch_tfms' in kwargs:
            if 'after_batch' not in kwargs:
                kwargs['after_batch'] = kwargs.pop('batch_tfms')
            else:
                raise ValueError(f'Cannot pass both `after_batch` and `batch_tfms` to `{self.__name__}`')

        self._split_idx = split_idx

        kwargs['after_batch'] = Pipeline(kwargs.get('after_batch', None), split_idx=self._split_idx)
        if do_setup:
            kwargs['after_batch'].setup(self)

        if device is None:
            self.device = default_device()
        else:
            self.device = device

        self._n_inp = n_inp

        self._callback_warning(kwargs)
        super().__init__(**kwargs)

    def decode(self, b:Tuple[Tensor,...]):
        "Decode batch `b`"
        return to_cpu(self.after_batch.decode(self._retain_dl(b)))

    def decode_batch(self, b:Tuple[Tensor,...], max_n:int=9):
        "Decode up to `max_n` input(s) from batch `b`"
        return self._decode_batch(self.decode(b), max_n)

    def one_batch(self):
        "Return one processed batch of input(s) and target(s)"
        raise NotImplementedError()

    def show_batch(self,
        b:Tuple[Tensor,...]|None=None, # Batch to show. If None calls `one_batch`
        max_n:int=9, # Maximum number of items to show
        ctxs=None, # List of `ctx` objects to show data. Could be matplotlib axis, DataFrame etc
        show:bool=True, # If False, return decoded batch instead of showing
        unique:bool=False, # Whether to show only one
        **kwargs
    ):
        "Show `max_n` input(s) and target(s) from the batch."
        if unique:
            old_get_idxs = self.get_idxs
            self.get_idxs = lambda: Inf.zeros
        if b is None:
            b = self.one_batch()
        if not show:
            return self._pre_show_batch(b, max_n=max_n)
        # Uses Type Dispatch to call the correct `show_batch` for b
        show_batch(*self._pre_show_batch(b, max_n=max_n), ctxs=ctxs, max_n=max_n, **kwargs)
        if unique:
            self.get_idxs = old_get_idxs

    def show_results(self,
        b, # Batch to show results for
        out, # Predicted output from model for the batch
        max_n:int=9, # Maximum number of items to show
        ctxs=None, # List of `ctx` objects to show data. Could be matplotlib axis, DataFrame etc
        show:bool=True, # If False, return decoded batch instead of showing
        **kwargs
    ):
        "Show `max_n` results with input(s), target(s) and prediction(s)."
        x,y,its = self.show_batch(b, max_n=max_n, show=False)
        b_out = type(b)(b[:self.n_inp] + (tuple(out) if is_listy(out) else (out,)))
        x1,_,outs = self.show_batch(b_out, max_n=max_n, show=False)
        if its is None:
            res = (x, x1, None, None)
        else:
            res = (x, y, its, outs.itemgot(slice(self.n_inp,None)))
        if not show:
            return res
        # Uses Type Dispatch to call the correct `show_results` for b
        show_results(*res, ctxs=ctxs, max_n=max_n, **kwargs)

    @property
    def n_inp(self) -> int:
        "Number of elements in a batch for model input"
        return self._n_inp

    @property
    def bs(self) -> int:
        "Number of items a batch"
        raise NotImplementedError()

    @property
    def device(self) -> torch.device:
        return self._device

    @device.setter
    def device(self, device:int|str|torch.device):
        # parse device
        device, *_ = torch._C._nn._parse_to(device=device)
        self._device = device
        # Device setter for DataLoaderMixin.batch_tfms
        if hasattr(self.after_batch, 'fs'):
            self._pipeline_device(self.after_batch.fs)

    def to(self, device:int|str|torch.device):
        "Sets `self.device=device`."
        self.device = device
        return self

    @property
    def split_idx(self) -> int|None:
        return self._split_idx

    @split_idx.setter
    def split_idx(self, split_idx:int):
        "Sets fastai batch transforms to train (split_idx=0) or valid (split_idx=1)"
        self._split_idx = split_idx
        if isinstance(self.after_batch, Pipeline):
            self.after_batch.split_idx = split_idx

    def __iter__(self):
        raise NotImplementedError()

    def _one_pass(self, b=None):
        if b is None:
            b = self.one_batch()
        self._types = explode_types(b)

    def _retain_dl(self, b):
        if not getattr(self, '_types', None):
            self._one_pass(b)
        return retain_types(b, typs=self._types)

    def _decode_batch(self, b, max_n=9):
        return L(batch_to_samples(b, max_n=max_n))

    def _pre_show_batch(self, b, max_n=9):
        "Decode `b` to be ready for `show_batch`"
        b = self.decode(b)
        if hasattr(b, 'show'):
            return b,None,None
        its = self._decode_batch(b, max_n)
        if not is_listy(b):
            b,its = [b],L((o,) for o in its)
        return detuplify(b[:self.n_inp]),detuplify(b[self.n_inp:]),its

    def _pipeline_device(self, pipe):
        "Device setter for fastai pipeline"
        for tfm in pipe:
            if hasattr(tfm, 'to') and callable(tfm.to):
                tfm.to(self.device, non_blocking=True)
            else:
                for a in L(getattr(tfm, 'parameters', None)):
                    setattr(tfm, a, getattr(tfm, a).to(self.device, non_blocking=True))

    def _callback_warning(self, kwargs):
        for name in ['item_tfms', 'after_item', 'before_batch']:
            if name in kwargs:
                msg = f"fastxtend's `BaseDataLoader` will not call any {name} " \
                      f"methods. {name} is for use with a fastai DataLoader.\n"
                warn(msg)
