__all__ = ['ExperimentManager']

import logging
import os
import tempfile
import datetime as dt
from matplotlib import pyplot as plt

import neptune

_log = logging.getLogger(__name__)


class ExperimentManager(object):
    """Wrapper around the neptune-client >= 1.0 API.

    When log=False the object is a no-op stub so that downstream code that
    calls log_* methods still runs without changes.
    """

    def __init__(self, log, project_name, experiment_name, experiment_params, experiment_tags):

        logging.getLogger('neptune').setLevel(logging.ERROR)

        self.log = log
        self.run = None
        self._experiment_name = None

        if self.log:
            self.run = neptune.init_run(
                project=project_name,
                name=experiment_name,
                tags=list(experiment_tags) if experiment_tags else [])

            if experiment_params:
                self.run['parameters'] = experiment_params

    @property
    def experiment_name(self):
        if self._experiment_name is None:
            if self.log and self.run:
                self._experiment_name = self.run['sys/id'].fetch()
            else:
                self._experiment_name = ExperimentManager.get_fake_name()

        return self._experiment_name

    @staticmethod
    def get_fake_name():
        return str(int(dt.datetime.now().timestamp()))

    def get_experiment_id(self):
        if self.log and self.run:
            return self.run['sys/id'].fetch()
        return None

    def add_experiment_tags(self, tag):
        if self.log and self.run:
            self.run['sys/tags'].add(tag)
            _log.info('[Tag] %s', tag)

    def set_experiment_property(self, key, value):
        if self.log and self.run:
            self.run[key] = value
            _log.info('[Property] %s: %s', key, value)

    def log_data_to_neptune(self, data, name, log_index=False, log_to_artifact=True):
        if self.log and self.run:
            if name is None:
                name = ExperimentManager.get_fake_name()

            if log_to_artifact:
                # Use NamedTemporaryFile (mode 0o600) so the file is not
                # world-readable while it exists on shared systems.
                import stat
                tmp = tempfile.NamedTemporaryFile(
                    suffix='.csv', delete=False, prefix='pwml_')
                file_path = tmp.name
                tmp.close()
                os.chmod(file_path, stat.S_IRUSR | stat.S_IWUSR)
                try:
                    data.to_csv(file_path, index=log_index)
                    self.run['data/{0}.csv'.format(name)].upload(file_path)
                finally:
                    if os.path.exists(file_path):
                        os.remove(file_path)

    def log_chart_to_neptune(self, figure, name, log_to_artifact=True):
        if self.log and self.run:
            if name is None:
                name = ExperimentManager.get_fake_name()

            if log_to_artifact:
                import stat
                tmp = tempfile.NamedTemporaryFile(
                    suffix='.jpg', delete=False, prefix='pwml_')
                file_path = tmp.name
                tmp.close()
                os.chmod(file_path, stat.S_IRUSR | stat.S_IWUSR)
                try:
                    figure.savefig(
                        fname=file_path,
                        dpi=96,
                        format='jpg',
                        bbox_inches='tight')
                    self.run['charts/{0}.jpg'.format(name)].upload(file_path)
                finally:
                    if os.path.exists(file_path):
                        os.remove(file_path)

    def stop(self):
        """Stop the Neptune run. Call when done logging."""
        if self.log and self.run:
            self.run.stop()
            self.run = None

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.stop()
