#
# Authors: Simon Vandenhende
# Licensed under the CC BY-NC 4.0 license (https://creativecommons.org/licenses/by-nc/4.0/)

import os
import sys


class Logger(object):
    def __init__(self, fpath=None,local_rank=0):
        self.console = sys.stdout
        self.file = None
        self.fpath = fpath
        self.local_rank = local_rank
        if fpath is not None and self.local_rank==0:
            if not os.path.exists(os.path.dirname(fpath)):
                os.makedirs(os.path.dirname(fpath))
                self.file = open(fpath, 'w')
            else:
                self.file = open(fpath, 'a')

    def __del__(self):
        self.close()

    def __enter__(self):
        pass

    def __exit__(self, *args):
        self.close()

    def write(self, msg):
        if self.local_rank==0:
            self.console.write(msg)
            if self.file is not None:
                self.file.write(msg)

    def flush(self):
        if self.local_rank==0:
            self.console.flush()
            if self.file is not None:
                self.file.flush()
                os.fsync(self.file.fileno())

    def close(self):
        self.console.close()
        if self.file is not None:
            self.file.close()
