"""A primrose model based around a sklearn model

Author(s):
    Carl Anderson (carl.anderson@weightwatchers.com)

"""
from primrose.base.model import AbstractModel
import importlib
import logging
from sklearn.metrics import roc_curve, auc
from sklearn import metrics
import datetime


class SklearnModel(AbstractModel):
    def __init__(self, configuration, instance_name):
        """A Sklearn-based model to train, evaluate and predict on dataframe feature data

        Args:
            configuration (Configuration): user defined configuration object
            instance_name (str): name to identify the correct configuration key

        """
        super(SklearnModel, self).__init__(configuration, instance_name)
        self.model = None
        self.scores = {}

    def load_model(self, data_object):
        """finds an upstream sklearn mode

        Args:
            data_object (DataObject): instance of DataObject

        Returns:
            Sklearn model

        """

        # check for a valid pipeline object inside the data object
        model = data_object.get_filtered_upstream_data(
            self.instance_name, filter_for_key="reader_data"
        )["reader_data"]

        if model is None:
            raise Exception(
                "No valid upstream model found in the data object for model."
            )

        logging.info("Model successfully loaded from upstream data_object.")

        return model

    @staticmethod
    def _instantiate_model(classname, args):
        """instantiates a Sklearn model from class name

        Args:
            classname (str): name of class without initial `sklearn.`, e.g. `cluster.KMeans`

            args (dict): dictionary of arguments/parameters for model instantiation

        Returns:
            instance of a sklearn mode

        """
        sk_module_name, sk_model_name = classname.split(".")
        sk_module = importlib.import_module("sklearn.{}".format(sk_module_name))

        try:
            m = getattr(sk_module, sk_model_name)
        except AttributeError:
            raise Exception(
                "Sklearn model {} not found in {} module".format(
                    sk_model_name, sk_module_name
                )
            )
        logging.info("Instantiating model " + classname)
        if args:
            return m(**args)
        else:
            return m()

    @staticmethod
    def _get_roc_score(actual_values, model_probability):
        """Get (binary) ROC store using actual target variables and predicted probabilities

            Assumes that the model is predicting 1 as the minority class

        Args:
            actual_values: series with correct target variable assignments
            model_probability: probability of target variable assignment from the model

        Returns: AUC for a binary classifier

        """
        fpr, tpr, _ = roc_curve(actual_values, [y[1] for y in model_probability])
        return auc(fpr, tpr)

    def _get_data(self, data_object):
        """get the upstream data, split into X and Y from config, return data frame

        Returns:
            dataframe (dataframe)

        """

        # get upstream data dict which contains the subkey data_test or data_train
        data = data_object.get_filtered_upstream_data(
            self.instance_name, "data_test"
        ) or data_object.get_filtered_upstream_data(self.instance_name, "data_train")

        X_train = None
        if "data_train" in data:
            X_train = data["data_train"]
            if "features" in self.node_config:
                X_train = X_train[self.node_config["features"]]

        y_train = None
        if "target_train" in data:
            y_train = data["target_train"]

        X_test = None
        if "data_test" in data:
            X_test = data["data_test"]
            if "features" in self.node_config:
                X_test = X_test[self.node_config["features"]]

        y_test = None
        if "target_test" in data:
            y_test = data["target_test"]

        def size(obj):
            if obj is None:
                return None
            return str(obj.shape)

        logging.info(
            "Get_data, X_train %s, y_train %s, X_test %s, y_test %s",
            size(X_train),
            size(y_train),
            size(X_test),
            size(y_test),
        )

        return X_train, y_train, X_test, y_test

    def train_model(self, data_object):
        """train the model

        Args:
            data_object (DataObject): instance of DataObject

        Returns:
            data_object (DataObject): instance of DataObject

        """
        self.X_train, self.y_train, self.X_test, self.y_test = self._get_data(
            data_object
        )

        args = None
        if "args" in self.node_config["model"]:
            args = self.node_config["model"]["args"]

        self.model = self._instantiate_model(self.node_config["model"]["class"], args)

        logging.info("Fitting model")
        self.fit_training_data()  # delegate down as args to self.model.fit() vary by model

        data_object.add(self, self.model, "model")

        return data_object

    def _make_predictions(self, data_object):
        """Make predictions from X_test

        Args:
            data_object (DataObject): instance of DataObject

        Returns:
            nothing. Predictions stored in self.predictions

        """
        if hasattr(self, "predictions") and self.predictions is not None:
            return

        if self.model is None:
            self.model = self.load_model(data_object)

        self.X_train, self.y_train, self.X_test, self.y_test = self._get_data(
            data_object
        )

        logging.info("Making predictions with model")
        self.predictions = self.model.predict(self.X_test)

    def predict(self, data_object, load_model=False, use_serial=False):
        """Predict y_test from X_test

        Args:
            data_object (DataObject): instance of DataObject

            load_model: load model object from gcs or not

        Returns:
            data_object (DataObject): instance of DataObject

        """
        self._make_predictions(data_object)

        data = self.X_test
        data["predictions"] = self.predictions
        if self.y_test is not None:
            data["actual"] = self.y_test

        data_object.add(self, data)
        return data_object

    def eval_model(self, data_object, load_model=False):
        """evalute a model by getting the scores

        Returns:
            data_object (DataObject): instance of DataObject

        """
        self._make_predictions(data_object)

        logging.info("Evaluating metrics")
        scores = self.get_scores()  # delegate down as appropriate metrics vary by model
        scores["eval_time"] = datetime.datetime.now()
        logging.info("Scores" + str(scores))

        data_object.add(self, scores, "scores")
        return data_object

    @staticmethod
    def evaluate_no_ground_truth_classifier_metrics(X, labels):
        """Compute a set of metric for a classifier where there is no ground truth

        Args:
            X (datafframe): the data
            labels: the predicted classes

        Returns:
            dictionary of score name: value

        """
        scores = {}
        scores["Silhouette Score"] = metrics.silhouette_score(
            X, labels, metric="euclidean"
        )
        scores["Calinski Harabasz Score"] = metrics.calinski_harabasz_score(X, labels)
        scores["Davies Bouldin Score"] = metrics.davies_bouldin_score(X, labels)
        return scores

    @staticmethod
    def evaluate_regression_metrics(actual, predictions):
        """compute set of metrics for a regression

        Args:
            actual: vector of actual values
            predictions: vector of predictions

        Returns:
            dictionary of score name: value

        """
        scores = {}
        scores["Explained variance"] = metrics.explained_variance_score(
            actual, predictions
        )
        scores["Max error"] = metrics.max_error(actual, predictions)
        scores["Mean absolute error"] = metrics.mean_absolute_error(actual, predictions)
        scores["MSE"] = metrics.mean_squared_error(actual, predictions)
        scores["Mean squared log error"] = metrics.mean_squared_log_error(
            actual, predictions
        )
        scores["Mean absolute error"] = metrics.median_absolute_error(
            actual, predictions
        )
        scores["R2"] = metrics.r2_score(actual, predictions)
        return scores
