from .callback import Callback
from ..loops import check_path
import sys
import os
from os import path as osp
import platform
from torch.utils.tensorboard import SummaryWriter
from torch.utils.tensorboard.summary import hparams



class LogCallback(Callback):
    def __init__(self, log_graph=True, log_dir='./logs'):
        super().__init__(priority=1)
        self.log_graph = log_graph
        self.graph_saved = False

        self.log_dir=os.path.abspath(log_dir)
        self.global_train_batch_idx = 0
        self.global_train_epoch_idx = 0
        self.global_val_batch_idx = 0
        self.global_val_epoch_idx = 0

    def log(self, metrics, stage, loop_phase, idx):
        for k, v in metrics.items():
            self.logger.add_scalar(f'{k}/{loop_phase}/{stage}', v, idx)

    def log_hparams(self, hyper_params, metrics):
        # self.logger.add_hparams(hyper_params, metrics)
        # 下面的方式可以将超参数写入已有文件之中
        exp, ssi, sei = hparams(hyper_params, metrics)
        self.logger.file_writer.add_summary(exp)
        self.logger.file_writer.add_summary(ssi)
        self.logger.file_writer.add_summary(sei)
        for k, v in metrics.items():
            self.logger.add_scalar(k, v)

    def on_before_fit(self, trainer, epochs):
        log_dir = osp.join(self.log_dir, trainer.running_id)
        check_path(log_dir)
        self.logger = SummaryWriter(log_dir=log_dir)

    def on_before_train_batch(self, trainer, batch_x, batch_y, batch_idx):
        if self.log_graph and not self.graph_saved:
            self.logger.add_graph(trainer.model, batch_x[0])
            self.graph_saved = True

    def on_after_test_epoch(self, trainer, task, metrics):
        if trainer.hyper_params is not None:
            self.log_hparams(trainer.hyper_params, metrics)

    def on_after_train_batch(self, trainer, metrics, batch_idx):
        self.global_train_batch_idx += 1
        self.log(metrics, 'train', 'batch', self.global_train_batch_idx)

    def on_after_train_epoch(self, trainer, task, metrics):
        self.global_train_epoch_idx += 1
        self.log(metrics, 'train', 'epoch', self.global_train_epoch_idx)

    def on_after_val_batch(self, trainer, metrics, batch_idx):
        self.global_val_batch_idx += 1
        self.log(metrics, 'val', 'batch', self.global_val_batch_idx)

    def on_after_val_epoch(self, trainer, task, metrics):
        self.global_val_epoch_idx += 1
        self.log(metrics, 'val', 'epoch', self.global_val_epoch_idx)

    def run_tensorboard(self):
        if 'Windows' in platform.system():
            tensorboard = osp.join(sys.prefix, 'Scripts', 'tensorboard')
        else:
            tensorboard = osp.join(sys.prefix, 'bin', 'tensorboard')
        try:
            os.system(f'{tensorboard} --logdir={self.log_dir}')
        except:  # pylint: disable=W0702
            pass
