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

# %% auto 0
__all__ = ['EpochIterator', 'AsyncEpochIterator']

# %% ../../nbs/ffcv.epoch_iterator.ipynb 1
from threading import Event
from queue import Queue
from contextlib import nullcontext
from typing import Sequence, TYPE_CHECKING

from ffcv.traversal_order.quasi_random import QuasiRandom
from ffcv.utils import chunks

from ffcv.loader.epoch_iterator import EpochIterator as _EpochIterator
from ffcv.loader.epoch_iterator import select_buffer, QUASIRANDOM_ERROR_MSG

from fastcore.transform import Pipeline

from ..imports import *

# %% ../../nbs/ffcv.epoch_iterator.ipynb 2
if TYPE_CHECKING:
    from fastxtend.ffcv.loader import Loader
IS_CUDA = torch.cuda.is_available()

# %% ../../nbs/ffcv.epoch_iterator.ipynb 3
class EpochIterator(_EpochIterator):
    def __init__(self, loader, order: Sequence[int]):
        super(_EpochIterator, self).__init__(daemon=True)
        self.loader: 'Loader' = loader
        self.order = order
        self.metadata = loader.reader.metadata
        self.current_batch_slot = 0
        batches = list(chunks(order, self.loader.batch_size))
        self.iter_ixes = iter(batches)
        self.closed = False
        self.output_queue = Queue(self.loader.batches_ahead)
        self.terminate_event = Event()
        self.memory_context = self.loader.memory_manager.schedule_epoch(
            batches)

        if IS_CUDA:
            self.current_stream = torch.cuda.current_stream()

        try:
            self.memory_context.__enter__()
        except MemoryError as e:
            if loader.traversal_order != QuasiRandom:
                print(QUASIRANDOM_ERROR_MSG)
                print('Full error below:')
            raise e

        self.storage_state = self.memory_context.state
        self.cuda_streams = self.loader.cuda_streams

        self.memory_allocations = self.loader.graph.allocate_memory(
            self.loader.batch_size,
            self.loader.batches_ahead + 2
        )
        self.start()

# %% ../../nbs/ffcv.epoch_iterator.ipynb 4
class AsyncEpochIterator(EpochIterator):
    def __init__(self, loader, order: Sequence[int], after_batch: Pipeline):
        self.after_batch = after_batch
        super().__init__(loader, order)


    def run_pipeline(self, b_ix, batch_indices, batch_slot, cuda_event):
        self.memory_context.start_batch(b_ix)
        args = []
        if IS_CUDA:
            stream = self.cuda_streams[batch_slot]
            ctx = torch.cuda.stream(stream)
        else:
            ctx = nullcontext()

        code, outputs = self.loader.code
        with ctx:
            self.loader._pipeline_device(self.after_batch)

            if IS_CUDA:
                if cuda_event:
                    cuda_event.wait()

            args = {
                'batch_indices': batch_indices,
                'storage_state': self.storage_state,
                'metadata': self.metadata,
                **{
                    f'memory_{k}':select_buffer(v, batch_slot, len(batch_indices))
                    for (k, v) in self.memory_allocations['operation'].items()
                },
                **{
                    f'shared_memory_{k}': select_buffer(v, batch_slot, len(batch_indices))
                    for (k, v) in self.memory_allocations['shared'].items()
                }
            }

            for stage_code, define_outputs in code:
                results = stage_code(**args)
                for node_id, result in zip(define_outputs, results):
                    args[f'result_{node_id}'] = result
                pass

            result = tuple(args[f'result_{x}'] for x in outputs)
            return self.after_batch(result)

    def close(self):
        self.loader._pipeline_device(self.after_batch)
        self.terminate_event.set()
        if not self.closed:
            self.memory_context.__exit__(None, None, None)
