from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514
import os
import time
import sys
import torch

USE_TENSORBOARD = True
try:
    import tensorboardX

    print("Using tensorboardX")
except:
    USE_TENSORBOARD = False


class Logger(object):
    def __init__(self, opt):
        """Create a summary writer logging to log_dir."""
        if not os.path.exists(opt.save_dir):
            os.makedirs(opt.save_dir)
        if not os.path.exists(opt.debug_dir):
            os.makedirs(opt.debug_dir)

        time_str = time.strftime("%Y-%m-%d-%H-%M")

        args = dict(
            (name, getattr(opt, name)) for name in dir(opt) if not name.startswith("_")
        )
        file_name = os.path.join(opt.save_dir, "opt.txt")
        with open(file_name, "wt") as opt_file:
            opt_file.write("==> torch version: {}\n".format(torch.__version__))
            opt_file.write(
                "==> cudnn version: {}\n".format(torch.backends.cudnn.version())
            )
            opt_file.write("==> Cmd:\n")
            opt_file.write(str(sys.argv))
            opt_file.write("\n==> Opt:\n")
            for k, v in sorted(args.items()):
                opt_file.write("  %s: %s\n" % (str(k), str(v)))

        log_dir = opt.save_dir + "/logs_{}".format(time_str)
        if USE_TENSORBOARD:
            self.writer = tensorboardX.SummaryWriter(log_dir=log_dir)
        else:
            if not os.path.exists(os.path.dirname(log_dir)):
                os.mkdir(os.path.dirname(log_dir))
            if not os.path.exists(log_dir):
                os.mkdir(log_dir)
        self.log = open(log_dir + "/log.txt", "w")
        try:
            os.system("cp {}/opt.txt {}/".format(opt.save_dir, log_dir))
        except:
            pass
        self.start_line = True

    def write(self, txt):
        if self.start_line:
            time_str = time.strftime("%Y-%m-%d-%H-%M")
            self.log.write("{}: {}".format(time_str, txt))
        else:
            self.log.write(txt)
        self.start_line = False
        if "\n" in txt:
            self.start_line = True
            self.log.flush()

    def close(self):
        self.log.close()

    def scalar_summary(self, tag, value, step):
        """Log a scalar variable."""
        if USE_TENSORBOARD:
            self.writer.add_scalar(tag, value, step)
