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

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

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

from pathlib import Path
from typing import Mapping, Sequence

import numpy as np

from ffcv.fields.base import Field
from ffcv.loader.loader import Loader as _Loader
from ffcv.loader.loader import OrderOption, ORDER_TYPE, DEFAULT_OS_CACHE, ORDER_MAP
from ffcv.pipeline.compiler import Compiler
from ffcv.pipeline.operation import Operation
from ffcv.transforms.ops import ToDevice as _ToDevice

from ..data.loader import DataLoaderMixin
from .epoch_iterator import EpochIterator, AsyncEpochIterator
from ..imports import *

# %% auto 0
__all__ = ['Loader', 'OrderOption']

# %% ../../nbs/ffcv.loader.ipynb 5
_all_ = ['OrderOption']

# %% ../../nbs/ffcv.loader.ipynb 7
class Loader(DataLoaderMixin, _Loader):
    "FFCV `Loader` with fastai Transformed DataLoader `TfmdDL` batch transforms"
    def __init__(self,
        fname:str|Path, # Path to the location of the dataset (FFCV beton format)
        batch_size:int, # Batch size
        num_workers:int=-1, # Number of CPU cores to use in parallel (default: All available up to 16)
        os_cache:bool=DEFAULT_OS_CACHE, # Leverage the OS for caching. Beneficial when there is enough memory to cache the dataset
        order:ORDER_TYPE=OrderOption.SEQUENTIAL, # Dataset traversal order, one of: `SEQEUNTIAL`, `RANDOM`, `QUASI_RANDOM`
        distributed:bool=False, # Emulates the behavior of PyTorch's DistributedSampler for distributed training
        seed:int|None=None, # Random seed for batch ordering
        indices:Sequence[int]|None=None, # Subset dataset by returning only these indices
        pipelines:Mapping[str, Sequence[Operation|nn.Module]]={}, # Dictionary defining for each field the sequence of Decoders and transforms to apply
        custom_fields:Mapping[str, Field]={}, # Dictonary informing `Loader` of the types associated to fields that are using a custom type
        drop_last:bool|None=None, # Drop non-full batch in each epoch. Defaults to True if order is `SEQEUNTIAL`
        batches_ahead:int=2, # Number of batches prepared in advance; balances latency and memory
        recompile:bool=False, # Recompile at every epoch. Required if FFCV augmentations change during training
        device:str|int|torch.device|None=None, # Device to place batch. Defaults to fastai's `default_device`
        async_tfms:bool=False, # Asynchronously run `batch_tfms` before batch is drawn.
        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 split_idx is None:
            split_idx = int(order==OrderOption.SEQUENTIAL)

        self.async_tfms = async_tfms and len(kwargs['after_batch'].fs) > 0
        self.cuda_streams = None

        if drop_last is None:
            drop_last != order==OrderOption.SEQUENTIAL

        _Loader.__init__(self,
            fname=str(Path(fname)), batch_size=batch_size, num_workers=num_workers,
            os_cache=os_cache, order=order, distributed=distributed, seed=seed,
            indices=indices, pipelines=pipelines, custom_fields=custom_fields,
            drop_last=drop_last, batches_ahead=batches_ahead, recompile=recompile
        )
        DataLoaderMixin.__init__(self,
            device=device, n_inp=n_inp, split_idx=split_idx, do_setup=do_setup, **kwargs
        )
        if self._n_inp is None:
            self._n_inp = len(pipelines) - 1

    def one_batch(self, batches_ahead:bool=False):
        "Return one processed batch of input(s) and target(s), optionally loading `batches_ahead`"
        for b in self._n_batches(self.batches_ahead + 2 if batches_ahead else 1):
            # need to return the yield from _n_batches so `Loader` can reset to iterate the entire epoch
            pass
        return b

    @property
    def bs(self) -> int:
        "Number of items a batch"
        return self.batch_size

    @property
    def device(self) -> torch.device:
        return super().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 Loader.batch_tfms
        if hasattr(self.after_batch, 'fs'):
            self._pipeline_device(self.after_batch.fs)
        # Device setter for FFCV Pipeline
        for p in self.pipeline_specs.values():
            for t in p.transforms:
                if isinstance(t, _ToDevice):
                    t.device = self._device

    def _iter(self):
        Compiler.set_num_threads(self.num_workers)
        order = self.next_traversal_order()
        selected_order = order[:len(self) * self.batch_size]
        self.next_epoch += 1

        # Compile at the first epoch
        if self.code is None or self.recompile:
            self.generate_code()

        # Asynchronous transforms require using the same Cuda streams for the entire run
        if self.cuda_streams is None:
            self.cuda_streams = [(torch.cuda.Stream() if torch.cuda.is_available() else None)
                                  for _ in range(self.batches_ahead + 2)]
        if self.async_tfms:
            return AsyncEpochIterator(self, selected_order, self.after_batch)
        else:
            return EpochIterator(self, selected_order)

    def __iter__(self):
        self.before_iter()
        if self.async_tfms:
            yield from self._iter()
        else:
            for b in self._iter():
                yield self.after_batch(b)
        self.after_iter()
        if hasattr(self, 'it'):
            del(self.it)

    def _n_batches(self, num_batches:int=1):
        orig_traversal_order = self.traversal_order
        orig_indices = self.indices
        orig_drop_last = self.drop_last

        # Set Loader to only return one batch per epoch
        if self._args['order'] == OrderOption.SEQUENTIAL:
            self.indices = np.arange(0, self.batch_size*num_batches)
        else:
            self.indices = np.random.random_integers(0, self.reader.num_samples, self.batch_size*num_batches)
        self.traversal_order = ORDER_MAP[OrderOption.SEQUENTIAL](self)
        self.drop_last = False

        # yield num_batches
        yield from self.__iter__()

        # Reset Loader state to its original status
        self.next_epoch -= 1
        self.indices = orig_indices
        self.drop_last = orig_drop_last
        self.traversal_order = orig_traversal_order

    def _callback_warning(self, kwargs):
        for name in ['item_tfms', 'after_item', 'before_batch']:
            if name in kwargs:
                if name != 'before_batch':
                    msg = f"fastxtend's `Loader` will not call any {name} methods. " \
                          f"{name} is for use with a fastai DataLoader.\n" \
                          f"Instead of passing fastai Item Transforms to {name}," \
                          f"initialize the fastxtend `Loader` pipeline with FFCV transforms."
                else:
                    msg = f"fastxtend's `Loader` will not call any {name} methods. " \
                          f"{name} are for use with a fastai DataLoader."
                warn(msg)
