# ----------------------------------------------------------------------------
# Copyright 2014 Nervana Systems Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""
Contains generic performance metric base class, metric comparison class, and
other utility functions.
"""

import hashlib
import math
import os
import sys

import neon
from neon.util.param import opt_param
from neon.util.persist import YAMLable, ensure_dirs_exist


class Metric(YAMLable):
    """
    A Metric quantitatively measures some aspect of model performance by
    contrasting the predictions generated by the model with actual expected
    outputs.

    Though metrics may examine device buffers, all computation is carried out
    on host, and the results are returned in host buffers.

    This base class defines which operations each metric must support
    to be used within our framework.
    """
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        opt_param(self, ['epoch_interval'])
        self.clear()

    def __str__(self):
        return self.__class__.__name__

    def add(self, reference, outputs):
        """
        Add the expected reference and predicted outputs passed to the set
        of values used to calculate this metric.

        Arguments:
            reference (neon.backend.Tensor): Ground truth, expected outcomes.
                                             If each outcome is a vector, we
                                             expect it to be a column vector,
                                             with each case in a separate
                                             column.
            outputs (neon.backend.Tensor): Predicted outputs.  Must have the
                                           same dimensions as reference.

        Raises:
            NotImplementedError: Can't be instantiated directly.
        """
        raise NotImplementedError()

    def report(self):
        """
        Report this metric's current calculated value(s).

        Returns:
            float or array-like: computed metric value

        Raises:
            NotImplementedError: Can't be instantiated directly.
        """
        raise NotImplementedError()

    def clear(self):
        """
        Reset this metric's calculated value(s)

        Raises:
            NotImplementedError: Can't be instantiated directly.
        """
        raise NotImplementedError()


def dump_metrics(dump_file, experiment_file, start_time, elapsed_time,
                 backend_name, metrics, field_sep="\t"):
    """
    Write or append collected metric values to the specified flat file.

    Arguments:
        dump_file (str): path to file to write. Will be created if doesn't
                         exist, or appended to (without header if it does)
        experiment_file (str): path to yaml file used to run this experiment
        start_time (str): date and time at which experiment was started.
        elapsed_time (float): time taken to run the experiment.
        metrics (dict): Collection of metric values, as returned from
                        FitPredictErrorExperiment.run() call.
        field_sep (str, optional): string used to separate each field in
                                   dump_file.  Defaults to tab character.
    """
    if dump_file is None or dump_file == '':
        df = sys.stdout()
    elif not os.path.exists(dump_file) or os.path.getsize(dump_file) == 0:
        ensure_dirs_exist(dump_file)
        df = open(dump_file, 'w')
        metric_names = []
        if isinstance(metrics, dict):
            metric_names = ["%s-%s" % (metric.lower(), dset.lower())
                            for metric in sorted(metrics.keys())
                            for dset in sorted(metrics[metric].keys())]
        df.write(field_sep.join(["host", "architecture", "os",
                                 "os_kernel_release", "neon_version",
                                 "backend",
                                 "yaml_name", "yaml_sha1", "start_time",
                                 "elapsed_time"] + metric_names) + "\n")
    else:
        df = open(dump_file, 'a')
    info = os.uname()
    trunc_exp_name = ("..." + os.path.sep +
                      os.path.dirname(experiment_file).split(os.path.sep)[-1] +
                      os.path.sep +
                      os.path.basename(experiment_file))
    # TODO: better handle situation where metrics recorded differ from those
    # already in file
    metric_vals = []
    if isinstance(metrics, dict):
        metric_vals = ["%.5f" % metrics[metric][dset] for metric in
                       sorted(metrics.keys()) for dset in
                       sorted(metrics[metric].keys())]
    df.write(field_sep.join([x.replace("\t", " ") for x in
                             [info[1], info[4], info[0], info[2],
                              neon.__version__, backend_name, trunc_exp_name,
                              hashlib.sha1(open(experiment_file,
                                                'rb').read()).hexdigest(),
                              start_time, "%.3f" % elapsed_time] +
                             metric_vals]) + "\n")
    df.close()


class MetricComparison(object):
    """
    Embodies the results of comparing one most recent run of a particular type
    of experiment against a prior set of runs.

    Arguments:
        dump_file (str): path of file to compare.
        experiment_file (str): path to yaml file used to run this experiment
        max_comps (int, optional): collect and compare statistics against
                                   max_comps most recent prior runs of the
                                   same example.  Defaults to 10.
        match_backend (bool, optional): Only compare metric results of the same
                                        backend.  Defaults to True.
        field_sep (str, optional): Dump file field separator.  Defaults to tab
                                   character.
    """
    def __init__(self, dump_file, experiment_file, max_comps=10,
                 match_backend=True, field_sep='\t'):
        self.dump_file = dump_file
        self.experiment_file = experiment_file
        self.experiment = None
        self.backend = None
        self.results = []
        if not os.path.exists(dump_file):
            raise OSError("file: %s doesn't exist.  Can't run comparisons" %
                          dump_file)

        data = open(dump_file).readlines()

        if len(data) < 1 or not data[0].startswith("host"):
            raise OSError("file: %s seems to have invalid format" % dump_file)

        self.experiment = ("..." + os.path.sep +
                           os.path.dirname(experiment_file)
                           .split(os.path.sep)[-1] + os.path.sep +
                           os.path.basename(experiment_file))
        line_num = len(data) - 1
        header = {x[1]: x[0] for x in
                  enumerate(data[0].rstrip('\r\n').split(field_sep))}
        latest = None
        comps = []
        while line_num > 0 and len(comps) < max_comps:
            if self.experiment in data[line_num]:
                this_line = data[line_num].rstrip('\r\n').split(field_sep)
                if (not match_backend or latest is None or
                        self.backend == this_line[header['backend']]):
                    if latest is None:
                        latest = this_line
                        self.backend = latest[header['backend']]
                    else:
                        comps.append(this_line)
            line_num -= 1
        if latest is None:
            raise ValueError("unable to find any lines containing %s" %
                             self.experiment)
        for name, idx in [(x, header[x]) for x in sorted(header.keys())
                          if x == "elapsed_time" or x.startswith("train-") or
                          x.startswith("test-") or
                          x.startswith("validation-")]:
            val = float(latest[idx])