##################################################################################
#  SENSIML CONFIDENTIAL                                                          #
#                                                                                #
#  Copyright (c) 2017  SensiML Corporation.                                      #
#                                                                                #
#  The source code contained or  described  herein and all documents related     #
#  to the  source  code ("Material")  are  owned by SensiML Corporation or its   #
#  suppliers or licensors. Title to the Material remains with SensiML Corpora-   #
#  tion  or  its  suppliers  and  licensors. The Material may contain trade      #
#  secrets and proprietary and confidential information of SensiML Corporation   #
#  and its suppliers and licensors, and is protected by worldwide copyright      #
#  and trade secret laws and treaty provisions. No part of the Material may      #
#  be used,  copied,  reproduced,  modified,  published,  uploaded,  posted,     #
#  transmitted, distributed,  or disclosed in any way without SensiML's prior    #
#  express written permission.                                                   #
#                                                                                #
#  No license under any patent, copyright,trade secret or other intellectual     #
#  property  right  is  granted  to  or  conferred upon you by disclosure or     #
#  delivery of the Materials, either expressly, by implication,  inducement,     #
#  estoppel or otherwise.Any license under such intellectual property rights     #
#  must be express and approved by SensiML in writing.                           #
#                                                                                #
#  Unless otherwise agreed by SensiML in writing, you may not remove or alter    #
#  this notice or any other notice embedded in Materials by SensiML or SensiML's #
#  suppliers or licensors in any way.                                            #
#                                                                                #
##################################################################################


import pandas as pd
from random import randint
import numpy as np
from numpy import NaN
import operator

UNK = "UNK"
UNC = "UNC"
STATIC_COLUMNS = [UNK, UNC]
DF_COLUMNS = ["ClassIn", "UNK", "UNC", "Support", "Sens(%)"]
INSERTION_INDEX = 1
AVG = "average"
ACC = "accuracy"
SE = "sensitivity"
PP = "positive_predictive_rate"
SPEC = "specificity"
PREC = "precision"
F1 = "f1_score"
HAS_KEYS = {"f1_score": True, "accuracy": False, "sensitivity": True, "specificity": False,
            "positive_predictive_rate": True, "precision": True}
EXPECTED_KEYS = ["ConfusionMatrix", "f1_score", "accuracy", "sensitivity",
                 "positive_predictive_rate", "precision", "specificity"]
NEEDED_PROPS = ["confusion_matrix", "y_true", "y_pred", "f1_score", "accuracy",
                "sensitivity", "positive_predictive_rate", "precision", "specificity"]


class ConfusionMatrixException(Exception):
    """Base exception for ConfusionMatrix and ConfusionMatrixList objects"""

    def __init__(self, e):
        self.e = e

    def __str__(self):
        return repr(self.e)


class ConfusionMatrix(object):
    """This object is a representation of a confusion matrix that contains properties for statistics that can be
    generated from a confusion matrix as well as a DataFrame representation for easy viewing"""

    def __init__(self, metrics_dict):
        """Initializes ConfusionMatrix object

            Args:
                metrics_dict (dict): dictionary of metrics values calculated by the server code when a test set
                is recognized against a model generating a confusion matrix

            Raise:
                ConfusionMatrixException: If arg is not correct type, or does not contain all needed keys returned by
                the server"""
        if not isinstance(metrics_dict, dict):
            raise ConfusionMatrixException(
                ("Expected a dictionary, instead got %s" % str(type(metrics_dict))))
        keys_missing = []
        for key in EXPECTED_KEYS:
            if key not in metrics_dict.keys():
                keys_missing.append(key)
        if len(keys_missing) >= 1:
            raise ConfusionMatrixException(
                ("Not all keys present in metrics dict, missing keys %s" % str(keys_missing)))
        # Init properties from metrics dict
        self._cm = metrics_dict["ConfusionMatrix"]
        self._accuracy = metrics_dict["accuracy"]
        self._sensitivity = metrics_dict["sensitivity"]
        self._positive_predictive_rate = metrics_dict["positive_predictive_rate"]
        self._precision = metrics_dict["precision"]
        self._specificity = metrics_dict["specificity"]
        self._f1_score = metrics_dict["f1_score"]
        self._y_true = metrics_dict["y_true"]
        self._y_pred = metrics_dict["y_pred"]
        # build DataFrame and determine se and pp by class rather than avg returned by server
        self._class_titles = None
        self._cm_data_frame = self._build_data_frame()

    def _order_columns(self, confusion_matrix):
        """Orders the columns such that the classes come first, then UNK and UNC."""
        columns = confusion_matrix.index.tolist() + ['UNK', 'UNC']
        return confusion_matrix[columns]

    def __str__(self):
        """Override string function to return the formatted string representation of the confusion matrix"""
        formatted_output = 'CONFUSION MATRIX:\n'
        cell_width = 10
        decimal_points = 1

        cell_width_str = str(cell_width)
        decimal_points_str = str(decimal_points)
        cm = self._order_columns(pd.DataFrame.from_dict(self._cm).transpose())
        row_format = ("{:>" + cell_width_str + "}") * \
            (len(cm.columns) + 3) + "\n"
        columns = cm.columns.tolist() + ['Support', 'Sens(%)']
        formatted_output += row_format.format('', *columns)
        decimal_cell_format = "{0:." + decimal_points_str + "f}"
        for i, r in cm.iterrows():
            row = [decimal_cell_format.format(a) for a in r] + [decimal_cell_format.format(r.sum() - r['UNC']),
                                                                decimal_cell_format.format(
                                                                    100 * cm.ix[i, i] / float(r.sum() - r['UNC']))]
            formatted_output += row_format.format(i, *row)
        formatted_output += '\n'
        sum_row = cm.sum().tolist() + [sum(cm.sum().tolist()[:-1]), '']
        formatted_output += row_format.format('Total', *sum_row) + '\n'

        pct_row_format = "{:>" + cell_width_str + "}" + ("{:>" + cell_width_str + "." + decimal_points_str + "f}") * (
            len(cm)) + \
            (
            "{:>" + cell_width_str + "}") * 3 + "{:>" + cell_width_str + "." + decimal_points_str + "f}" + "\n"
        pp_row = 100 * np.array(cm).diagonal().astype(float) / \
            cm.sum()[:-2].astype(float)
        pp_row = pp_row.tolist()

        pp_row += ['', '', 'Acc(%)',
                   100. * sum(np.array(cm).diagonal()) / (sum(cm.sum()) - sum(cm['UNC']))]
        formatted_output += pct_row_format.format('PosPred(%)', *pp_row)

        return formatted_output

    def __html__(self):
        """Override string function to return the formatted string representation of the confusion matrix"""
        style = """<style>
                table {
                font-family: arial, sans-serif;
                border-collapse: collapse;
                width: 90%;
                }

                th {
                border: 1px solid #040911;
                text-align: center;
                padding: 8px;
                background-color: #dddddd;
                }

                td {
                border: 1px solid #040911;
                text-align: center;
                padding: 8px;
                background-color: #f9f9f9;
                }

                </style>
                """

        formatted_output = '<table style="width:90%">'
        cell_width = 10
        decimal_points = 1

        cell_width_str = str(cell_width)
        decimal_points_str = str(decimal_points)
        cm = self._order_columns(pd.DataFrame.from_dict(self._cm).transpose())
        row_format = ("{:>" + cell_width_str + "}") * \
            (len(cm.columns) + 3) + "\n"
        th = '<th>{}</th>'
        td = '<td>{}</td>'
        tr = '<tr>{}</tr>'
        columns = cm.columns.tolist() + ['Support', 'Sens(%)']
        formatted_output += tr.format(''.join([th.format(x)
                                               for x in ['']+columns]))

        decimal_cell_format = "{0:.2f}"
        int_format_cell = "{0:.0f}"
        index = 0
        for i, r in cm.iterrows():
            row = th.format(columns[index])
            row += ''.join([td.format(int_format_cell.format(x))
                            for x in r])
            row += td.format(int_format_cell.format(r.sum() - r['UNC']))
            row += td.format(decimal_cell_format.format(100 *
                                                        cm.ix[i, i] / float(r.sum() - r['UNC'])))
            index += 1

            formatted_output += tr.format(row)

        sum_row = cm.sum().tolist() + [sum(cm.sum().tolist()[:-1]), '']
        sum_row = [td.format(x) for x in sum_row]
        sum_row = [th.format('Total')] + sum_row

        formatted_output += tr.format(''.join(sum_row))

        pp_row = 100 * np.array(cm).diagonal().astype(float) / \
            cm.sum()[:-2].astype(float)

        pp_row = [td.format(decimal_cell_format.format(x))
                  for x in pp_row.tolist()]

        pp_row += [th.format(x) for x in ['', '', 'Acc(%)']]

        pp_row += [td.format(decimal_cell_format.format(100. * sum(
            np.array(cm).diagonal()) / (sum(cm.sum()) - sum(cm['UNC']))))]

        pp_row = [th.format('PosPred(%)')]+pp_row

        formatted_output += tr.format(''.join(pp_row))

        formatted_output += '</table>'
        return style+formatted_output

    def __repr__(self):
        """Override repr function to return str value of object"""
        return str(self)

    @property
    def sensitivity(self):
        return self._sensitivity

    @property
    def positive_predictive_rate(self):
        return self._positive_predictive_rate

    @property
    def accuracy(self):
        return self._accuracy

    @property
    def specificity(self):
        return self._specificity

    @property
    def precision(self):
        return self._precision

    @property
    def f1_score(self):
        return self._f1_score

    @property
    def y_true(self):
        return self._y_true

    @property
    def y_pred(self):
        return self._y_pred

    @property
    def confusion_matrix(self):
        return self._cm

    @property
    def confusion_matrix_data_frame(self):
        return self._cm_data_frame

    def se_by_class_name(self, class_name):
        """Returns the sensitivity of a given class in the confusion matrix

            Args:
                class_name (str): name of the class to return sensitivity of

            Returns:
                se_by_class: if class name exists in matrix return SE, else NaN"""
        return self._sensitivity[class_name] if class_name in self._class_titles else NaN

    def pp_by_class_name(self, class_name):
        """Returns the positive predictivity of a given class in the confusion matrix

            Args:
                class_name (str): name of the class to return positive predictivity of

            Returns:
                pp_by_class: if class name exists in matrix return PP, else NaN"""
        return self._positive_predictive_rate[class_name] if class_name in self._class_titles else NaN

    def _build_data_frame(self):
        """Builds a DataFrame from the dictionary representation of a confusion matrix returned by the server
        primarily for viewing purposes.

            Returns:
                cm_data_frame (pandas.DataFrame): DataFrame representation of dictionary confusion matrix"""
        self._class_titles = [
            x for x in self._cm.keys() if x not in STATIC_COLUMNS]
        self._class_titles.sort()
        row_index = 0
        df_cols = ["ClassIn"] + self._class_titles + \
            ["UNK", "UNC", "Support", "Sens(%)"]
        ret = pd.DataFrame(columns=df_cols)
        support_total = 0.0
        for item in self._class_titles:
            c_vals = []
            for it in self._class_titles + [UNK]:
                c_vals.append(float(self._cm[item][it]))
            row_total = sum(c_vals)
            support_total += row_total
            c_vals.append(float(self._cm[item][UNC]))
            c_vals.append(row_total)
            if row_total == 0.0:
                c_vals.append(NaN)
            else:
                c_vals.append(float(self._cm[item][item] / row_total) * 100)
            n_row = [item] + c_vals
            ret.loc[row_index] = n_row
            row_index += 1
        n_row = ["UNK"] + ([0] * (len(ret.columns.tolist()) - 2)) + [NaN]
        ret.loc[row_index] = n_row
        row_index += 1
        pp_vals = []
        sum_vals = []
        for item in self._class_titles + [UNK]:
            col_total = 0.0
            if item == UNK:
                for un in self._class_titles:
                    col_total += self._cm[un][UNK]
                sum_vals.append(col_total)
            else:
                for it in self._class_titles:
                    col_total += self._cm[it][item]
                sum_vals.append(col_total)
            if item == UNK:
                pp_vals.append(NaN)
            else:
                if col_total > 0.0:
                    pp_vals.append(
                        float(self._cm[item][item] / col_total) * 100)
                else:
                    pp_vals.append(NaN)
        n_row = ["Total"] + sum_vals + [sum(ret[UNC]), support_total, 0]
        ret.loc[row_index] = n_row
        row_index += 1
        temp_total = 0.0
        row_total = 0.0
        for item in self._class_titles:
            temp_total += self._cm[item][item]
            for it in self._class_titles + [UNK]:
                row_total += self._cm[item][it]
        acc = float(temp_total / row_total) * 100
        n_row = ["PosPred(%)"] + pp_vals + [0] + ["Acc(%)"] + [acc]
        ret.loc[row_index] = n_row
        self._class_titles.sort(reverse=True)
        return ret


class ConfusionMatrixList(ConfusionMatrix):
    """ConfusionMatrix list object represents a group of confusion matrices and is itself a confusion matrix.
    Enables investigation of child matrices to determine best model by the statistics of the confusion matrix"""

    def __init__(self, confusion_matrices):
        """Initializes ConfusionMatrixList object
            Args:
                confusion_matrices (dict): keys should be a way to identify each matrix, an index for example
                values are ConfusionMatrix objects

            Raises:
                ConfusionMatrixException: If arg are incorrect type, needs to be a dict of confusion matrices"""
        if isinstance(confusion_matrices, dict):
            for key, value in confusion_matrices.items():
                if not isinstance(value, ConfusionMatrix):
                    raise ConfusionMatrixException(
                        ("Expected a ConfusionMatrix, instead got %s" % str(type(value))))
        else:
            raise ConfusionMatrixException(
                ("Expected a dictionary, instead got %s" % str(type(confusion_matrices))))
        # compile child matrices, and average their statistics to generate this parent matrix
        self._confusion_matrices = confusion_matrices
        self._class_titles = self._inspect_confusion_matrix()
        compiled_metrics = self._average_child_metrics
        super(ConfusionMatrixList, self).__init__(compiled_metrics)
        self._sorted_matrices = [(k, v)
                                 for k, v in self._confusion_matrices.items()]

    def highest_avg_se_child(self):
        """Get key, and child matrix from dict with the highest average sensitivity

            Returns:
                key: key value provided in initialize dictionary pertaining to the confusion matrix
                confusion_matrix (ConfusionMatrix): The confusion matrix object with the highest average sensitivity"""
        return self._query_children(operator.gt, SE, True)

    def highest_avg_pp_child(self):
        """Get key, and child matrix from dict with the highest average positive predictivity

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the highest average PP"""
        return self._query_children(operator.gt, PP, True)

    def highest_accuracy_child(self):
        """Get key, and child matrix from dict with the highest accuracy

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the highest accuracy"""
        return self._query_children(operator.gt, ACC, False)

    def highest_specificity_child(self):
        """Get key, and child matrix from dict with the highest specificity

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the highest specificity"""
        return self._query_children(operator.gt, SPEC, False)

    def highest_precision_child(self):
        """Get key, and child matrix from dict with the highest precision

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the highest precision"""
        return self._query_children(operator.gt, PREC, True)

    def highest_f1_score_child(self):
        """Get key, and child matrix from dict with the highest f1_score

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the highest f1_score"""
        return self._query_children(operator.gt, F1, True)

    def lowest_avg_se_child(self):
        """Get key, and child matrix from dict with the lowest average sensitivity

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the lowest average sensitivity"""
        return self._query_children(operator.lt, SE, True)

    def lowest_avg_pp_child(self):
        """Get key, and child matrix from dict with the lowest average positive predictivity

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the lowest average positive predictivity"""
        return self._query_children(operator.lt, PP, True)

    def lowest_accuracy_child(self):
        """Get key, and child matrix from dict with the lowest accuracy

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the lowest accuracy"""
        return self._query_children(operator.lt, ACC, False)

    def lowest_specificity_child(self):
        """Get key, and child matrix from dict with the lowest specificity

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the lowest specificity"""
        return self._query_children(operator.lt, SPEC, False)

    def lowest_precision_child(self):
        """Get key, and child matrix from dict with the lowest precision

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the lowest precision"""
        return self._query_children(operator.lt, PREC, True)

    def lowest_f1_score_child(self):
        """Get key, and child matrix from dict with the lowest f1_score

        Returns:
            key: key value provided in initialize dictionary pertaining to the confusion matrix
            confusion_matrix (ConfusionMatrix): The confusion matrix object with the lowest f1_score"""
        return self._query_children(operator.lt, F1, True)

    def head(self, length=5):
        """Returns the head of the list of confusion matrix objects, size specified

        Args:
            length (int): Number of values to return from confusion matrix list

        Returns:
            Head of size specified if all parameters are correct
            None if length is invalid
            Entire list if length is greater than the list size"""
        try:
            length = int(length)
        except ValueError:
            return None
        if length <= len(self._sorted_matrices):
            return self._sorted_matrices[0:length]
        else:
            return self._sorted_matrices

    def sort_by_metric(self, metric):
        """Sort the list of confusion matrices by a given metric
            Args:
                metric (str): The metric name to sort by, must be a property of the confusion matrix

            Returns:
                sorted_matrices (list): a list of index, object tuples after sorting is complete

            Raises:
                ConfusionMatrixException if metric value passed in is not recognized"""
        if metric not in EXPECTED_KEYS[1:len(EXPECTED_KEYS)]:
            raise ConfusionMatrixException(("Metric %s is not a supported metric for sorting, use one of the following %s"
                                            % (str(metric), str(EXPECTED_KEYS[1:len(EXPECTED_KEYS)]))))
        self._lambda_child_sort(metric, HAS_KEYS[metric])
        return self._sorted_matrices

    def _lambda_child_sort(self, metric, has_keys):
        """Uses lambda functions to sort the confusion matrix list by a given metric,
         or property of the child matrices. Sorting is done in place.

            Args:
                metric (str): the name of the metric to sort by
                has_keys (bool): If the metric to sort by is a dictionary value True, else False
                """
        temp = self._sorted_matrices
        # If the metric is a dictionary, sort by the average value
        if has_keys:
            temp = sorted(temp, key=lambda x: getattr(
                x[1], metric)[AVG], reverse=True)
        else:
            temp = sorted(temp, key=lambda x: getattr(
                x[1], metric), reverse=True)
        self._sorted_matrices = temp

    @property
    def _average_child_metrics(self):
        """averages the metrics of all of the children provided in order to build dict needed to init a ConfusionMatrix
        object

            Returns:
                metrics_dict (dict): averaged metrics dictionary from all children"""
        ret = dict.fromkeys(NEEDED_PROPS)
        y_true = []
        y_pred = []
        for key in ret.keys():
            if key == NEEDED_PROPS[1]:
                for k, cm in self._confusion_matrices.items():
                    y_true += getattr(cm, key)
                ret[key] = y_true
            elif key == NEEDED_PROPS[2]:
                for k, cm in self._confusion_matrices.items():
                    y_pred += getattr(cm, key)
                ret[key] = y_pred
            # create a list for each statistical measure, and append each CM's value
            elif key != NEEDED_PROPS[0]:
                for index, cm in self._confusion_matrices.items():
                    if isinstance(getattr(cm, key), dict):
                        if not isinstance(ret[key], dict):
                            ret[key] = dict.fromkeys(getattr(cm, key).keys())
                            for k in ret[key].keys():
                                ret[key][k] = []
                        for sub in ret[key].keys():
                            ret[key][sub].append(getattr(cm, key)[sub])
                    else:
                        if not isinstance(ret[key], list):
                            ret[key] = []
                        ret[key].append(getattr(cm, key))
            else:
                ret["ConfusionMatrix"] = self._sum_confusion_matrices()

        # After all of the statistical measures across all confusion matrices have been filled in, average them
        for to_avg in NEEDED_PROPS[3:]:
            if isinstance(ret[to_avg], dict):
                for k in ret[to_avg].keys():
                    ret[to_avg][k] = np.nanmean(ret[to_avg][k])
            else:
                ret[to_avg] = np.nanmean(ret[to_avg])
        return ret

    def _sum_confusion_matrices(self):
        """Sums all of the dictionary confusion matrices of children to create a top level combined confusion matrix

            Returns:
                confusion_matrix (dict): dictionary representation of a confusion matrix"""
        ret = dict.fromkeys(self._class_titles)
        for key in ret.keys():
            for item in self._class_titles + [UNK, UNC]:
                if not isinstance(ret[key], dict):
                    ret[key] = {}
                ret[key][item] = 0.0
        for index, conf in self._confusion_matrices.items():
            for item in self._class_titles:
                for it in self._class_titles + [UNK, UNC]:
                    ret[item][it] += float(conf.confusion_matrix[item][it])
        return ret

    def _inspect_confusion_matrix(self):
        """Inspects the child matrices given to the init function to determine the classes contained in them

            Returns:
                class_titles (list): list of class titles contained in the confusion matrix"""
        k = self._confusion_matrices.keys()
        random_matrix_index = randint(0, len(self._confusion_matrices) - 1)
        to_inspect = self._confusion_matrices[k[random_matrix_index]]
        class_titles = [
            x for x in to_inspect.confusion_matrix.keys() if x not in STATIC_COLUMNS]
        assert(len(class_titles) > 1)
        class_titles.sort(reverse=True)
        return class_titles

    @staticmethod
    def _get_truth(val_a, opr, val_b):
        """returns the truth of two values and a boolean operator

            Returns:
                True if val_a operator val_b is True, Else False"""
        return opr(val_a, val_b)

    def _query_children(self, opr, prop, has_keys=False):
        """Iterates over children to find the child at the top of the list given the operator and property provided

            Example:
                If operator.gt and property f1_score are passed in, the index and confusion matrix with the highest
                f1_score will be returned

            Args:
                opr (operator): Must be a boolean operator
                prop (str): must be a property of the ConfusionMatrix object

            Returns:
                result (tuple): Tuple of length 2, index 0 == index of confusion matrix, index 1 == confusion matrix"""
        is_first = True
        ret = None
        for i, conf in self._confusion_matrices.items():
            if is_first:
                ret = (i, conf)
                is_first = False
            else:
                if not has_keys:
                    if self._get_truth(getattr(conf, prop), opr, getattr(ret[1], prop)):
                        ret = (i, conf)
                else:
                    if self._get_truth(getattr(conf, prop)[AVG], opr, getattr(ret[1], prop)[AVG]):
                        ret = (i, conf)
        return ret
