# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01e_pipes.iter.nstep.ipynb.

# %% auto 0
__all__ = ['NStepper', 'n_steps_expected']

# %% ../nbs/01e_pipes.iter.nstep.ipynb 3
# Python native modules
import os
# Third party libs
from fastcore.all import *
import torchdata.datapipes as dp
import typing
from fastai.torch_basics import *
from fastai.torch_core import *
# Local modules
from ...core import *
from ..core import *
from ...fastai.data.block import *
from ..core import *

# %% ../nbs/01e_pipes.iter.nstep.ipynb 5
class NStepper(dp.iter.IterDataPipe):
    def __init__(self, source_datapipe, n=1) -> None:
        self.source_datapipe = source_datapipe
        self.n = n
        self.env_buffer = {}
        
    def __iter__(self) -> typing.Tuple[typing.NamedTuple]:
        self.env_buffer = {}
        for step in self.source_datapipe:
            if not issubclass(step.__class__,StepType):
                raise Exception(f'Expected typing.NamedTuple object got {type(step)}\n{step}')
    
            env_id,terminated = int(step.env_id),bool(step.terminated)
        
            if env_id in self.env_buffer:
                self.env_buffer[env_id].append(step)
            else:
                self.env_buffer[env_id] = [step]
                
            if not terminated and len(self.env_buffer[env_id])<self.n: continue
            
            while terminated and len(self.env_buffer[env_id])!=0:
                yield tuple(self.env_buffer[env_id])
                self.env_buffer[env_id].pop(0)
                
            if not terminated:
                yield tuple(self.env_buffer[env_id])
                self.env_buffer[env_id].pop(0)
add_docs(
    NStepper,
    """Accepts a `source_datapipe` or iterable whose `next()` produces a `typing.NamedTuple` of 
       max size `n` that will contain steps from a single environment with 
       a subset of fields from `SimpleStep`, namely `terminated` and `env_id`.""",
)

# %% ../nbs/01e_pipes.iter.nstep.ipynb 17
def n_steps_expected(
    default_steps:int, # The number of steps the episode would run without n_steps
    n:int # The n-step value that we are planning ot use
):
    return (default_steps * n) - sum(range(n))
    
n_steps_expected.__doc__=r"""
Produces the expected number of steps, assuming a fully deterministic episode based on `default_steps` and `n`

Given `n=2`, given 1 envs, knowing that `CartPole-v1` when `seed=0` will always run 18 steps, the total 
steps will be:

$$
18 * n - \sum_{0}^{n - 1}(i)
$$
"""    
