"""
IBM Confidential
OCO Source Materials
5737-H76, 5725-W78, 5900-A1R
(c) Copyright IBM Corp. 2020 All Rights Reserved.
The source code for this program is not published or otherwise divested of its trade secrets,
irrespective of what has been deposited with the U.S. Copyright Office.
"""
import joblib
import numpy as np
import logging

from diffprivlib.models import GaussianNB

from ibmfl.exceptions import ModelInitializationException, ModelException
from ibmfl.model.model_update import ModelUpdate
from ibmfl.model.sklearn_fl_model import SklearnFLModel

logger = logging.getLogger(__name__)


class NaiveBayesFLModel(SklearnFLModel):
    """
    Wrapper class for diffprivlib.models.GaussianNB classifier, which itself
    is an implementation of sklearn.naive_bayes.GaussianNB with differential
    privacy.
    """

    def __init__(self, model_name, model_spec, model=None, **kwargs):
        """
        Create a `NaiveBayesFLModel` instance from a diffprivlib.models.GaussianNB
        model. If model is provided, it will use it; otherwise it will take
        the model_spec to create the model.

        :param model_name: A name specifying the type of model, e.g., \
        naive_bayes
        :type model_name: `str`
        :param model_spec: A dictionary contains model specification
        :type model_spec: `dict`
        :param model: Compiled GaussianNB model
        :type model: `diffprivlib.models.GaussianNB`
        :param kwargs: A dictionary contains other parameter settings on \
         to initialize a Naive Bayes model.
        :type kwargs: `dict`
        """
        if model:
            if not issubclass(type(model), GaussianNB):
                raise ValueError('Compiled GaussianNB needs to be provided'
                                 '(diffprivlib.models.GaussianNB). '
                                 'Type provided' + str(type(model)))

        super().__init__(model_name,
                         model_spec,
                         sklearn_model=model,
                         **kwargs)

        self.old_vals = {"theta": None, "sigma": None, "class_count": None}
        self.model_trained = False
        self.model_type = 'Sklearn-NaiveBayes'

    def fit_model(self, train_data, fit_params=None, **kwargs):
        """
        Fits current model with provided training data.

        :param train_data: Training data a tuple given in the form \
        (x_train, y_train).
        :type train_data: `np.ndarray`
        :param fit_params: Not used, present here for API consistency by \
        convention.
        :return: None
        """
        self.model.fit(train_data[0], train_data[1])
        self.model_trained = True

    def update_model(self, model_update, **kwargs):
        """
        Update GaussianNB model with provided model_update.

        :param model_update: `ModelUpdate` object that contains \
        the required information to update the GaussianNB model.
        :type model_update: `ModelUpdate`
        :return: None
        """
        if model_update.get("class_count") is None:
            return

        self.model.theta_ = model_update.get("theta")
        self.model.sigma_ = model_update.get("sigma")
        self.model.class_count_ = model_update.get("class_count")

        self._store_old_values(self.model.theta_, self.model.sigma_,
                               self.model.class_count_)

    def get_model_update(self):
        """
        Generates a `ModelUpdate` object that will be sent to other entities.

        :return: ModelUpdate
        :rtype: `ModelUpdate`
        """
        if (not self.model_trained or
                not hasattr(self.model, "class_count_") or
                np.all(self.old_vals['class_count'] ==
                       self.model.class_count_)):
            return ModelUpdate(theta=None, sigma=None, class_count=None)
        if self.old_vals['class_count'] is None:
            return ModelUpdate(theta=self.model.theta_,
                               sigma=self.model.sigma_,
                               class_count=self.model.class_count_)

        class_count_update = self.model.class_count_ - \
            self.old_vals['class_count']

        nm = self.model.class_count_[:, np.newaxis]
        n = self.old_vals['class_count'][:, np.newaxis]
        m = class_count_update[:, np.newaxis]

        theta_update = self.model.theta_ * nm - self.old_vals['theta'] * n
        theta_update = theta_update / m

        sigma_update = nm * self.model.sigma_ - n * self.old_vals['sigma'] - \
            (n * m) / nm * (self.old_vals['theta'] - theta_update) ** 2

        return ModelUpdate(theta=theta_update,
                           sigma=sigma_update,
                           class_count=class_count_update)

    def evaluate(self, test_dataset, **kwargs):
        """
        Evaluates the model given testing data.
        :param test_dataset: Testing data, a tuple given in the form \
        (x_test, test) or a datagenerator of of type `keras.utils.Sequence`, 
        `keras.preprocessing.image.ImageDataGenerator`
        :type test_dataset: `np.ndarray`

        :param kwargs: Dictionary of metrics available for the model
        :type kwargs: `dict`
        """

        if type(test_dataset) is tuple:
            x_test = test_dataset[0]
            y_test = test_dataset[1]

            return self.evaluate_model(x_test, y_test)

        else:
            raise ModelException("Invalid test dataset!")

    def evaluate_model(self, x, y, **kwargs):
        """
        Evaluates the model given test data x and the corresponding labels y.

        :param x: Samples with shape as expected by the model.
        :type x: `np.ndarray`
        :param y: Corresponding true labels to x
        :type y: `np.ndarray`
        :param kwargs: Dictionary of model-specific arguments for evaluating \
        models. For example, sample weights accepted by model.score.
        :return: Dictionary with all evaluation metrics provided by \
        specific implementation.
        :rtype: `dict`
        """
        acc = {'score': self.model.score(x, y, **kwargs)}
        return acc

    def _store_old_values(self, theta, sigma, class_count):
        """
        Store old values of training parameters to allow reconstruction of
        parameters for newly-trained examples when model update is sought.
        """
        self.old_vals = {"theta": theta,
                         "sigma": sigma,
                         "class_count": class_count}

    @staticmethod
    def load_model_from_spec(model_spec):
        """
        Loads model from provided model_spec, where model_spec is a `dict`
        that contains the following items: model_spec['model_definition']
        contains the model definition as type diffprivlib.GaussianNB.

        :param model_spec: Model specification contains \
        a complied sklearn model.
        :param model_spec: `dict`
        :return: model
        :rtype: `sklearn.cluster`
        """
        model = None
        try:
            if 'model_definition' in model_spec:
                path = model_spec['model_definition']
                with open(path, 'rb') as f:
                    model = joblib.load(f)

                if not issubclass(type(model), GaussianNB):
                    raise ValueError('Provided complied model in model_spec '
                                     'should be of type GaussianNB.'
                                     'Instead they are:' + str(type(model)))
        except Exception as ex:
            raise ModelInitializationException('Model specification was '
                                               'badly formed. ', str(ex))
        return model
