# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_agent.ipynb (unless otherwise specified).

__all__ = ['agent_event', 'Agent', 'AgentCallback']

# Cell
# Python native modules
import os
# Third party libs
from fastcore.all import *
# Local modules
from .callback.core import *

# Cell
_events = L.split('create preprocess action noise')
_events = _events.map(lambda s:'after_cancel_'+s)+ _events.map(lambda s:'after_cancel_'+s) + _events.map(lambda s:'before_'+s) + _events.map(lambda s:'after_'+s)

mk_class('agent_event', **_events.map_dict(),
         doc="All possible events as attributes to get tab-completion and typo-proofing")

#nbdev_comment _all_ = ['agent_event']

_inner_loop = "".split()

# Cell
class Agent(Loop):
    _loop=L(['Start Action','before_preprocess','after_preprocess',
                            'before_action','after_action',
                            'before_noise','after_noise','End Action'])
    _default='agent'
    end_event='after_noise'
    @delegates(Loop)
    def __init__(self,model=None,**kwargs):
        super().__init__(**kwargs)
        self.experience=None
        self.action=None
        self.model=model

    def state_dict(self): return self.model.state_dict()
    def parameters(self): return self.model.parameters()
    def children(self): return self.model.children()
    def train(self,*args,**kwargs): return self.model.train(*args,**kwargs)
    def to(self,*args,**kwargs):
        self.model.to(*args,**kwargs)
        return self

    def _call_one(self, event_name):
        if not hasattr(agent_event, event_name): raise Exception(f'missing {event_name}')
        for cb in self.cbs.sorted('order'): cb(event_name)

    def __call__(self, event_name=None,**kwargs):
        if isinstance(event_name,str): L(event_name).map(self._call_one)
        elif isinstance(event_name,dict):return self.do_action(**event_name)

    def do_action(self,**kwargs):
        self.experience=kwargs
        for event in L.split('preprocess action noise'):
            self._with_events(noop,event,Exception)
        return self.action,self.experience

# Cell
class AgentCallback(LoopCallback):
    _default,agent='agent',None
    _methods=_events

    def after_cancel_preprocess(self):raise
    def after_cancel_action(self):    raise
    def after_cancel_noise(self):     raise