# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/06a_memory.experience_replay.ipynb (unless otherwise specified).

__all__ = ['ExperienceReplayException', 'ExperienceReplay', 'ExperienceReplayCallback']

# Cell
# Python native modules
import os
from typing import *
# Third party libs
from fastcore.all import *
from fastai.learner import *
from fastai.torch_basics import *
from fastai.torch_core import *
from fastai.callback.all import *
# Local modules
from ..core import *
from ..callback.core import *
from ..data.block import *

# Cell
class ExperienceReplayException(Exception): pass

class ExperienceReplay(object):
    def __init__(self,
                 bs=16,         # Number of entries to query from memory
                 max_sz=200,    # Maximum number of entries to hold. Will start overwriting after.
                 warmup_sz=100  # Minimum number of entries needed to continue with a batch
                 ):
        "Stores `BD`s in a rotating list `self.memory`"
        store_attr()
        test_lt(warmup_sz-1,max_sz)
        self.memory=None
        self.pointer=0

    def __add__(self,other:BD):
        "In-place add `other` to memory, overwriting if len(self.memory)>self.max_sz"
        if isinstance(other,tuple) and len(other)==1: other=other[0]
        elif isinstance(other,tuple):                 raise ExperienceReplayException('records need to be `BD`s or 1 element tuples')
        if isinstance(other,dict): other=BD(other)

        if self.memory is None:
            if other.bs()>self.max_sz:
                self.memory=other[:self.max_sz]
                self.pointer=0           # Keep the pointer 0 since we have basically replaced the memory
                self+other[self.max_sz:] # Recursively add the rest of the batch
            else:
                self.memory=other
                self.pointer=self.memory.bs() # remember that pointer is not an index but number of elements
        else:
            n_over=(other.bs()+self.pointer)-self.max_sz
            if n_over>0: # e.g.: max_sz 200, pointer 195, other is 5.
                self.memory=self.memory[:self.pointer]+other[:-n_over]
                self.pointer=0
                self+other[other.bs()-n_over:]
            else:
                # If the number of elements is not over
                next_pointer=self.pointer+other.bs()
                self.memory=self.memory[:self.pointer]+other+self.memory[next_pointer:]
                self.pointer=next_pointer
        return self

    def __radd__(self,other:BD): raise ExperienceReplayException('You can only do experience_reply+[some other element]')

    def __len__(self): return self.memory.bs() if self.memory is not None else 0

    def sample(self)->BD:
        "Returns a sample of size `self.bs`"
        with torch.no_grad():
            idxs=np.random.randint(0,self.memory.bs(),self.bs).tolist()
            samples=self.memory[idxs].mapv(to_device)

        if self.memory.bs()<self.warmup_sz: raise CancelBatchException
        return samples

# Cell
class ExperienceReplayCallback(Callback):
    @delegates(ExperienceReplay)
    def __init__(self,**kwargs):
        "Stores `BD`s in a rotating list `self.memory`"
        store_attr()
        self.experience_replay=ExperienceReplay(**kwargs)

    def after_pred(self):
        "Adds `learn.xb` to memory, then sets `learn.xb=experience_replay.sample()`"
        xb=BD(self.learn.xb[0]).mapv(to_detach)
        self.experience_replay+xb

        self.learn.xb=self.experience_replay.sample()