# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_agents.core.ipynb.

# %% auto 0
__all__ = ['AgentBase', 'is_agent_base', 'find_agent_base', 'AgentHead', 'StepModelFeeder']

# %% ../nbs/07_agents.core.ipynb 3
# Python native modules
import os
# Third party libs
from fastcore.all import *
import torchdata.datapipes as dp
import torch
from fastai.torch_basics import *
from fastai.torch_core import *
# Local modules
from ..core import *
from ..pipes.core import *
from ..pipes.map.demux import *
from ..pipes.map.mux import *

# %% ../nbs/07_agents.core.ipynb 5
class AgentBase(dp.iter.IterDataPipe):
    def __init__(self,
            model:Module, # The base NN that we getting raw action values out of.
            action_iterator:list=None # A reference to an iterator that contains actions to process.
    ):
        self.model = model
        self.iterable = ifnone(action_iterator,[])
        self.agent_base = self
        
    def __iter__(self):
        while self.iterable:
            yield self.iterable.pop(0)
            
add_docs(
    AgentBase,
    """Acts as the footer of the Agent pipeline. 
    Maintains important state such as the `model` being used for get actions from.
    Also optionally allows passing a reference list of `action_iterator` which is a
    persistent list of actions for the entire agent pipeline to process through.
    
    > Important: Must be at the start of the pipeline, and be used with AgentHead at the end.
    
    > Important: `action_iterator` is stored in the `iterable` field. However the recommended
    way of passing actions to the pipeline is to call an `AgentHead` instance.
    """
    
) 

# %% ../nbs/07_agents.core.ipynb 6
def is_agent_base(pipe): return isinstance(pipe,AgentBase)
def find_agent_base(pipe):
    "Basically just find_pipes+is_agent_base with exception handling"
    agent_base = find_pipes(pipe,is_agent_base)
    if not agent_base:
        raise Exception('`AgentBase` must be at the start of the pipeline, but it seems to be missing.')
    return agent_base[0]

class AgentHead(dp.iter.IterDataPipe):
    def __init__(self,source_datapipe):
        self.source_datapipe = source_datapipe
        self.agent_base = find_agent_base(self.source_datapipe)

    def __call__(self,steps:list):
        if issubclass(steps.__class__,StepType):
            raise Exception(f'Expected List[{StepType}] object got {type(steps)}\n{steps}')
        self.agent_base.iterable.extend(steps)
        return self

    def __iter__(self): yield from self.source_datapipe
    
add_docs(
    AgentHead,
    """Acts as the head of the Agent pipeline. 
    Used for conveniently adding actions to the pipeline to process.
    
    > Important: Must be paired with `AgentBase`
    """
)  

# %% ../nbs/07_agents.core.ipynb 15
class StepModelFeeder(dp.iter.IterDataPipe):
    def __init__(self,
                 source_datapipe, # next() must produce a `StepType`,
                 keys:List[str] # A list of field names to grab and push into `self.agent_base.model`
                ): 
        self.source_datapipe = source_datapipe
        self.keys = keys
        self.agent_base = find_agent_base(self.source_datapipe)

    def __iter__(self):
        for o in self.source_datapipe: 
            
            if not issubclass(b.__class__,StepType):
                raise Exception(f'Expected {StepType} object got {type(step)}\n{step}')
            
            tensors = tuple(getattr(o,k) for k in self.keys)
            
            try: yield self.agent_base.model(tensors)
            except Exception:
                print('Failed on ',o)
                raise
        
add_docs(
    StepModelFeeder,
    """Converts `StepTypes` into unified tensors using `keys` and feeds them into `self.agent_base.model`
    """
)  
    
