# Copyright (c) Facebook, Inc. and its affiliates.
import base64
import logging
import os
import sys

from tensorboardX import SummaryWriter

from pythia.utils.distributed_utils import is_main_process
from pythia.utils.general import (ckpt_name_from_core_args,
                                  foldername_from_config_override)
from pythia.utils.timer import Timer


class Logger:
    def __init__(self, config):
        self.logger = None
        self.summary_writer = None

        if not is_main_process():
            return

        self.timer = Timer()
        self.config = config
        self.save_dir = config.training_parameters.save_dir
        self.log_folder = ckpt_name_from_core_args(config)
        self.log_folder += foldername_from_config_override(config)
        time_format = "%Y-%m-%dT%H:%M:%S"
        self.log_filename = ckpt_name_from_core_args(config) + "_"
        self.log_filename += self.timer.get_time_hhmmss(None, format=time_format)
        self.log_filename += ".log"

        self.log_folder = os.path.join(self.save_dir, self.log_folder, "logs")

        arg_log_dir = self.config.get("log_dir", None)
        if arg_log_dir:
            self.log_folder = arg_log_dir

        if not os.path.exists(self.log_folder):
            os.makedirs(self.log_folder)

        tensorboard_folder = os.path.join(self.log_folder, "tensorboard")
        self.summary_writer = SummaryWriter(tensorboard_folder)

        self.log_filename = os.path.join(self.log_folder, self.log_filename)

        print("Logging to:", self.log_filename)

        logging.captureWarnings(True)

        self.logger = logging.getLogger(__name__)
        self._file_only_logger = logging.getLogger(__name__)
        warnings_logger = logging.getLogger("py.warnings")

        # Set level
        level = config["training_parameters"].get("logger_level", "info")
        self.logger.setLevel(getattr(logging, level.upper()))
        self._file_only_logger.setLevel(getattr(logging, level.upper()))

        formatter = logging.Formatter(
            "%(asctime)s %(levelname)s: %(message)s", datefmt="%Y-%m-%dT%H:%M:%S"
        )

        # Add handler to file
        channel = logging.FileHandler(filename=self.log_filename, mode="a")
        channel.setFormatter(formatter)

        self.logger.addHandler(channel)
        self._file_only_logger.addHandler(channel)
        warnings_logger.addHandler(channel)

        # Add handler to stdout
        channel = logging.StreamHandler(sys.stdout)
        channel.setFormatter(formatter)

        self.logger.addHandler(channel)
        warnings_logger.addHandler(channel)

        should_not_log = self.config["training_parameters"]["should_not_log"]
        self.should_log = not should_not_log

        # Single log wrapper map
        self._single_log_map = set()

    def __del__(self):
        if getattr(self, "summary_writer", None) is not None:
            self.summary_writer.close()

    def write(self, x, level="info", donot_print=False):
        if self.logger is None:
            return
        # if it should not log then just print it
        if self.should_log:
            if hasattr(self.logger, level):
                if donot_print:
                    getattr(self._file_only_logger, level)(str(x))
                else:
                    getattr(self.logger, level)(str(x))
            else:
                self.logger.error("Unknown log level type: %s" % level)
        else:
            print(str(x) + "\n")

    def single_write(self, x, level="info"):
        if x + "_" + level in self._single_log_map:
            return
        else:
            self.write(x, level)

    def add_scalar(self, key, value, iteration):
        if self.summary_writer is None:
            return
        self.summary_writer.add_scalar(key, value, iteration)

    def add_scalars(self, scalar_dict, iteration):
        if self.summary_writer is None:
            return
        for key, val in scalar_dict.items():
            self.summary_writer.add_scalar(key, val, iteration)

    def add_histogram_for_model(self, model, iteration):
        if self.summary_writer is None:
            return
        for name, param in model.named_parameters():
            np_param = param.clone().cpu().data.numpy()
            self.summary_writer.add_histogram(name, np_param, iteration)

[((23, 21, 23, 28), 'pythia.utils.timer.Timer', 'Timer', ({}, {}), '()', False, 'from pythia.utils.timer import Timer\n'), ((26, 26, 26, 58), 'pythia.utils.general.ckpt_name_from_core_args', 'ckpt_name_from_core_args', ({(26, 51, 26, 57): 'config'}, {}), '(config)', False, 'from pythia.utils.general import ckpt_name_from_core_args, foldername_from_config_override\n'), ((27, 27, 27, 66), 'pythia.utils.general.foldername_from_config_override', 'foldername_from_config_override', ({(27, 59, 27, 65): 'config'}, {}), '(config)', False, 'from pythia.utils.general import ckpt_name_from_core_args, foldername_from_config_override\n'), ((33, 26, 33, 78), 'os.path.join', 'os.path.join', ({(33, 39, 33, 52): 'self.save_dir', (33, 54, 33, 69): 'self.log_folder', (33, 71, 33, 77): '"""logs"""'}, {}), "(self.save_dir, self.log_folder, 'logs')", False, 'import os\n'), ((42, 29, 42, 73), 'os.path.join', 'os.path.join', ({(42, 42, 42, 57): 'self.log_folder', (42, 59, 42, 72): '"""tensorboard"""'}, {}), "(self.log_folder, 'tensorboard')", False, 'import os\n'), ((43, 30, 43, 63), 'tensorboardX.SummaryWriter', 'SummaryWriter', ({(43, 44, 43, 62): 'tensorboard_folder'}, {}), '(tensorboard_folder)', False, 'from tensorboardX import SummaryWriter\n'), ((45, 28, 45, 76), 'os.path.join', 'os.path.join', ({(45, 41, 45, 56): 'self.log_folder', (45, 58, 45, 75): 'self.log_filename'}, {}), '(self.log_folder, self.log_filename)', False, 'import os\n'), ((49, 8, 49, 37), 'logging.captureWarnings', 'logging.captureWarnings', ({(49, 32, 49, 36): '(True)'}, {}), '(True)', False, 'import logging\n'), ((51, 22, 51, 49), 'logging.getLogger', 'logging.getLogger', ({(51, 40, 51, 48): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((52, 33, 52, 60), 'logging.getLogger', 'logging.getLogger', ({(52, 51, 52, 59): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((53, 26, 53, 58), 'logging.getLogger', 'logging.getLogger', ({(53, 44, 53, 57): '"""py.warnings"""'}, {}), "('py.warnings')", False, 'import logging\n'), ((60, 20, 62, 9), 'logging.Formatter', 'logging.Formatter', (), '', False, 'import logging\n'), ((65, 18, 65, 75), 'logging.FileHandler', 'logging.FileHandler', (), '', False, 'import logging\n'), ((73, 18, 73, 51), 'logging.StreamHandler', 'logging.StreamHandler', ({(73, 40, 73, 50): 'sys.stdout'}, {}), '(sys.stdout)', False, 'import logging\n'), ((20, 15, 20, 32), 'pythia.utils.distributed_utils.is_main_process', 'is_main_process', ({}, {}), '()', False, 'from pythia.utils.distributed_utils import is_main_process\n'), ((29, 28, 29, 60), 'pythia.utils.general.ckpt_name_from_core_args', 'ckpt_name_from_core_args', ({(29, 53, 29, 59): 'config'}, {}), '(config)', False, 'from pythia.utils.general import ckpt_name_from_core_args, foldername_from_config_override\n'), ((39, 15, 39, 46), 'os.path.exists', 'os.path.exists', ({(39, 30, 39, 45): 'self.log_folder'}, {}), '(self.log_folder)', False, 'import os\n'), ((40, 12, 40, 40), 'os.makedirs', 'os.makedirs', ({(40, 24, 40, 39): 'self.log_folder'}, {}), '(self.log_folder)', False, 'import os\n')]