# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/12l_agents.dqn.asynchronous.ipynb.

# %% auto 0
__all__ = ['ModelSubscriber', 'ModelPublisher', 'DQNLearner', 'CacheLoggerBase', 'DQNAgent']

# %% ../nbs/12l_agents.dqn.asynchronous.ipynb 3
# Python native modules
import os
from collections import deque
# Third party libs
from fastcore.all import *
import torchdata.datapipes as dp
from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta
import torch.multiprocessing as mp
import torch
from torch.nn import *
import torch.nn.functional as F
from torch.optim import *

from fastai.torch_basics import *
from fastai.torch_core import *
from torchdata.dataloader2.graph import find_dps,traverse
# Local modules

from ...core import *
from ..core import *
from ...pipes.core import *
from ...fastai.data.block import *
from ...memory.experience_replay import *
from ..core import *
from ..discrete import *
from ...loggers.core import *
from ...loggers.jupyter_visualizers import *
from ...learner.core import *
from .basic import *
from ...dataloader2_ext import *
from torchdata.dataloader2 import DataLoader2,DataLoader2Iterator

# %% ../nbs/12l_agents.dqn.asynchronous.ipynb 8
class ModelSubscriber(dp.iter.IterDataPipe):
    "If an agent is passed to another process and 'spawn' start method is used, then this module is needed."
    def __init__(self,source_datapipe): 
        super().__init__()
        self.source_datapipe = source_datapipe
        self.model = find_dp(traverse(self.source_datapipe,only_datapipe=True),AgentBase).model
        
    def __iter__(self):
        for x in self.source_datapipe:
            if type(x)==GetInputItemRequest and x.key.startswith('model_state_dict_publish_'):
                self.model.load_state_dict(x.value)
                continue
            yield x

# %% ../nbs/12l_agents.dqn.asynchronous.ipynb 9
class ModelPublisher(dp.iter.IterDataPipe):
    def __init__(self,
            source_datapipe,
            publish_freq:int=1,
            # Sometimes its not possible to share current model due to cuda issues.
            # `do_deepcopy` will copy and move the model to cpu in order to publish it.
            do_deepcopy:bool=False
        ):
        super().__init__()
        self.source_datapipe = source_datapipe
        self.model = find_dp(traverse(self,only_datapipe=True),LearnerBase).model
        self.publish_freq = publish_freq
        self.protocol_clients = []
        self._expect_response = []
        self.initialized = False
        self.do_deepcopy = do_deepcopy
 
    def _reset(self):
        for dl in find_dp(traverse(self,only_datapipe=True),LearnerBase).iterable:
            for q_wrapper in dl.datapipe.iterable.datapipes:
                self.protocol_clients.append(q_wrapper.protocol)
                self._expect_response.append(False)
        self.initialized = True

    def __iter__(self):
        for i,batch in enumerate(self.source_datapipe):
            # print('Got batch: ',batch)
            if not self.initialized: self._reset()
            #  (this batch we should publish) and (there are protocols) and (there are some that are ready)
            if type(batch)==GetInputItemResponse and batch.value.startswith('model_state_dict_publish_'): 
                client_num = int(batch.value.replace('model_state_dict_publish_',''))

                if self._expect_response[client_num]:
                    self._expect_response[client_num] = False

                continue
            if i%self.publish_freq==0 and self.protocol_clients and not all(self._expect_response):
                with torch.no_grad():
                    # We need to deepcopy the model itself since `cpu` is an inplace op.
                    # We cant keep the model in cuda because mp.Manager passes around the 
                    # tensors too much and causes errors ref: https://github.com/pytorch/pytorch/issues/30401
                    # This is alos why we cant just call state_dict directly. It returns references
                    # to cuda tensors.
                    if self.do_deepcopy:
                        state = deepcopy(self.model).to(device=self.device).state_dict()
                    else:
                        state = self.model.state_dict()

                        
                    for client_id,client in enumerate(self.protocol_clients):
                        if not self._expect_response[client_id]: 
                            # print('PUBLISHING!!!!')
                            client.request_input_item(
                                key=f'model_state_dict_publish_{client_id}',value=state
                            )
                            self._expect_response[client_id] = True
            yield batch

# %% ../nbs/12l_agents.dqn.asynchronous.ipynb 10
def DQNLearner(
    model,
    dls,
    agent,
    logger_bases=None,
    loss_func=MSELoss(),
    opt=AdamW,
    lr=0.005,
    bs=128,
    max_sz=10000,
    nsteps=1,
    device=None,
    batches=1000,
    publish_freq=1
) -> LearnerHead:
    learner = LearnerBase(model,dls,batches=batches,loss_func=MSELoss(),opt=opt(model.parameters(),lr=lr))
    learner = ModelPublisher(learner,publish_freq=publish_freq)
    learner = BatchCollector(learner,logger_bases=logger_bases,batch_on_pipe=LearnerBase)
    learner = EpocherCollector(learner,logger_bases=logger_bases)
    for logger_base in L(logger_bases): learner = logger_base.connect_source_datapipe(learner)
    if logger_bases: 
        learner = RollingTerminatedRewardCollector(learner,logger_bases)
        learner = EpisodeCollector(learner,logger_bases)
    learner = ExperienceReplay(learner,bs=bs,max_sz=max_sz)
    learner = StepBatcher(learner,device=device)
    learner = QCalc(learner,nsteps=nsteps)
    learner = ModelLearnCalc(learner)
    if logger_bases: 
        learner = LossCollector(learner,logger_bases)
    learner = LearnerHead(learner)
    return learner

# %% ../nbs/12l_agents.dqn.asynchronous.ipynb 11
class CacheLoggerBase(LoggerBase):
    "Short lived logger base meant to dump logs"
    def reset(self):
        # This logger will be exhausted frequently if used in an agent.
        # We need to get the buffer alive so we dont lose reference
        pass
    
    def __iter__(self):
        print('Iterating through buffer of len: ',len(self.buffer))
        yield from self.buffer
        self.buffer.clear()

# %% ../nbs/12l_agents.dqn.asynchronous.ipynb 12
def DQNAgent(
    model,
    logger_bases=None,
    min_epsilon=0.02,
    max_epsilon=1,
    max_steps=1000,
    device='cpu'
)->AgentHead:
    agent_base = AgentBase(model,logger_bases=ifnone(logger_bases,[CacheLoggerBase()]))
    agent = StepFieldSelector(agent_base,field='state')
    agent = InputInjester(agent)
    agent = ModelSubscriber(agent)
    agent = SimpleModelRunner(agent,device=device)
    agent = ArgMaxer(agent)
    selector = EpsilonSelector(agent,min_epsilon=min_epsilon,max_epsilon=max_epsilon,max_steps=max_steps,device=device)

    agent = EpsilonCollector(selector,agent_base.logger_bases)

    agent = ArgMaxer(agent,only_idx=True)
    agent = NumpyConverter(agent)
    agent = PyPrimativeConverter(agent)
    agent = AgentHead(agent)
    return agent
