# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05b_data.block_simple.ipynb (unless otherwise specified).

__all__ = ['parse_events', 'source_events', 'Source', 'SourceDataset', 'VerboseChecked', 'IterableTfmdLists',
           'IterableDatasets', 'IterableDataBlock', 'TransformBlock', 'Source', 'vstack_collate', 'fr_collate',
           'IterableTfmdDL', 'get_sources', 'SourceDataBlock']

# Cell
# Python native modules
import os
from collections import deque
from copy import deepcopy
from time import sleep
# Third party libs
from fastcore.all import *
from fastai.torch_basics import *
from fastai.data.all import *
from fastai.basics import *
from fastai.callback.all import *
from torch.utils.data import Dataset
from torch import nn
import torch
import gym
import numpy as np
# Local modules
from ..core import *
from ..callback.core import *
from ..agent import *

# Cell
def parse_events(loop:L):
    return loop.filter(lambda s:in_('event.',s)).map(lambda s:s.replace('event.',''))

# Cell
_loop=L(['event.after_create','Start Setup','event.initialize','End Setup',
             'event.before_episodes',
             'Start Episodes',
                 'event.reset',
                 'event.do_action',
                 'event.do_step',
                 'event.render',
                 'event.history',
             'End Episodes',
             'event.after_episodes'
             ])

mk_class('source_events', **parse_events(_loop).map_dict(),
         doc="All possible events as attributes to get tab-completion and typo-proofing")

#nbdev_comment _all_=['source_events']

# Cell
class Source(Loop):
    _loop=_loop
    _events=source_events
    _default='source'
    end_event=parse_events(_loop)[-1]

    @delegates(Loop)
    def __init__(self,cbs=None,**kwargs):
        super().__init__(cbs=cbs,**kwargs)
        store_attr(but='cbs')
        self.idx=0

    def after_create(self):
        self('initialize')
        return self

    def __iter__(self):
        while True:
            self.idx+=1
            self('do_step')
            yield {'this':torch.full((1,5),self.idx)}

# Cell
class SourceDataset(IterableDataset):
    "Iterates through a `source` object. Allows for re-initing source connections when `num_workers>0`"
    def __init__(self,source=None): self.source=source
    def __iter__(self):
        source=iter(self.source)
        yield from source

# Cell
class VerboseChecked(LoopCallback):
    _methods=source_events
    def __init__(self,show_worker_id=True,show_env_id=True):store_attr()

    def initialize(self):
        worker_id=get_worker_info()
        worker_id=worker_id.id if worker_id is not None else 0
        if self.show_worker_id: print('Worker id: ',worker_id)

        self.env=gym.make('HumanoidPyBulletEnv-v0')
        self.env.reset()
    def do_step(self):
        if self.show_env_id: print('Env Id: ',id(self.env))

# Cell
from itertools import cycle

class IterableTfmdLists(TfmdLists):
    def __iter__(self): return cycle(self[i] for i in range(len(self)))

class IterableDatasets(Datasets):
    def __iter__(self): return cycle(self[i] for i in range(len(self)))

# Cell
class IterableDataBlock(DataBlock):
    tl_type = TfmdLists
    datasets_type = Datasets

    @delegates(DataBlock)
    def __init__(self,blocks=None,datasets_type=None,n_inp=None,**kwargs):
        blocks = L(self.blocks if blocks is None else blocks)
        for b in blocks:
            if getattr(b, 'tl_type', None) is not None: self.tl_type = b.tl_type
        if datasets_type is not None: self.datasets_type=datasets_type
        if (not is_listy(blocks) or len(blocks)==1) and n_inp is not None: n_inp=0
        super().__init__(blocks=blocks,n_inp=n_inp, **kwargs)

    def datasets(self, source, verbose=False):
        self.source = source                     ; pv(f"Collecting items from {source}", verbose)
        items = (self.get_items or noop)(source) ; pv(f"Found {len(items)} items", verbose)
        splits = (self.splitter or RandomSplitter())(items)
#         pv(f"{len(splits)} datasets of sizes {','.join([str(len(s)) for s in splits])}", verbose)
        return self.datasets_type(items, tfms=self._combine_type_tfms(), splits=splits, dl_type=self.dl_type,
                                  n_inp=self.n_inp, verbose=verbose, tl_type=self.tl_type)

@patch
def __init__(self:Datasets, items=None, tfms=None, tls=None, n_inp=None, dl_type=None,tl_type=TfmdLists, **kwargs):
    super(Datasets,self).__init__(dl_type=dl_type)
    self.tls = L(tls if tls else [tl_type(items, t, **kwargs) for t in L(ifnone(tfms,[None]))])
    self.n_inp = ifnone(n_inp, max(1, len(self.tls)-1))

class TransformBlock():
    "A basic wrapper that links defaults transforms for the data block API"
    def __init__(self, type_tfms=None, item_tfms=None, batch_tfms=None, dl_type=None, dls_kwargs=None, tl_type=None):
        self.type_tfms  =            L(type_tfms)
        self.item_tfms  = ToTensor + L(item_tfms)
        self.batch_tfms =            L(batch_tfms)
        self.dl_type,self.dls_kwargs = dl_type,({} if dls_kwargs is None else dls_kwargs)
        self.tl_type = tl_type

# Cell
class Source(Loop):
    _loop=_loop
    _events=source_events
    _default='source'
    end_event=parse_events(_loop)[-1]

    @delegates(Loop)
    def __init__(self,cbs=None,**kwargs):
        super().__init__(cbs=cbs,**kwargs)
        store_attr(but='cbs')
        self.idx=0

    def after_create(self):
        self('initialize')
        print('init')
        return self

    def __iter__(self):
        while True:
            self.idx+=1
            self('do_step')
            yield {'this':torch.full((1,5),self.idx)}

# Cell
from torch.utils.data._utils.collate import default_collate_err_msg_format
from torch.utils.data._utils.collate import *

def vstack_collate(batch):
    "99% similar to default_collate, however vstacks tensors thus assuming they already have a batch dim"

    elem = batch[0]
    elem_type = type(elem)
    if isinstance(elem, torch.Tensor):
        out = None
        if torch.utils.data.get_worker_info() is not None:
            # If we're in a background process, concatenate directly into a
            # shared memory tensor to avoid an extra copy
            numel = sum([x.numel() for x in batch])
            storage = elem.storage()._new_shared(numel)
            out = elem.new(storage)
        return torch.vstack(batch, out=out)
    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
            and elem_type.__name__ != 'string_':
        if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':
            # array of string classes and object
            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                raise TypeError(default_collate_err_msg_format.format(elem.dtype))

            return vstack_collate([torch.as_tensor(b) for b in batch])
        elif elem.shape == ():  # scalars
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float64)
    elif isinstance(elem, int_classes):
        return torch.tensor(batch)
    elif isinstance(elem, string_classes):
        return batch
    elif isinstance(elem, container_abcs.Mapping):
        return {key: vstack_collate([d[key] for d in batch]) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
        return elem_type(*(default_collate(samples) for samples in zip(*batch)))
    elif isinstance(elem, container_abcs.Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError('each element in list of batch should be of equal size')
        transposed = zip(*batch)
        return [vstack_collate(samples) for samples in transposed]

    raise TypeError(default_collate_err_msg_format.format(elem_type))

# Cell
_collate_types = (ndarray, Tensor, typing.Mapping, str)

def fr_collate(t):
    "A replacement for PyTorch `default_collate` which maintains types and handles `Sequence`s"
    b = t[0]
    return (vstack_collate(t) if isinstance(b, _collate_types)
            else type(t[0])([fr_collate(s) for s in zip(*t)]) if isinstance(b, Sequence)
            else vstack_collate(t))

# Cell

from fastai.data.load import _loaders

class IterableTfmdDL(TfmdDL):
    def __init__(self, dataset,bs=64, shuffle=False, num_workers=None, verbose=False, do_setup=True,**kwargs):
        super().__init__(dataset, bs=bs, shuffle=shuffle, num_workers=num_workers,**kwargs)
        self._loader=None

    def create_batch(self, b): return (fr_collate,fa_convert)[self.prebatched](b)

# Cell
def get_sources(_,ls): return [o.after_create() for o in ls]

class SourceDataBlock(IterableDataBlock):
    datasets_type=IterableDatasets
    get_items=get_sources
    blocks=TransformBlock(
        type_tfms=[lambda o:iter(o)],
        item_tfms=lambda o:next(o),
        tl_type=IterableTfmdLists,
        dl_type=IterableTfmdDL,
        dls_kwargs={'indexed':False,'shuffle':False,'persistent_workers':True})

# Cell
@patch
def _do_epoch_validate(self:Learner,*args,**kwargs): return 0

# Cell
@patch
def after_create(self:Callback):
    for cb in self.learn.cbs:
        if hasattr(cb,'train_metrics'): cb.train_metrics=True


# Cell
@patch
def after_batch(self:Recorder):
    "Update all metrics and records lr and smooth loss in training"
    if len(self.yb) == 0 and len(self.xb) == 0: return
    mets = self._train_mets if self.training else self._valid_mets
    for met in mets: met.accumulate(self.learn)
    if not self.training: return
    self.lrs.append(self.opt.hypers[-1]['lr'])
    self.losses.append(self.smooth_loss.value)
    self.learn.smooth_loss = self.smooth_loss.value
