import pandas as pd
import fastai.vision.all as fai


class DataframeLogger(fai.Callback):
    order = 60
    
    def before_fit(self):
        if hasattr(self, "gather_preds"): return
        self.df = pd.DataFrame(columns=self.recorder.metric_names)
        self.old_logger, self.learn.logger = self.logger, self._record_line
        
    def _record_line(self, log):
        if self._is_training(log):
            self.df.loc[log[0]] = log
        self.old_logger(log)
    
    def _is_training(self, log):
        return self.df.columns.size == len(log)
    
    def read_log(self):
        return self.df
    
    def after_fit(self):
        if hasattr(self, "gather_preds"): return
        self.learn.logger = self.old_logger