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

__all__ = ['TstCallback', 'source_events', 'Src', 'SrcCallback', 'GymSrc', 'SrcProfiler', 'FirstLast', 'SourceDataset',
           'IterableTfmdLists', 'vstack_collate', 'fr_collate', 'IterableTfmdDL', 'IterableDataBlock', 'TransformBlock',
           'SourceBlock', 'map_iters', 'noop_splitter', 'get_all_keys', '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 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
class TstCallback(AgentCallback):
    def __init__(self,action_space=None,constant=None): store_attr()
    def before_noise(self):
        bs=self.experience['state'].shape[0]
        self.agent.action=Tensor([[ifnone(self.constant,self.action_space.sample())] for _ in range(bs)])
        self.agent.experience=D(merge(self.experience,{'random_action':np.random.randint(0,3,(bs,1))}))

# Cell
_events=L.split('create reset do_action do_step render history episodes initialize')
_events= _events.map(lambda s:'before_'+s) + _events.map(lambda s:'after_'+s)+_events.map(lambda s:'after_cancel_'+s)
_events+=L.split('reset do_step render initialize history do_action')

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

#nbdev_comment _all_=['source_events']

_loop=L(['Start Setup','before_initialize','initialize','after_initialize','End Setup',
             'before_episodes',
             'Start Episodes','before_reset','reset','after_reset',
                 'before_do_action','do_action','after_do_action',
                 'before_do_step','do_step','after_do_step',
                 'before_render','render','after_render',
                 'before_history','history','after_history',
             'End Episodes',
             'after_episodes'
             ])
_default_shapes={'done':(1,1),'all_done':(1,1),'step':(1,1),'env_id':(1,1),
                 'p_id':(1,1),'episode_num':(1,1),'accum_rewards':(1,1)}

# Cell
class Src(Loop):
    _loop=_loop
    _events=source_events
    _default='source'
    end_event='after_episodes'
    initial_experience_flds='done step state next_state env_id'

    @delegates(Loop)
    def __init__(self,env:object,agent=None,n_envs:int=1,steps_count:int=1,steps_delta:int=1,
                 seed:int=None,mode=None,num_workers=0,but='',cbs=None,shape_map=None,**kwargs):
        super().__init__(cbs=cbs,**kwargs)
        store_attr(but='cbs,shape_map')
        self.env_kwargs=kwargs
        self.pool=L()
        self.histories=L()
        self.current_history=L()
        self.shape_map=merge(_default_shapes,ifnone(shape_map,{}))
        self.all_exp:BD=None
        self.imask=torch.zeros((1,))
        if self.mode is None: self.but+=',image'
        # Fields
        self.done=None
        self.all_done=None
        self.step=None
        self.env_id=None
        self.p_id=None

    def _init_state(self):     self('initialize')
    def _do_reset(self):       self('reset')
    def _do_action(self):      self('do_action')
    def _do_step(self):        self('do_step')
    def _do_render(self):      self('render')
    def process_history(self,history):
        self.current_history=history
        self('history')
        return self.current_history

    def yield_histories(self):
        self('before_history')
        for i,idx in enumerate(self.imask.nonzero().reshape(-1,)):
            self.histories[idx].append(deepcopy(self.active_exp()[i]))
            if len(self.histories[idx])==self.steps_count and int(self.step[i])%self.steps_delta==0 and not bool(self.done.reshape(-1,)[i]):
                yield self.process_history(sum(self.histories[idx]))
            if bool(self.done.reshape(-1,)[i]):
#                 print(self.histories[idx],'\n')
#                 if 0<len(self.histories[idx])<self.steps_count:
#                     yield self.process_history(sum(self.histories[idx]))
                while len(self.histories[idx])>=1:
                    yield self.process_history(sum(self.histories[idx]))
                    self.histories[idx].popleft()
        self('after_history')

    def active_exp(self,bs=None):
        return BD({k:TensorBatch(getattr(self,k),bs=ifnone(bs,self.imask.sum()))
                   for k in self.shape_map})

    def __iter__(self):
        "Iterates through a list of environments."
        if not self.pool:self._with_events(self._init_state,'initialize',Exception)
        self('before_episodes')
        while True:
            self._with_events(self._do_reset,'reset',Exception)
            self._with_events(self._do_action,'do_action',Exception)
            self._with_events(self._do_step,'do_step',Exception)
            if self.mode is not None: self._with_events(self._do_render,'render',Exception)
            yield from self.yield_histories()
        self('after_episodes')

# Cell
class SrcCallback(LoopCallback):
    _default='source'
    _methods=_events

    def after_cancel_initialize(self):raise
    def after_cancel_step(self):      raise
    def after_cancel_do_action(self): raise
    def after_cancel_do_step(self):   raise
    def after_cancel_episodes(self):  raise
    def after_cancel_episode(self):   raise
    def after_cancel_history(self):   raise
    def after_cancel_render(self):    raise
    def after_cancel_create(self):    raise
    def after_cancel_reset(self):     raise

# Cell
def _add_batch(t):
    if len(t)==1:
        if type(t[0])==float: return (1,1)
        else:                 return (1,t[0])
    elif t[0]!=1: return (1,*t)
    return t

def _fix_shape_map(t:Tensor,expected_shape,bs=None):
    if not hasattr(t,'shape'): t=Tensor([t])
    if bs is not None:
        expected_shape=(bs,*expected_shape[1:])
    if expected_shape!=tuple(t.shape):
        if isinstance(t,np.ndarray): t=Tensor(t)
        t=t.unsqueeze(0)
        test_eq(expected_shape,t.shape)
    return t

def _env_seed(o,seed): return o.seed(seed)

def _env_reset(o,shape_map=None):
    state=o.reset()
    return _fix_shape_map(state,shape_map['state'])

def _env_render(o,mode='rgb_array',shape_map=None):
    image=o.render(mode=mode).copy()
    if shape_map is None: return image
    return _fix_shape_map(image,shape_map['image'])

def _env_step(o,action,*args,shape_map=None,is_int=False,**kwargs):
    a=action.detach().cpu().numpy()[0]
    if np.isscalar(a):
        a=int(a) if is_int else float(a)
    s,r,d,info=o.step(a,*args,**kwargs)
    s=_fix_shape_map(s,shape_map['state'])
    r=_fix_shape_map(r,shape_map['reward'])
    d=_fix_shape_map(d,shape_map['done'])
    return (s,r,d,info)

def _batchwise_zero(shape,bs): return torch.zeros((bs,*shape[1:]))

class GymSrc(SrcCallback):
    def initialize(self):
        self.source.histories,self.source.pool=L((deque(maxlen=self.steps_count),
                          gym.make(self.env,**self.env_kwargs))
                          for _ in range(self.n_envs)).zip().map(L)
        if self.source.agent is None:
            test_cb=TstCallback(action_space=self.pool[0].action_space)
            self.source.agent=Agent(cbs=test_cb)
        self.action_is_int=False
       # Extra fields
        self.source.next_state=None
        self.source.state=None
        self.source.image=None
        self.source.action=None
        self.source.episode_num=None
        self.source.accum_rewards=None
        self.init_shapes()
        self('reset')
        if self.mode is not None: self.init_render_shapes()
        self.source.imask=torch.zeros((self.n_envs,)).bool()

    def init_render_shapes(self):
        "Set the image shapes with a batch dim."
        image_shape=_env_render(self.pool[0],self.mode).shape
        self.shape_map['image']=_add_batch(image_shape)

    def init_shapes(self):
        "Set the reward shape, state shapes."
        self.shape_map['reward']=(1,1)
        obs_shape=_add_batch(self.pool[0].observation_space.shape)
        for k in ('state','next_state'): self.shape_map[k]=obs_shape
        action=self.pool[0].action_space.sample()
        if isinstance(action,(int,float)):
            self.shape_map['action']=(1,1)
            if isinstance(action,int): self.action_is_int=True
        else:
            self.shape_map['action']=_add_batch(action)

    def reset(self):
        if self.imask.sum()==0:
            _shape_map={k:v for k,v in self.shape_map.items() if k!='episode_num'}
            reset_exps=D(_shape_map).mapv(_batchwise_zero,bs=self.n_envs)
            for k,v in reset_exps.items(): setattr(self.source,k,v)
            self.source.pool.map(_env_seed,seed=self.seed)
            self.source.state=self.pool.map(_env_reset,shape_map=self.shape_map)
            self.source.state=TensorBatch.vstack(tuple(self.state))
            self.source.done=self.source.done.bool()
            self.source.all_done=self.source.all_done.bool()
            self.source.env_id=self.source.env_id.long()
            self.source.p_id=self.source.p_id.long()
            if self.mode is not None:
                self.source.image=self.pool.map(_env_render,mode=self.mode)
            if self.action_is_int:
                self.source.action=self.source.action.long()

    def after_reset(self):
        if self.imask.sum()==0:
            if self.source.episode_num is None:
                self.source.episode_num=self.source.env_id.detach().clone()+1
            else:
                self.source.episode_num+=self.source.episode_num.max()+1-self.source.episode_num
            self.source.all_exp=deepcopy(self.active_exp(self.n_envs))
            self.source.imask=torch.ones((self.n_envs,)).bool()

    def do_action(self):
        self.source.action,exp=self.agent.do_action(**self.active_exp())
        # TODO: Im not sure we need this section, maybe continuous actions need this here though?
#         self.source.action=_fix_shape_map(self.action,self.shape_map['action'],bs=int(self.imask.sum()))
        if not isinstance(self.action,TensorBatch):
            self.source.action=TensorBatch(self.source.action)
        if self.action_is_int:
            self.source.action=self.source.action.long()
        for k in exp:
            if k=='action': continue
            setattr(self.source,k,exp[k])

    def do_step(self):
        step_res=self.pool[self.imask].zipwith(self.action).starmap(_env_step,
                                                                    shape_map=self.shape_map,
                                                                    is_int=self.action_is_int)
        next_states,rewards,dones=step_res.zip()[:3]
        if self.mode is not None:
            self.source.image=self.pool[self.imask].map(_env_render,mode=self.mode)
        self.source.next_state=TensorBatch.vstack(next_states)
        self.source.reward=TensorBatch.vstack(rewards)
        self.source.accum_rewards+=TensorBatch.vstack(rewards)
        self.source.done=TensorBatch.vstack(dones).bool()
        running_mask=self.imask.nonzero().reshape(-1,1)
        self.source.env_id=TensorBatch(running_mask,bs=self.imask.sum()).long()
        worker_id=get_worker_info()
        worker_id=worker_id.id if worker_id is not None else 0
        self.source.p_id=TensorBatch(torch.full(running_mask.shape,worker_id),bs=self.imask.sum()).long()

    def history(self):
        if self.current_history['all_done'].shape[0]==1 and self.current_history['done'][0][0]:
#             print(self.current_history)
            self.current_history['all_done'][0][0]=True

    def after_history(self):
        active_exp=self.active_exp()
        for k in self.all_exp:
            if k in ['episode_num']:continue
            self.all_exp[k][self.imask]=torch.clone(active_exp[k])
        self.source.imask=~self.all_exp['done'].reshape(-1,)
        for k in self.all_exp:
            if k in ['episode_num']:continue
            setattr(self.source,k,self.all_exp[k][self.imask])
        self.source.state=torch.clone(self.next_state)
        self.source.step+=1

# Cell
class SrcProfiler(SrcCallback):
    "Checks time between callbacks."
    def __init__(self):
        self.callback_time:Dict={}
        self.callback_deltas:Dict={}

    def before_set(self,callback_name):
        self.callback_time[callback_name]=datetime.now()
    def after_set_delta(self,callback_name):
        delta=(self.callback_time[callback_name]-datetime.now())
        seconds=0 if delta.seconds==86399 else delta.seconds
        if callback_name not in self.callback_deltas:self.callback_deltas[callback_name]=0
        self.callback_deltas[callback_name]+=seconds+delta.microseconds/1000000

    def before_initialize(self): self.before_set('initialize')
    def after_initialize(self):  self.after_set_delta('initialize')
    def before_reset(self): self.before_set('reset')
    def after_reset(self):  self.after_set_delta('reset')
    def before_do_step(self): self.before_set('do_step')
    def after_do_step(self):  self.after_set_delta('do_step')
    def before_render(self): self.before_set('render')
    def after_render(self):  self.after_set_delta('render')
    def before_history(self): self.before_set('history')
    def after_history(self):  self.after_set_delta('history')
    def before_do_action(self): self.before_set('do_action')
    def after_do_action(self):  self.after_set_delta('do_action')

    def relative_time(self):
        max_v=max(self.callback_deltas.values())
        return D(self.callback_deltas).mapv(lambda o:o/max_v)

# Cell
class FirstLast(SrcCallback):
    gamma=0.99
    def history(self):
        element=self.current_history[0]
        if self.current_history.bs()!=1:
            remainder=self.current_history[1:]
            reward=element['reward']
            for e in reversed(remainder['reward']):
                reward*=self.gamma
                reward+=e
            element['reward']=reward
            element['next_state']=self.current_history[-1]['next_state']
            element['done']=self.current_history[-1]['done']
            self.source.current_history=element
        else:
            self.source.current_history=element

# Cell
class SourceDataset(IterableDataset):
    "Iterates through a `source` object. Allows for re-initing source connections when `num_workers>0`"
    def __init__(self,source=None):
        if not is_listy(source): self.source=source
        elif is_iter(source[0]): self.source=list(map(iter,self.source))
        else:                    self.source=source

    def __iter__(self):          return iter(self.source)
    def wif(self):               self.source._init_state()

# Cell
class IterableTfmdLists(TfmdLists):
    def _after_item(self, o): return self.tfms(next(o))

# 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.all import *

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

# Cell
class IterableDataBlock(DataBlock):
    tl_type = TfmdLists

    @delegates(DataBlock)
    def __init__(self,blocks=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
        super().__init__(blocks=blocks, **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 Datasets(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)

# Cell
@patch
def __iter__(self:Datasets):
    for i in cycle(range(len(self))): yield self[i]

@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))

# Cell
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
def _load_wif(f): return f

class SourceBlock(TransformBlock):
    def __init__(self,source,dls_kwargs=None,**kwargs):
        dls_kwargs=ifnone(dls_kwargs,{})
        dls_kwargs['indexed']=False
        dls_kwargs['shuffle']=False
        dls_kwargs['wif']=partial(source._init_state)
        dls_kwargs['persistent_workers']=True
        super().__init__(dls_kwargs=dls_kwargs,dl_type=IterableTfmdDL,tl_type=IterableTfmdLists,**kwargs)

def map_iters(_,ls): return list(map(iter,ls))
def noop_splitter(_,o): return [[0]]
def get_all_keys(_,o): return o

class SourceDataBlock(IterableDataBlock):
    get_items=map_iters
    splitter=noop_splitter
    get_y=get_all_keys

    @delegates(DataBlock)
    def __init__(self,blocks=None,n_inp=None,**kwargs):
        if not is_listy(blocks) or len(blocks)==1: n_inp=0
        super().__init__(blocks=blocks,n_inp=n_inp,**kwargs)