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

__all__ = ['DQN', 'init_experience', 'cast_dtype', 'FakeAgent', 'ExperienceSource', 'SourceDataset',
           'FirstLastExperienceSource']

# Cell
# Python native modules
import os
from collections import deque
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
# Local modules
from ..core import *

# Cell
class DQN(Module):
    def __init__(self):
        self.policy=nn.Sequential(
            nn.Linear(4,50),
            nn.ReLU(),
            nn.Linear(50,2),
            nn.ReLU()
        )

    def forward(self,x):
        return torch.argmax(self.policy(x),dim=0)

# Cell
def init_experience(but='',**kwargs):
    "Returns dictionary with default values that can be overridden."
    experience=D(
        state=0,action=0,next_state=0,reward=0,done=False,
        step=0,env=0,image=0
    )
    for s in but.split(','):
        if s in experience: del experience[s]
    return BD.merge(experience,kwargs)

# Cell
def _state2experience(s,**kwargs):   return init_experience(state=s,next_state=s,step=torch.zeros((1,1)),**kwargs)
def _env_reset(o):                   return o.reset()
def _env_seed(o,seed):               return o.seed(seed)
def _env_render(o,mode='rgb_array'): return TensorBatch(o.render(mode=mode).copy())
def _env_step(o,*args,**kwargs):     return o.step(*args,**kwargs)

def cast_dtype(t,dtype):
    if dtype==torch.float:    return t.float()
    elif dtype==torch.double: return t.double()
    elif dtype==torch.long:   return t.long()

class FakeAgent:
    def __init__(self,action_space): store_attr()
    def __call__(self,state,**kwargs):
        return (L([self.action_space.sample() for _ in range(state.shape[0])]),
                D(merge(kwargs,{'random_action':np.random.randint(0,3,(state.shape[0],1))})))

class ExperienceSource(Stateful):
    _stateattrs=('pool',)
    def __init__(self,env:str,agent=None,n_envs:int=1,steps_count:int=1,steps_delta:int=1,
                 seed:int=None,render=None,num_workers=0,but='',**kwargs):
        store_attr()
        self.env_kwargs=kwargs
        self.pool=L()
        if self.render is None: self.but+=',image'

    def _init_state(self):
        "Inits the histories, experiences, and the environment pool when sent to a `Process`"
        self.history,self.pool=L((deque(maxlen=self.steps_count),
                                  gym.make(self.env,**self.env_kwargs))
                                  for _ in range(self.n_envs)).zip().map(L)
        self.pool.map(_env_seed,seed=self.seed)
        if self.agent is None: self.agent=FakeAgent(self.pool[0].action_space)
        self.reset_all()

    def reset_all(self):
        self.experiences=self.pool.map(_env_reset)
        self.experiences=self.experiences.map(_state2experience,but=self.but)
        self.experiences=sum(self.experiences)
        self.attempt_render(self.experiences)

    def attempt_render(self,experiences,indexes=None):
        if self.render is not None:
            pool=self.pool if indexes is None else self.pool[indexes]
            renders=pool.map(_env_render,mode=self.render)
            # No idea why we have to do this, but multiprocessing hangs forever otherwise
            if self.num_workers>0:sleep(0.1)
            experiences['image']=torch.stack(tuple(renders)).unsqueeze(0)

    def __iter__(self):
        "Iterates through a list of environments."
        if not self.pool:self._init_state()
        while True:
#             try:
                # Only work on envs that are not done
            not_done_idxs=(self.experiences['done']==False).nonzero()[:,0]
            if len(not_done_idxs)==0:
                self.reset_all()
                not_done_idxs=(self.experiences['done']==False).nonzero()[:,0]
            not_done_idxs=not_done_idxs.reshape(-1,)
            not_done_experiences=self.experiences[not_done_idxs]
            # Pass current experiences into agent
            actions,experiences=self.agent(**not_done_experiences)
            # Step through all envs.
            step_res=self.pool[not_done_idxs].zipwith(actions).starmap(_env_step)
            next_states,rewards,dones=step_res.zip()[:3].map(TensorBatch)
            rewards,dones=(v.reshape(len(not_done_idxs),-1) for v in (rewards,dones))
            # Add the image field if available
            self.attempt_render(self.experiences,not_done_idxs)
            new_exp=BD(next_state=next_states,reward=rewards,done=dones,
                       env=not_done_idxs.reshape(not_done_experiences.bs,-1),
                       step=not_done_experiences['step']+1)

            experiences=BD.merge(not_done_experiences,experiences,new_exp)
            for i,idx in enumerate(not_done_idxs):
#                 print(experiences,idx)
                self.history[idx].append(experiences[i])
                if len(self.history[idx])==self.steps_count and \
                       int(experiences[i]['step'][0])%self.steps_delta==0:
#                     print(self.history[idx])
                    yield tuple(self.history[idx])

                if bool(experiences[i]['done'][0]):
                    if 0<len(self.history[idx])<self.steps_count:
                        yield tuple(self.history[idx])
                    while len(self.history[idx])>1:
                        self.history[idx].popleft()
                        yield tuple(self.history[idx])
            experiences['state']=experiences['next_state']

            for k in experiences:
                dtype=experiences[k].dtype
                if k not in self.experiences:
                    self.experiences[k]=TensorBatch(torch.zeros(self.experiences.bs,
                                                    *experiences[k].shape[1:]))
                if self.experiences[k].dtype!=dtype:
                    self.experiences[k]=cast_dtype(self.experiences[k],dtype)
                self.experiences[k][not_done_idxs]=experiences[k]
#             except ValueError:
#                 self.reset_all()

add_docs(ExperienceSource,
        """Iterates through `n_envs` of `env` feeding experience or states into `agent`.
           If `agent` is None, then random actions will be taken instead.
           It will return `steps_count` experiences every `steps_delta`.
           At the end of an env, it will return `steps_count-1` experiences per next. """,
        reset_all="resets the envs and experience",
        attempt_render="Updates `experiences` with images if `render is not None`. Optionally indexes can be passed.")

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

# Cell
class FirstLastExperienceSource(ExperienceSource):
    gamma=0.99
    def __iter__(self):
        for res in super().__iter__():
            element,remainder=res[0],() if len(res)==1 else res[1:]
            reward=element['reward']
            for e in reversed(remainder):
                reward*=self.gamma
                reward+=e['reward']
            element.bs=1
#             print(element,element.bs)
            yield (element,)