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

# %% auto 0
__all__ = ['DataPipeAugmentationFn', 'DQN', 'DQNAgent', 'QCalc', 'TargetCalc', 'LossCalc', 'ModelLearnCalc', 'StepBatcher',
           'EpisodeCollector', 'LossCollector', 'RollingTerminatedRewardCollector', 'DQNLearner']

# %% ../nbs/12g_agents.dqn.basic.ipynb 3
# Python native modules
import os
from collections import deque
# Third party libs
from fastcore.all import *
import torchdata.datapipes as dp
from torchdata.dataloader2 import DataLoader2
from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta
from torchdata.dataloader2.graph import find_dps,traverse,DataPipe
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 *
# 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 *

# %% ../nbs/12g_agents.dqn.basic.ipynb 6
class DQN(Module):
    def __init__(self,
                 state_sz:int,  # The input dim of the state
                 action_sz:int, # The output dim of the actions
                 hidden=512,    # Number of neurons connected between the 2 input/output layers
                 head_layer:Module=Linear, # DQN extensions such as Dueling DQNs have custom heads
                 activition_fn:Module=ReLU # The activiation fn used by `DQN`
                ):
        self.layers=Sequential(
            Linear(state_sz,hidden),
            activition_fn(),
            head_layer(hidden,action_sz),
        )
    def forward(self,x): return self.layers(x)


# %% ../nbs/12g_agents.dqn.basic.ipynb 8
DataPipeAugmentationFn = Callable[[DataPipe],Optional[DataPipe]]

def DQNAgent(
    model,
    logger_bases=None,
    min_epsilon=0.02,
    max_epsilon=1,
    max_steps=1000,
    device='cpu',
    dp_augmentation_fns:Optional[List[DataPipeAugmentationFn]]=None
)->AgentHead:
    agent_base = AgentBase(model,logger_bases=logger_bases)
    agent = StepFieldSelector(agent_base,field='state')
    agent = SimpleModelRunner(agent,device=device)
    agent = ArgMaxer(agent)
    agent = EpsilonSelector(agent,min_epsilon=min_epsilon,max_epsilon=max_epsilon,max_steps=max_steps,device=device)
    if agent_base.logger_bases is not None: 
        agent = EpsilonCollector(agent,agent_base.logger_bases)
    agent = ArgMaxer(agent,only_idx=True)
    agent = NumpyConverter(agent)
    agent = PyPrimativeConverter(agent)
    agent = AgentHead(agent)
    
    for fn in ifnone(dp_augmentation_fns,[]):
        result = fn(agent)
        if result is not None: agent = result
    return agent

# %% ../nbs/12g_agents.dqn.basic.ipynb 15
class QCalc(dp.iter.IterDataPipe):
    def __init__(self,source_datapipe):
        self.source_datapipe = source_datapipe
        self.learner = find_dp(traverse(self),LearnerBase)
        
    def __iter__(self):
        for batch in self.source_datapipe:
            self.learner.done_mask = batch.terminated.reshape(-1,)
            self.learner.next_q = self.learner.model(batch.next_state)
            self.learner.next_q = self.learner.next_q.max(dim=1).values.reshape(-1,1)
            self.learner.next_q[self.learner.done_mask] = 0 
            yield batch

# %% ../nbs/12g_agents.dqn.basic.ipynb 16
class TargetCalc(dp.iter.IterDataPipe):
    def __init__(self,source_datapipe,discount=0.99,nsteps=1):
        self.source_datapipe = source_datapipe
        self.discount = discount
        self.nsteps = nsteps
        self.learner = None
        
    def __iter__(self):
        self.learner = find_dp(traverse(self),LearnerBase)
        for batch in self.source_datapipe:
            self.learner.targets = batch.reward+self.learner.next_q*(self.discount**self.nsteps)
            self.learner.pred = self.learner.model(batch.state)
            self.learner.target_qs = self.learner.pred.clone()
            self.learner.target_qs.scatter_(1,batch.action.long(),self.learner.targets)
            yield batch

# %% ../nbs/12g_agents.dqn.basic.ipynb 17
class LossCalc(dp.iter.IterDataPipe):
    def __init__(self,source_datapipe,discount=0.99,nsteps=1):
        self.source_datapipe = source_datapipe
        self.discount = discount
        self.nsteps = nsteps
        self.learner = find_dp(traverse(self),LearnerBase)
        
    def __iter__(self):
        for batch in self.source_datapipe:
            self.learner.loss_grad = self.learner.loss_func(self.learner.pred, self.learner.target_qs)
            yield batch

# %% ../nbs/12g_agents.dqn.basic.ipynb 18
class ModelLearnCalc(dp.iter.IterDataPipe):
    def __init__(self,source_datapipe):
        self.source_datapipe = source_datapipe
        self.learner = find_dp(traverse(self),LearnerBase)
        
    def __iter__(self):
        for batch in self.source_datapipe:
            self.learner.loss_grad.backward()
            self.learner.opt.step()
            self.learner.opt.zero_grad()
            self.learner.loss = self.learner.loss_grad.clone()
            yield self.learner.loss

# %% ../nbs/12g_agents.dqn.basic.ipynb 19
class StepBatcher(dp.iter.IterDataPipe):
    def __init__(self,
            source_datapipe,
            device=None
        ):
        "Converts multiple `StepType` into a single `StepType` with the fields concated."
        self.source_datapipe = source_datapipe
        self.device = device
        
    def vstack_by_fld(self,batch,fld):
        try:
            if self.device is None: return torch.vstack(tuple(getattr(step,fld) for step in batch))
            return torch.vstack(tuple(getattr(step,fld) for step in batch)).to(torch.device(self.device))
        except RuntimeError as e:
            print(f'Failed to stack {fld} given batch: {batch}')
            raise
        
        
    def __iter__(self):
        for batch in self.source_datapipe:
            cls = batch[0].__class__
            yield cls(**{fld:self.vstack_by_fld(batch,fld) for fld in cls._fields})

# %% ../nbs/12g_agents.dqn.basic.ipynb 20
class EpisodeCollector(LogCollector):
    
    def episode_detach(self,step): 
        try:
            v = step.episode_n.cpu().detach().numpy()
            if len(v.shape)==0: return int(v)
            return v[0]
        except IndexError:
            print(f'Got IndexError getting episode_n which is unexpected: \n{step}')
            raise
    
    def __iter__(self):
        for q in self.main_buffers: q.append(Record('episode',None))
        for steps in self.source_datapipe:
            if isinstance(steps,dp.DataChunk):
                for step in steps:
                    for q in self.main_buffers: q.append(Record('episode',self.episode_detach(step)))
            else:
                for q in self.main_buffers: q.append(Record('episode',self.episode_detach(steps)))
            yield steps

# %% ../nbs/12g_agents.dqn.basic.ipynb 21
class LossCollector(LogCollector):
    def __init__(self,
         source_datapipe, # The parent datapipe, likely the one to collect metrics from
         logger_bases:List[LoggerBase] # `LoggerBase`s that we want to send metrics to
        ):
        self.source_datapipe = source_datapipe
        self.main_buffers = [o.buffer for o in logger_bases]
        self.learner = find_dp(traverse(self),LearnerBase)
        
    def __iter__(self):
        for q in self.main_buffers: q.append(Record('loss',None))
        for steps in self.source_datapipe:
            for q in self.main_buffers: q.append(Record('loss',self.learner.loss.cpu().detach().numpy()))
            yield steps

# %% ../nbs/12g_agents.dqn.basic.ipynb 22
class RollingTerminatedRewardCollector(LogCollector):
    debug=False
    def __init__(self,
         source_datapipe, # The parent datapipe, likely the one to collect metrics from
         logger_bases:List[LoggerBase], # `LoggerBase`s that we want to send metrics to
         rolling_length:int=100
        ):
        self.source_datapipe = source_datapipe
        self.main_buffers = [o.buffer for o in logger_bases]
        self.rolling_rewards = deque([],maxlen=rolling_length)
        
    def step2terminated(self,step): return bool(step.terminated)

    def reward_detach(self,step): 
        try:
            v = step.total_reward.cpu().detach().numpy()
            if len(v.shape)==0: return float(v)
            return v[0]
        except IndexError:
            print(f'Got IndexError getting reward which is unexpected: \n{step}')
            raise

    def __iter__(self):
        for q in self.main_buffers: q.append(Record('rolling_reward',None))
        for steps in self.source_datapipe:
            if self.debug: print(f'RollingTerminatedRewardCollector: ',steps)
            if isinstance(steps,dp.DataChunk):
                for step in steps:
                    if self.step2terminated(step):
                        self.rolling_rewards.append(self.reward_detach(step))
                        for q in self.main_buffers: q.append(Record('rolling_reward',np.average(self.rolling_rewards)))
            elif self.step2terminated(steps):
                self.rolling_rewards.append(self.reward_detach(steps))
                for q in self.main_buffers: q.append(Record('rolling_reward',np.average(self.rolling_rewards)))
            yield steps

# %% ../nbs/12g_agents.dqn.basic.ipynb 23
def DQNLearner(
    model,
    dls,
    logger_bases=None,
    loss_func=MSELoss(),
    opt=AdamW,
    lr=0.005,
    bs=128,
    max_sz=10000,
    nsteps=1,
    device=None,
    batches=None,
    dp_augmentation_fns:Optional[List[DataPipeAugmentationFn]]=None
) -> LearnerHead:
    learner = LearnerBase(model,dls,batches=batches,loss_func=loss_func,opt=opt(model.parameters(),lr=lr))
    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)
    learner = TargetCalc(learner,nsteps=nsteps)
    learner = LossCalc(learner)
    learner = ModelLearnCalc(learner)
    if logger_bases: 
        learner = LossCollector(learner,logger_bases)
    learner = LearnerHead(learner)
    
    for fn in ifnone(dp_augmentation_fns,[]):
        result = fn(learner)
        if result is not None: learner = result
    
    return learner
