# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06a_memory.experience_replay.ipynb.

# %% auto 0
__all__ = ['ExperienceReplay']

# %% ../nbs/06a_memory.experience_replay.ipynb 3
# Python native modules
import os
from typing import *
from warnings import warn
# Third party libs
from fastcore.all import *
from fastai.learner import *
from fastai.torch_basics import *
from fastai.torch_core import *
import torchdata.datapipes as dp
# Local modules
from ..core import *

# %% ../nbs/06a_memory.experience_replay.ipynb 5
class ExperienceReplay(dp.iter.IterDataPipe):
    debug=False
    def __init__(self,
            source_datapipe,
            learn=None,
            bs=1,
            max_sz=100,
            return_idxs=False,
            # If being run with n_workers>0, shared_memory, and fork, this MUST be true. This is needed because
            # otherwise the tensors in the memory will remain shared with the tensors created in the 
            # dataloader.
            clone_detach:bool=False 
        ):
        self.memory = np.array([None]*max_sz)
        self.source_datapipe = source_datapipe
        self.learn = learn
        if learn is not None:
            self.learn.experience_replay = self
        self.bs = bs
        self.max_sz = max_sz
        self._sz_tracker = 0
        self._idx_tracker = 0
        self._cycle_tracker = 0
        self.return_idxs = return_idxs
        self.clone_detach = clone_detach
    
    def sample(self,bs=None):  
        idxs = np.random.choice(range(self._sz_tracker),size=(ifnone(bs,self.bs),),replace=False)
        if self.return_idxs: return self.memory[idxs],idxs
        return self.memory[idxs]
    
    def __repr__(self):
        return str({k:v if k!='memory' else f'{len(self)} elements' for k,v in self.__dict__.items()})

    def __len__(self): return self._sz_tracker
    
    def __iter__(self):
        for i,b in enumerate(self.source_datapipe):
            if self.debug: print('Experience Replay Adding: ',b)
            
            if not issubclass(b.__class__,(StepType,list,tuple)):
                raise Exception(f'Expected typing.NamedTuple,list,tuple object got {type(step)}\n{step}')
            
            if issubclass(b.__class__,StepType):   self.add(b)
            elif issubclass(b.__class__,(list,tuple)): 
                for step in b: self.add(step)
            else:
                raise Exception(f'This should not have occured: {self.__dict__}')
        
            if self._sz_tracker<self.bs: continue
            yield self.sample()

    def add(self,step:StepType): 
        if self.clone_detach: step = step.clone().detach()
        
        if self._sz_tracker==0: 
            self.memory[self._idx_tracker] = step
            self._sz_tracker += 1
            self._idx_tracker = 1
        elif 0<self._sz_tracker<self.max_sz:
            self.memory[self._idx_tracker] = step
            self._sz_tracker += 1
            self._idx_tracker += 1
        elif self._sz_tracker>=self.max_sz:
            if self._idx_tracker>=self.max_sz:
                self._idx_tracker = 0
                self._cycle_tracker += 1
            self.memory[self._idx_tracker] = step
            self._idx_tracker += 1
        else:
            raise Exception(f'This should not have occured: {self.__dict__}')
            
add_docs(
    ExperienceReplay,
    """Simplest form of memory. Takes steps from `source_datapipe` to stores them in `memory`. 
       It outputs `bs` steps.""",
    sample="Returns `bs` steps from `memory` in a uniform distribution.",
    add="Adds new steps to `memory`. If `memory` reaches size `max_sz` then `step` will be added in earlier steps."
)
