# -*- coding: utf-8 -*-

import operator

import numpy as np
import scipy.sparse as sp
from sklearn.base import copy
from sklearn.naive_bayes import (BaseDiscreteNB, BaseNB, BernoulliNB,
                                 GaussianNB, MultinomialNB)

from splearn import BlockRDD
from splearn.base import SparkClassifierMixin
from splearn.utils.validation import check_rdd


class SparkBaseNB(BaseNB, SparkClassifierMixin):

    """Abstract base class for distributed naive Bayes estimators"""

    def __radd__(self, other):
        return self if other == 0 else self.__add__(other)

    def predict(self, X):
        """
        Perform classification on an RDD containing arrays of test vectors X.

        Parameters
        ----------
        X : RDD containing array-like items, shape = [m_samples, n_features]

        Returns
        -------
        C : RDD with arrays, shape = [n_samples]
            Predicted target values for X
        """
        check_rdd(X, (sp.spmatrix, np.ndarray))
        return X.map(
            lambda X: super(SparkBaseNB, self).predict(X))

    def predict_proba(self, X):
        """
        Return probability estimates for the RDD containing test vector X.

        Parameters
        ----------
        X : RDD containing array-like items, shape = [m_samples, n_features]

        Returns
        -------
        C : RDD with array-like items , shape = [n_samples, n_classes]
            Returns the probability of the samples for each class in
            the models for each RDD block. The columns correspond to the classes
            in sorted order, as they appear in the attribute `classes_`.
        """
        check_rdd(X, (sp.spmatrix, np.ndarray))
        return X.map(
            lambda X: super(SparkBaseNB, self).predict_proba(X))

    def predict_log_proba(self, X):
        """
        Return log-probability estimates for the RDD containing the
        test vector X.

        Parameters
        ----------
        X : RDD containing array-like items, shape = [m_samples, n_features]

        Returns
        -------
        C : RDD with array-like items, shape = [n_samples, n_classes]
            Returns the log-probability of the samples for each class in
            the model for each RDD block. The columns correspond to the classes
            in sorted order, as they appear in the attribute `classes_`.
        """
        # required, scikit call self.predict_log_proba(X) in predict_proba
        # and thus this function is call, it must have the same behavior when
        # not called by sparkit-learn
        if not isinstance(X, BlockRDD):
            return super(SparkBaseNB, self).predict_log_proba(X)

        check_rdd(X, (sp.spmatrix, np.ndarray))
        return X.map(
            lambda X: super(SparkBaseNB, self).predict_log_proba(X))

    def _to_scikit(self, cls):
        new = cls()
        new.__dict__ = self.__dict__
        return new


class SparkGaussianNB(GaussianNB, SparkBaseNB):

    """
    Distributed Gaussian Naive Bayes (SparkGaussianNB)

    Based on sklearn's GaussianNB, model averaging is based on the work of
    the Stanford CS tech report STAN-CS-79-773 by Chan, Golub, and LeVeque:

        http://i.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf

    Attributes
    ----------
    class_prior_ : array, shape (n_classes,)
        probability of each class.

    class_count_ : array, shape (n_classes,)
        number of training samples observed in each class.

    theta_ : array, shape (n_classes, n_features)
        mean of each feature per class

    sigma_ : array, shape (n_classes, n_features)
        variance of each feature per class

    Examples
    --------
    TODO!
    """

    def fit(self, Z, classes=None):
        """Fit Gaussian Naive Bayes according to X, y

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training vectors, where n_samples is the number of samples
            and n_features is the number of features.

        y : array-like, shape (n_samples,)
            Target values.

        Returns
        -------
        self : object
            Returns self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray), 'y': (sp.spmatrix, np.ndarray)})
        models = Z[:, ['X', 'y']].map(
            lambda X_y: self.partial_fit(X_y[0], X_y[1], classes))
        avg = models.reduce(operator.add)
        self.__dict__.update(avg.__dict__)
        return self

    def __add__(self, other):
        this = copy.deepcopy(self)
        this_class2idx = {cls: idx for idx, cls in enumerate(this.classes_)}
        other_class2idx = {cls: idx for idx, cls in enumerate(other.classes_)}
        for class_i in this.classes_:
            i = this_class2idx[class_i]
            j = other_class2idx[class_i]
            N_x = this.class_count_[i]
            N_y = other.class_count_[j]
            mu_x = this.theta_[i, :]
            mu_y = other.theta_[j, :]
            sigma_x = this.sigma_[i, :]
            sigma_y = this.sigma_[j, :]
            N_total = N_x + N_y

            mu_xy = N_x * mu_x + N_y * mu_y
            sigma_xy = (sigma_x * N_x + sigma_y * N_y +
                        (N_x * N_y * (mu_x - mu_y) ** 2) / N_total)

            this.theta_[i, :] = mu_xy / N_total
            this.sigma_[i, :] = sigma_xy / N_total
            this.class_count_[i] += N_y

        this.class_prior_[:] = this.class_count_ / np.sum(this.class_count_)
        return this

    def to_scikit(self):
        return self._to_scikit(GaussianNB)


class SparkBaseDiscreteNB(BaseDiscreteNB, SparkBaseNB):

    """
    Abstract base class for distributed naive Bayes on discrete/categorical
    data. It provides the necessary methods for model averaging.

    """

    def __add__(self, other):
        """
        Add method for DiscreteNB models.

        Parameters
        ----------
        other : fitted splearn multinomilal NB model with class_count_
                and feature_count_ attribute
            Model to add.

        Returns
        -------
        model : splearn Naive Bayes model
            Model with updated coefficients.
        """
        # The rdd operator add does not consider __radd__ :(
        if other == 0:
            return self
        model = copy.deepcopy(self)
        model.class_count_ += other.class_count_
        model.feature_count_ += other.feature_count_
        model._update_class_log_prior()
        model._update_feature_log_prob()
        return model

    def fit(self, Z, classes=None):
        """
        TODO fulibacsi fix docstring
        Fit Multinomial Naive Bayes according to (X,y) pair
        which is zipped into TupleRDD Z.

        Parameters
        ----------
        Z : TupleRDD containing X [array-like, shape (m_samples, n_features)]
            and y [array-like, shape (m_samples,)] tuple
            Training vectors, where ,_samples is the number of samples in the
            block and n_features is the number of features, and y contains
            the target values.

        Returns
        -------
        self : object
            Returns self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray), 'y': (sp.spmatrix, np.ndarray)})
        if 'w' in Z.columns:
            models = Z[:, ['X', 'y', 'w']].map(
                lambda X_y_w: self.partial_fit(
                    X_y_w[0], X_y_w[1], classes, X_y_w[2]
                )
            )
        else:
            models = Z[:, ['X', 'y']].map(
                lambda X_y: self.partial_fit(X_y[0], X_y[1], classes))
        avg = models.sum()
        self.__dict__.update(avg.__dict__)
        return self


class SparkMultinomialNB(MultinomialNB, SparkBaseDiscreteNB):
    def to_scikit(self):
        return self._to_scikit(MultinomialNB)


class SparkBernoulliNB(BernoulliNB, SparkBaseDiscreteNB):
    def to_scikit(self):
        return self._to_scikit(BernoulliNB)
