#!/usr/bin/env python
# -*- coding: utf-8; -*-

# Copyright (c) 2021 Oracle and/or its affiliates.
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/

from __future__ import print_function, absolute_import, division

import json

import numpy as np
import pandas as pd
from sklearn import metrics
from sklearn.preprocessing import LabelEncoder
import scipy.stats as st
from ads.common import utils

__all__ = ['ModelEvaluator']


"""
The ModelEvalutor Class takes in the true values and the predicted values
and returns a pandas df of the every we support

First it determines the problem type (binary, regression, etc), then it
class the appropriate evaluation function: _get_<problem type>_metrics
"""

class ModelEvaluator:
    """
    ModelEvaluator takes in the true adn predicted values and returns a pandas dataframe

    Attributes
    ----------
    y_true : array-like object holding the true values for the model
    y_pred : array-like object holding the predicted values for the model
    model_name (str): the name of the model
    classes (list): list of target classes
    positive_class (str): label for positive outcome from model
    y_score : array-like object holdind the scores for true values for the model
    metrics (dict): dictionary object holding model data

    Methods
    -------
    get_metrics()
        Gets the metrics information in a dataframe based on the number of classes
    safe_metrics_call(scoring_functions, *args)
        Applies sklearn scoring functions to parameters in args        

    """

    def __init__(self, y_true, y_pred, model_name, classes=None, positive_class=None, y_score=None):
        self.y_true = y_true
        self.y_pred = np.squeeze(y_pred)  # This is a bug where y_pred is shape (1,n) from AutoML, rather than (n)
        self.y_score = y_score
        self.model_name = model_name
        self.classes = classes
        self.positive_class = positive_class
        self.metrics = {'model_name': model_name, 'classes': classes, 'positive_class': positive_class,
                        'y_true': y_true, 'y_pred': y_pred, 'y_score': y_score}

    def get_metrics(self):
        """
        Gets the metrics information in a dataframe based on the number of classes

        Parameters
        ----------
            self: (`ModelEvaluator` instance)
                The `ModelEvaluator` instance with the metrics.

        Returns
        -------
        :class:`pandas.DataFrame`
            Pandas dataframe containing the metrics
        """

        if self.classes is None:
            self._get_regression_metrics()
            return pd.DataFrame.from_dict(self.metrics, orient='index', columns=[self.model_name])
        else:
            if len(self.classes) == 2:  # 'binary'
                self.positive_class = self.positive_class if self.positive_class else self.classes[0]
                self._get_binary_metrics()
                return pd.DataFrame.from_dict(self.metrics, orient='index', columns=[self.model_name])
            else:  # 'multiclass'
                self._get_multiclass_metrics()
                return pd.DataFrame.from_dict(self.metrics, orient='index', columns=[self.model_name])

    def _get_general_metrics(self):
        args = [self.y_true, self.y_pred, self.classes]
        scoring_functions = {'classification_report': ('classification_report', len(args)),
                             'kappa_score': ('cohen_kappa_score', len(args)),
                             'raw_confusion_matrix': ('confusion_matrix', len(args)),
                             'hamming_loss': ('hamming_loss', len(args)-1),
                             'hinge_loss': ('hinge_loss', len(args))}

        self.safe_metrics_call(scoring_functions, *args)

        args = [self.y_true, self.y_pred]
        scoring_functions_without_labs = {'accuracy': ('accuracy_score', len(args)),
                                          'zero_one_loss': ('zero_one_loss', len(args))}
        self.safe_metrics_call(scoring_functions_without_labs, *args)

        # could still add: brier_score_loss, log_loss, matthews_corrcoef

        cm = self.metrics['raw_confusion_matrix']
        a = cm.astype('float')
        b = cm.sum(axis=1)[:, np.newaxis]
        normalized_cm = np.divide(
            a, b, out=np.zeros_like(a), where=b != 0).tolist()
        self.metrics['confusion_matrix'] = normalized_cm

        vc = self.y_true.value_counts()
        self.metrics['class_balance'] = min(vc) / max(vc)

    def _get_binary_metrics(self):
        self._get_general_metrics()
        # add a metric to determine imbalance
        self.metrics['balanced accuracy'] = metrics.balanced_accuracy_score(self.y_true, self.y_pred)
        self.metrics['precision'] = metrics.precision_score(self.y_true, self.y_pred,
                                                            pos_label=self.positive_class)
        self.metrics['recall'] = metrics.recall_score(self.y_true, self.y_pred,
                                                      pos_label=self.positive_class)
        self.metrics['f1'] = metrics.f1_score(self.y_true, self.y_pred,
                                              pos_label=self.positive_class, average='binary')

        if self.y_score is not None:
            if not all(0 >= x >= 1 for x in self.y_score):
                self.y_score = np.asarray([0 if x < 0 else 1 if x > 1 else x for x in self.y_score])
            self.y_score = list(self.y_score)
            self.metrics['false_positive_rate'], self.metrics['true_positive_rate'], _ = \
                metrics.roc_curve(self.y_true, self.y_score, pos_label=self.positive_class)
            self.metrics['youden_j'] = self.metrics['true_positive_rate'] - self.metrics['false_positive_rate']
            best_idx = np.argmax(self.metrics['youden_j'])
            self.metrics['roc_best_model_score'] = (self.metrics['false_positive_rate'][best_idx],
                                                    self.metrics['true_positive_rate'][best_idx])
            self.metrics['precision_values'], self.metrics['recall_values'], _ = \
                metrics.precision_recall_curve(self.y_true, self.y_score, pos_label=self.positive_class)
            pr_best_idx = np.argmax(self.metrics['precision_values'] + self.metrics['recall_values'])
            self.metrics['pr_best_idx'] = pr_best_idx
            self.metrics['pr_best_model_score'] = (self.metrics['recall_values'][pr_best_idx],
                                                   self.metrics['precision_values'][pr_best_idx])
            self.metrics['auc'] = metrics.auc(self.metrics['false_positive_rate'], self.metrics['true_positive_rate'])
            self.metrics['average_precision_score'] = metrics.average_precision_score(self.y_true, self.y_score,
                                                                                      pos_label=self.positive_class)

            self.metrics['brier score'] = metrics.brier_score_loss(self.y_true, self.y_score)
            self._get_lift_and_gain()
            # Compute KS Statistic curves
            self._binary_ks_curve()

    def _get_lift_and_gain(self):
        # make y_true a boolean vector
        y_true, y_score = np.asarray(self.y_true), np.asarray(self.y_score)
        y_true = (y_true == self.positive_class)

        sorted_indices = np.argsort(y_score)[::-1]
        y_true = y_true[sorted_indices]
        gains = np.cumsum(y_true)

        percentages = np.arange(start=1, stop=len(y_true) + 1)
        tp = sum(y_true)
        perfect = np.append(np.arange(start=1, stop=tp + 1) / float(tp), np.ones(len(y_true) - tp))

        gains = gains / float(np.sum(y_true))
        percentages = percentages / float(len(y_true))

        self.metrics['cumulative_gain'] = np.insert(gains, 0, [0]) * 100
        self.metrics['percentages'] = np.insert(percentages, 0, [0]) * 100
        self.metrics['perfect_gain'] = np.insert(perfect, 0, [0]) * 100

        percentages = percentages[1:]
        self.metrics['lift'] = gains[1:] / percentages
        self.metrics['perfect_lift'] = perfect[1:] / percentages

    def _binary_ks_curve(self):
        """This function generates the points necessary to calculate the KS
        Statistic curve.
        Args:
            y_true (array-like, shape (n_samples)): True labels of the data.
            y_probas (array-like, shape (n_samples)): Probability predictions of
                the positive class.
        Returns:
            thresholds (numpy.ndarray): An array containing the X-axis values for
                plotting the KS Statistic plot.
            pct1 (numpy.ndarray): An array containing the Y-axis values for one
                curve of the KS Statistic plot.
            pct2 (numpy.ndarray): An array containing the Y-axis values for one
                curve of the KS Statistic plot.
            ks_statistic (float): The KS Statistic, or the maximum vertical
                distance between the two curves.
            max_distance_at (float): The X-axis value at which the maximum vertical
                distance between the two curves is seen.
            classes (np.ndarray, shape (2)): An array containing the labels of the
                two classes making up `y_true`.
        Raises:
            ValueError: If `y_true` is not composed of 2 classes. The KS Statistic
                is only relevant in binary classification.
        """
        y_true, y_probas = np.asarray(self.y_true), np.asarray(self.y_score)
        lb = LabelEncoder()
        encoded_labels = lb.fit_transform(y_true)
        if len(lb.classes_) != 2:
            raise ValueError('Cannot calculate KS statistic for data with '
                             '{} category/ies'.format(len(lb.classes_)))
        idx = encoded_labels == 0
        data1 = np.sort(y_probas[idx])
        data2 = np.sort(y_probas[np.logical_not(idx)])

        ctr1, ctr2 = 0, 0
        thresholds, pct1, pct2 = [], [], []
        while ctr1 < len(data1) or ctr2 < len(data2):

            # Check if data1 has no more elements
            if ctr1 >= len(data1):
                current = data2[ctr2]
                while ctr2 < len(data2) and current == data2[ctr2]:
                    ctr2 += 1

            # Check if data2 has no more elements
            elif ctr2 >= len(data2):
                current = data1[ctr1]
                while ctr1 < len(data1) and current == data1[ctr1]:
                    ctr1 += 1

            else:
                if data1[ctr1] > data2[ctr2]:
                    current = data2[ctr2]
                    while ctr2 < len(data2) and current == data2[ctr2]:
                        ctr2 += 1

                elif data1[ctr1] < data2[ctr2]:
                    current = data1[ctr1]
                    while ctr1 < len(data1) and current == data1[ctr1]:
                        ctr1 += 1

                else:
                    current = data2[ctr2]
                    while ctr2 < len(data2) and current == data2[ctr2]:
                        ctr2 += 1
                    while ctr1 < len(data1) and current == data1[ctr1]:
                        ctr1 += 1

            thresholds.append(current)
            pct1.append(ctr1)
            pct2.append(ctr2)

        thresholds = np.asarray(thresholds)
        pct1 = np.asarray(pct1) / float(len(data1))
        pct2 = np.asarray(pct2) / float(len(data2))

        if thresholds[0] != 0:
            thresholds = np.insert(thresholds, 0, [0.0])
            pct1 = np.insert(pct1, 0, [0.0])
            pct2 = np.insert(pct2, 0, [0.0])
        if thresholds[-1] != 1:
            thresholds = np.append(thresholds, [1.0])
            pct1 = np.append(pct1, [1.0])
            pct2 = np.append(pct2, [1.0])

        differences = pct1 - pct2
        self.metrics['ks_statistic'], self.metrics['max_distance_at'] = (np.max(differences),
                                                                         thresholds[np.argmax(differences)])

        self.metrics['ks_thresholds'], self.metrics['ks_pct1'], self.metrics['ks_pct2'], \
        self.metrics['ks_labels'] = thresholds, pct1, pct2, lb.classes_

    def _get_multiclass_metrics(self):
        self._get_general_metrics()
        y_true, y_pred, y_score = np.asarray(self.y_true), np.asarray(self.y_pred), np.asarray(self.y_score)
        if not all([x in self.classes for x in np.unique(y_true)]):
            utils.print_user_message(f"At least 1 value found in the test data was not specified by the 'classes' "
                                     f"attribute of model: {self.model_name}. This may lead to erroneous "
                                     f"calculations.", display_type='warning')

        self.metrics['precision_weighted'] = metrics.precision_score(y_true, y_pred, self.classes, average='weighted')
        self.metrics['precision_micro'] = metrics.precision_score(y_true, y_pred, self.classes, average='micro')
        self.metrics['precision_by_label'] = metrics.precision_score(y_true, y_pred, self.classes,
                                                                     average=None).tolist()
        self.metrics['recall_weighted'] = metrics.recall_score(y_true, y_pred, self.classes, average='weighted')
        self.metrics['recall_micro'] = metrics.recall_score(y_true, y_pred, self.classes, average='micro')
        self.metrics['recall_by_label'] = metrics.recall_score(y_true, y_pred, self.classes, average=None).tolist()
        self.metrics['f1_weighted'] = metrics.f1_score(y_true, y_pred, self.classes, average='weighted')
        self.metrics['f1_micro'] = metrics.f1_score(y_true, y_pred, self.classes, average='micro')
        self.metrics['f1_by_label'] = metrics.f1_score(y_true, y_pred, self.classes, average=None).tolist()
        self.metrics['jaccard_weighted'] = metrics.jaccard_score(y_true, y_pred, self.classes, average='weighted')
        self.metrics['jaccard_micro'] = metrics.jaccard_score(y_true, y_pred, self.classes, average='micro')
        self.metrics['jaccard_by_label'] = metrics.jaccard_score(y_true, y_pred, self.classes, average=None).tolist()

        if self.y_score is not None:
            # Multiclass ROC
            self.metrics['fpr_by_label'], self.metrics['tpr_by_label'], self.metrics['auc'], \
            self.metrics['roc_best_model_score'] = dict(), dict(), dict(), dict()
            for i, label in enumerate(self.classes):
                self.metrics['fpr_by_label'][i], self.metrics['tpr_by_label'][i], _ = metrics.roc_curve(
                    y_true, y_score[:, i], pos_label=self.classes[i])
                self.metrics['auc'][i] = metrics.auc(self.metrics['fpr_by_label'][i], self.metrics['tpr_by_label'][i])
                youden_j = self.metrics['tpr_by_label'][i] - self.metrics['fpr_by_label'][i]
                best_idx = np.argmax(youden_j)
                self.metrics['roc_best_model_score'][i] = (self.metrics['fpr_by_label'][i][best_idx],
                                                           self.metrics['tpr_by_label'][i][best_idx])
            # Multiclass PR
            self.metrics['recall_values'], self.metrics['precision_values'], \
            self.metrics['pr_best_model_score'] = dict(), dict(), dict()
            for i, label in enumerate(self.classes):
                self.metrics['precision_values'][i], self.metrics['recall_values'][
                    i], _ = metrics.precision_recall_curve(
                    y_true, y_score[:, i], pos_label=self.classes[i])
                pr_best_idx = np.argmax(self.metrics['precision_values'][i] + self.metrics['recall_values'][i])
                self.metrics['pr_best_model_score'][i] = (self.metrics['recall_values'][i][pr_best_idx],
                                                          self.metrics['precision_values'][i][pr_best_idx])

    def _get_regression_metrics(self):
        self.y_true = np.array(self.y_true)
        self.y_pred = np.array(np.squeeze(self.y_pred))
        args = [self.y_true, self.y_pred]
        scoring_functions = {'r2_score': ('r2_score', len(args)),
                             'explained_variance_score': ('explained_variance_score', len(args)),
                             'max_error': ('max_error', len(args)),
                             'mae': ('mean_absolute_error', len(args)),
                             'mse': ('mean_squared_error', len(args)),
                             'median_absolute_error': ('median_absolute_error', len(args))}

        self.safe_metrics_call(scoring_functions, *args)

        self.metrics['root_mean_squared_error'] = np.sqrt(self.metrics['mse'])
        self.metrics['residuals'] = self.y_true - self.y_pred
        self.metrics['mean_residuals'] = np.mean(self.metrics['residuals'])

        # For QQ Plot:
        portions = min(len(self.metrics['residuals']), 100) + 1
        norm_quantiles = (np.arange(portions) / portions)[1:]
        self.metrics['norm_quantiles'] = st.norm.ppf(norm_quantiles)
        resid_quant = [np.quantile(self.metrics['residuals'], p) for p in norm_quantiles]
        self.metrics['residual_quantiles'] = st.zscore(resid_quant)

    def safe_metrics_call(self, scoring_functions, *args):
        """Applies the sklearn function in `scoring_functions` to parameters in `args`.

        Parameters
        ----------
            scoring_functions: (dict)
                Scoring functions dictionary
            args: (keyword arguments)
                Arguments passed to the sklearn function from metrics

        Returns:
            Nothing

        Raises:
            Exception: If an error is enountered applying the sklearn function fn to arguments.
        """

        for name, (fn, n_params) in scoring_functions.items():
            try:
                self.metrics[name] = getattr(metrics, fn)(*(args[:n_params]))
            except Exception as e:
                self.metrics[name] = f"Error unable to compute {fn}, due to: {e}"
