import os
import shutil
import logging

from . import Callback
from typing import Dict


class LogWriter(Callback):
    def __init__(self, name, verbose=2, value_format=None):
        super(LogWriter, self).__init__()

        self.log = logging.getLogger(name)
        self.verbose = verbose
        self.value_format = value_format

        self.epochs = None
        self._e = None
        self.batch_nums = None
        self._i = None

    def format_value(self, value):
        if self.value_format is not None:
            return f'%{self.value_format}' % value
        return str(value)

    def on_epoch_begin(self, epoch_number: int, logs: Dict):
        self._e = epoch_number
        self.epochs = logs.get('epochs')

    def on_train_epoch_begin(self, logs):
        self.batch_nums = logs.get('batch_nums')

    def on_train_batch_begin(self, batch_number: int, logs: Dict):
        self._i = batch_number

    def on_train_batch_end(self, batch_number: int, logs: Dict):
        if self.verbose > 1:
            self.log.info(f'e:{self._e}/{self.epochs}  b{self._i}/{self.batch_nums}  ' +
                          '  '.join([f'{name}:{self.format_value(value)}' for name, value in logs.items()]))

    def on_eval_end(self, logs: Dict):
        if self.verbose > 0:
            self.log.info(f'e:{self._e}/{self.epochs}  -EVAL- ' +
                          '  '.join([f'{name}:{self.format_value(value)}' for name, value in logs.items()]))

