from pathlib import Path

from pytorch_lightning.logging.base import LightningLoggerBase
from pytorch_lightning.logging.neptune import NeptuneLogger
from pytorch_lightning.logging.test_tube import TestTubeLogger

from lib.utils.config import Config


class Logger(LightningLoggerBase):

    @property
    def experiment(self):
        if self.debug:
            return self.testtubelogger.experiment
        else:
            return self.neptunelogger.experiment

    @property
    def name(self):
        return self.config.model.type

    @property
    def project_name(self):
        return f"{self.config.project.owner}/{self.config.project.name}"

    @property
    def version(self):
        return f"version_{self.config.get('main', 'seed')}"

    @property
    def outpath(self):
        # ToDo: Add further path modification such as dataset config etc.
        return Path(self.config.train.outpath)

    def __init__(self, config: Config, debug=False):
        """
        params (dict|None): Optional. Parameters of the experiment. After experiment creation params are read-only.
           Parameters are displayed in the experiment’s Parameters section and each key-value pair can be
           viewed in experiments view as a column.
        properties (dict|None): Optional default is {}. Properties of the experiment.
           They are editable after experiment is created. Properties are displayed in the experiment’s Details and
           each key-value pair can be viewed in experiments view as a column.
        tags (list|None): Optional default []. Must be list of str. Tags of the experiment.
           They are editable after experiment is created (see: append_tag() and remove_tag()).
           Tags are displayed in the experiment’s Details and can be viewed in experiments view as a column.
        """
        super(Logger, self).__init__()

        self.debug = debug
        self.config = config
        self._testtube_kwargs = dict(save_dir=self.outpath, version=self.version, name=self.name)
        self._neptune_kwargs = dict(offline_mode= self.debug,
                                    api_key=self.config.project.neptune_key,
                                    project_name=self.project_name,
                                    name=self.name,
                                    upload_source_files=list())
        self.neptunelogger = NeptuneLogger(**self._neptune_kwargs)
        self.testtubelogger = TestTubeLogger(**self._testtube_kwargs)

    def log_hyperparams(self, params):
        self.neptunelogger.log_hyperparams(params)
        self.testtubelogger.log_hyperparams(params)
        pass

    def log_metrics(self, metrics, step_num):
        self.neptunelogger.log_metrics(metrics, step_num)
        self.testtubelogger.log_metrics(metrics, step_num)
        pass