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

# %% ../../nbs/experimental/mp.ipynb 1
from __future__ import print_function, division, annotations
from ..imports import *
from ..datasets import ArrayDataset, JAXDataset
from ..loaders import BaseDataLoader
from ..utils import get_config, asnumpy
from ..tests import *
import jax_dataloader as jdl
from threading import Thread, Event
from queue import Queue, Full
import multiprocessing as mp
import weakref

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

# %% ../../nbs/experimental/mp.ipynb 2
def chunk(seq: Sequence, size: int) -> List[Sequence]:
    return [seq[pos:pos + size] for pos in range(0, len(seq), size)]  


# %% ../../nbs/experimental/mp.ipynb 3
class EpochIterator(Thread):
    """[WIP] Multiprocessing Epoch Iterator"""
    
    def __init__(self, data, batch_size: int, indices: Sequence[int]):
        super().__init__()
        self.data = data
        batches = chunk(indices, batch_size)
        self.iter_idx = iter(batches)
        self.output_queue = Queue(5) # TODO: maxsize
        self.terminate_event = Event()
        self.start()

    def run(self):
        try:
            while True:
                # get data
                result = self.get_data()
                # put result in queue
                while True:
                    try: 
                        self.output_queue.put(result, block=True, timeout=0.5)
                        break
                    except Full: pass
                
                    if self.terminate_event.is_set(): return                

        except StopIteration:
            self.output_queue.put(None)

    def __next__(self):
        result = self.output_queue.get()
        if result is None:
            self.close()
            raise StopIteration()
        return result
    
    def __iter__(self):
        return self
    
    def __del__(self):
        self.close()

    def close(self):
        self.terminate_event.set()

    def get_data(self):
        batch_idx = next(self.iter_idx)
        batch = self.data[batch_idx]
        return batch

