# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/08_loggers.core.ipynb.

# %% auto 0
__all__ = ['LoggerBase', 'LogCollector', 'is_pipe_instance', 'find_pipe_instance', 'Record', 'ProgressBarLogger',
           'RewardCollector']

# %% ../nbs/08_loggers.core.ipynb 3
# Python native modules
import os,typing
# Third party libs
from fastcore.all import *
from torch.multiprocessing import Queue
import torchdata.datapipes as dp
from fastprogress.fastprogress import *
# Local modules
from ..pipes.core import *

# %% ../nbs/08_loggers.core.ipynb 5
class LoggerBase(dp.iter.IterDataPipe):
    def __init__(self,source_datapipe=None):
        self.source_datapipe = source_datapipe
        self.main_queue = Queue()
        
    def connect_source_datapipe(self,pipe):
        self.source_datapipe = pipe
        return self
        

# %% ../nbs/08_loggers.core.ipynb 7
class LogCollector(dp.iter.IterDataPipe):
    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_queues = [o.main_queue for o in logger_bases]
        
    def __iter__(self): raise NotImplementedError

# %% ../nbs/08_loggers.core.ipynb 9
def is_pipe_instance(pipe,cls): return isinstance(pipe,cls) 
def find_pipe_instance(main_pipe,pipe_cls):
    return find_pipes(main_pipe,partial(is_pipe_instance,cls=pipe_cls))[0]

# %% ../nbs/08_loggers.core.ipynb 10
class Record(typing.NamedTuple):
    name:str
    value:typing.Any

# %% ../nbs/08_loggers.core.ipynb 11
class ProgressBarLogger(LoggerBase):
    def __init__(self,
                 # This does not need to be immediately set since we need the `LogCollectors` to 
                 # first be able to reference its queues.
                 source_datapipe=None, 
                 # For this, and many LoggerBase objects, they likely will show their metrics
                 # at the end of each epoch.
                 epochs=None,
                 # For this, and many LoggerBase objects, they likely will show progress of an epoch
                 # at the end of each batch.
                 batches=None,
                 # For automatic pipe attaching, we can designate which pipe this should be
                 # referneced for information on which epoch we are on
                 epoch_on_pipe:dp.iter.IterDataPipe=None,
                 # For automatic pipe attaching, we can designate which pipe this should be
                 # referneced for information on which batch we are on
                 batch_on_pipe:dp.iter.IterDataPipe=None
                ):
        self.source_datapipe = source_datapipe
        self.main_queue = Queue()
        self.epochs = epochs
        self.batches = batches
        self.epoch_on_pipe = epoch_on_pipe
        self.batch_on_pipe = batch_on_pipe
    
    def dequeue(self): 
        while not self.main_queue.empty(): yield self.main_queue.get()
        
    def __iter__(self):
        epochs = find_pipe_instance(self,self.epoch_on_pipe).epochs if self.epochs is None else self.epochs
        batches = find_pipe_instance(self,self.batch_on_pipe).batches if self.batches is None else self.batches
        epoch_pipe = find_pipe_instance(self,self.epoch_on_pipe)
        mbar = master_bar(list(range(epochs))) 
        pbar = progress_bar(list(range(batches)),parent=mbar,leave=False)
        attached_collectors = {o.name:o.value for o in self.dequeue()}
        mbar.write(attached_collectors, table=True)

        source_datapipe_iter = iter(self.source_datapipe)
        
        for epoch in mbar:
            for batch in pbar:
                
                batch_step = next(source_datapipe_iter)
                
                
                yield batch_step
                for o in self.dequeue(): attached_collectors[o.name] = o.value
                
                if epoch_pipe.epoch!=epoch:
                    print('breaking',epoch_pipe.epoch,epoch)
                    break
                
            mbar.write([f'{l:.6f}' if isinstance(l, float) else str(l)
                        for l in attached_collectors.values()], table=True)

# %% ../nbs/08_loggers.core.ipynb 12
class RewardCollector(LogCollector):
    def __iter__(self):
        for q in self.main_queues: q.put(Record('reward',None))
        for steps in self.source_datapipe:
            if isinstance(steps,dp.DataChunk):
                for step in steps:
                    for q in self.main_queues: q.put(Record('reward',step.reward.detach().numpy()))
            else:
                for q in self.main_queues: q.put(Record('reward',steps.reward.detach().numpy()))
            yield steps
