import numpy as np
import dask.array as da
from dask import delayed
from sklearn.base import BaseEstimator
from sklearn import naive_bayes as _naive_bayes

from dask_ml.base import _BigPartialFitMixin, _copy_partial_doc


class GaussianNB(BaseEstimator):
    """
    Fit a naive bayes model with a Gaussian likelihood

    Examples
    --------
    >>> from dask_ml import datasets
    >>> from dask_ml.naive_bayes import GaussianNB
    >>> X, y = datasets.make_classification(chunks=10)
    >>> gnb = GaussianNB()
    >>> gnb.fit(X, y)
    """

    def __init__(self, priors=None, classes=None):
        self.priors = priors
        self.classes = classes
        self.classes_ = classes
        self.class_prior_ = None
        self.class_count_ = None
        self.theta_ = None
        self.sigma_ = None

    def fit(self, X, y=None):
        if self.classes is None:
            # TODO: delayed
            self.classes_ = np.unique(y)

        thetas = []
        sigmas = []
        counts = []
        N, P = X.shape
        K = len(self.classes_)

        for i, c in enumerate(self.classes_):
            X_c = X[y == c]
            thetas.append(X_c.mean(axis=0))
            sigmas.append(X_c.var(axis=0))
            counts.append(delayed(len)(X_c))

        thetas = da.from_delayed(delayed(np.array)(thetas), (K, P), np.float64)
        sigmas = da.from_delayed(delayed(np.array)(sigmas), (K, P), np.float64)
        counts = da.from_delayed(delayed(np.array)(counts, np.float64),
                                 (P,), np.float64)
        priors = counts / N

        # Should these be explicitly cached on self?
        self.theta_ = thetas
        self.sigma_ = sigmas
        self.class_count_ = counts
        self.class_prior_ = priors

        return self

    def predict(self, X):
        """
        Perform classification on an array of test vectors X.
        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
        Returns
        -------
        C : array, shape = [n_samples]
            Predicted target values for X
        """
        jll = self._joint_log_likelihood(X)
        return delayed(self.classes_)[da.argmax(jll, axis=1)]

    def predict_log_proba(self, X):
        """
        Return log-probability estimates for the test vector X.
        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
        Returns
        -------
        C : array-like, shape = [n_samples, n_classes]
            Returns the log-probability of the samples for each class in
            the model. The columns correspond to the classes in sorted
            order, as they appear in the attribute `classes_`.
        """
        jll = self._joint_log_likelihood(X)
        # normalize by P(x) = P(f_1, ..., f_n)
        log_prob_x = logsumexp(jll, axis=1)
        return (jll - log_prob_x.reshape(-1, 1))

    def predict_proba(self, X):
        """
        Return probability estimates for the test vector X.
        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]
        Returns
        -------
        C : array-like, shape = [n_samples, n_classes]
            Returns the probability of the samples for each class in
            the model. The columns correspond to the classes in sorted
            order, as they appear in the attribute `classes_`.
        """
        return da.exp(self.predict_log_proba(X))

    def _joint_log_likelihood(self, X):
        jll = []
        for i in range(np.size(self.classes_)):
            jointi = da.log(self.class_prior_[i])
            n_ij = - 0.5 * da.sum(da.log(2. * np.pi * self.sigma_[i, :]))
            n_ij -= 0.5 * da.sum(((X - self.theta_[i, :]) ** 2) /
                                 (self.sigma_[i, :]), 1)
            jll.append(jointi + n_ij)

        joint_log_likelihood = da.stack(jll).T
        return joint_log_likelihood


@_copy_partial_doc
class PartialMultinomialNB(_BigPartialFitMixin, _naive_bayes.MultinomialNB):
    _init_kwargs = ['classes']
    _fit_kwargs = ['classes']


@_copy_partial_doc
class PartialBernoulliNB(_BigPartialFitMixin, _naive_bayes.BernoulliNB):
    _init_kwargs = ['classes']
    _fit_kwargs = ['classes']


def logsumexp(arr, axis=0):
    """Computes the sum of arr assuming arr is in the log domain.
    Returns log(sum(exp(arr))) while minimizing the possibility of
    over/underflow.
    Examples
    --------
    >>> import numpy as np
    >>> from sklearn.utils.extmath import logsumexp
    >>> a = np.arange(10)
    >>> np.log(np.sum(np.exp(a)))
    9.4586297444267107
    >>> logsumexp(a)
    9.4586297444267107
    """
    if axis == 0:
        pass
    elif axis == 1:
        arr = arr.T
    else:
        raise NotImplementedError
    # Use the max to normalize, as with the log this is what accumulates
    # the less errors
    vmax = arr.max(axis=0)
    out = da.log(da.sum(da.exp(arr - vmax), axis=0))
    out += vmax
    return out
