import numpy as np
import theano.tensor as tt

from pymc3.util import get_variable_name
from ..math import logsumexp
from .dist_math import bound, random_choice
from .distribution import (Discrete, Distribution, draw_values,
                           generate_samples, _DrawValuesContext)
from .continuous import get_tau_sd, Normal


def all_discrete(comp_dists):
    """
    Determine if all distributions in comp_dists are discrete
    """
    if isinstance(comp_dists, Distribution):
        return isinstance(comp_dists, Discrete)
    else:
        return all(isinstance(comp_dist, Discrete) for comp_dist in comp_dists)


class Mixture(Distribution):
    R"""
    Mixture log-likelihood

    Often used to model subpopulation heterogeneity

    .. math:: f(x \mid w, \theta) = \sum_{i = 1}^n w_i f_i(x \mid \theta_i)

    ========  ============================================
    Support   :math:`\cap_{i = 1}^n \textrm{support}(f_i)`
    Mean      :math:`\sum_{i = 1}^n w_i \mu_i`
    ========  ============================================

    Parameters
    ----------
    w : array of floats
        w >= 0 and w <= 1
        the mixture weights
    comp_dists : multidimensional PyMC3 distribution (e.g. `pm.Poisson.dist(...)`)
        or iterable of one-dimensional PyMC3 distributions the
        component distributions :math:`f_1, \ldots, f_n`

    Example
    -------
    .. code-block:: python

        # 2-Mixture Poisson distribution
        with pm.Model() as model:
            lam = pm.Exponential('lam', lam=1, shape=(2,))  # `shape=(2,)` indicates two mixtures.

            # As we just need the logp, rather than add a RV to the model, we need to call .dist()
            components = pm.Poisson.dist(mu=lam, shape=(2,))

            w = pm.Dirichlet('w', a=np.array([1, 1]))  # two mixture component weights.

            like = pm.Mixture('like', w=w, comp_dists=components, observed=data)

        # 2-Mixture Poisson using iterable of distributions.
        with pm.Model() as model:
            lam1 = pm.Exponential('lam1', lam=1)
            lam2 = pm.Exponential('lam2', lam=1)

            pois1 = pm.Poisson.dist(mu=lam1)
            pois2 = pm.Poisson.dist(mu=lam2)

            w = pm.Dirichlet('w', a=np.array([1, 1]))

            like = pm.Mixture('like', w=w, comp_dists = [pois1, pois2], observed=data)
    """

    def __init__(self, w, comp_dists, *args, **kwargs):
        shape = kwargs.pop('shape', ())

        self.w = w = tt.as_tensor_variable(w)
        self.comp_dists = comp_dists

        defaults = kwargs.pop('defaults', [])

        if all_discrete(comp_dists):
            dtype = kwargs.pop('dtype', 'int64')
        else:
            dtype = kwargs.pop('dtype', 'float64')

            try:
                self.mean = (w * self._comp_means()).sum(axis=-1)

                if 'mean' not in defaults:
                    defaults.append('mean')
            except AttributeError:
                pass

        try:
            comp_modes = self._comp_modes()
            comp_mode_logps = self.logp(comp_modes)
            self.mode = comp_modes[tt.argmax(w * comp_mode_logps, axis=-1)]

            if 'mode' not in defaults:
                defaults.append('mode')
        except (AttributeError, ValueError, IndexError):
            pass

        super(Mixture, self).__init__(shape, dtype, defaults=defaults,
                                      *args, **kwargs)

    def _comp_logp(self, value):
        comp_dists = self.comp_dists

        try:
            value_ = value if value.ndim > 1 else tt.shape_padright(value)

            return comp_dists.logp(value_)
        except AttributeError:
            return tt.squeeze(tt.stack([comp_dist.logp(value)
                                        for comp_dist in comp_dists],
                                       axis=1))

    def _comp_means(self):
        try:
            return tt.as_tensor_variable(self.comp_dists.mean)
        except AttributeError:
            return tt.squeeze(tt.stack([comp_dist.mean
                                        for comp_dist in self.comp_dists],
                                       axis=1))

    def _comp_modes(self):
        try:
            return tt.as_tensor_variable(self.comp_dists.mode)
        except AttributeError:
            return tt.squeeze(tt.stack([comp_dist.mode
                                        for comp_dist in self.comp_dists],
                                       axis=1))

    def _comp_samples(self, point=None, size=None):
        try:
            samples = self.comp_dists.random(point=point, size=size)
        except AttributeError:
            samples = np.column_stack([comp_dist.random(point=point, size=size)
                                       for comp_dist in self.comp_dists])

        return np.squeeze(samples)

    def logp(self, value):
        w = self.w

        return bound(logsumexp(tt.log(w) + self._comp_logp(value), axis=-1),
                     w >= 0, w <= 1, tt.allclose(w.sum(axis=-1), 1),
                     broadcast_conditions=False)

    def random(self, point=None, size=None):
        with _DrawValuesContext() as draw_context:
            w = draw_values([self.w], point=point)[0]
            comp_tmp = self._comp_samples(point=point, size=None)
        if np.asarray(self.shape).size == 0:
            distshape = np.asarray(np.broadcast(w, comp_tmp).shape)[..., :-1]
        else:
            distshape = np.asarray(self.shape)

        # Normalize inputs
        w /= w.sum(axis=-1, keepdims=True)

        w_samples = generate_samples(random_choice,
                                     p=w,
                                     broadcast_shape=w.shape[:-1] or (1,),
                                     dist_shape=distshape,
                                     size=size).squeeze()
        if (size is None) or (distshape.size == 0):
            with draw_context:
                comp_samples = self._comp_samples(point=point, size=size)
            if comp_samples.ndim > 1:
                samples = np.squeeze(comp_samples[np.arange(w_samples.size), ..., w_samples])
            else:
                samples = np.squeeze(comp_samples[w_samples])
        else:
            if w_samples.ndim == 1:
                w_samples = np.reshape(np.tile(w_samples, size), (size,) + w_samples.shape)
            samples = np.zeros((size,)+tuple(distshape))
            with draw_context:
                for i in range(size):
                    w_tmp = w_samples[i, :]
                    comp_tmp = self._comp_samples(point=point, size=None)
                    if comp_tmp.ndim > 1:
                        samples[i, :] = np.squeeze(comp_tmp[np.arange(w_tmp.size), ..., w_tmp])
                    else:
                        samples[i, :] = np.squeeze(comp_tmp[w_tmp])

        return samples


class NormalMixture(Mixture):
    R"""
    Normal mixture log-likelihood

    .. math::

        f(x \mid w, \mu, \sigma^2) = \sum_{i = 1}^n w_i N(x \mid \mu_i, \sigma^2_i)

    ========  =======================================
    Support   :math:`x \in \mathbb{R}`
    Mean      :math:`\sum_{i = 1}^n w_i \mu_i`
    Variance  :math:`\sum_{i = 1}^n w_i^2 \sigma^2_i`
    ========  =======================================

    Parameters
    ----------
    w : array of floats
        w >= 0 and w <= 1
        the mixture weights
    mu : array of floats
        the component means
    sd : array of floats
        the component standard deviations
    tau : array of floats
        the component precisions
    comp_shape : shape of the Normal component
        notice that it should be different than the shape
        of the mixture distribution, with one axis being
        the number of components.

    Note: You only have to pass in sd or tau, but not both.
    """

    def __init__(self, w, mu, comp_shape=(), *args, **kwargs):
        _, sd = get_tau_sd(tau=kwargs.pop('tau', None),
                           sd=kwargs.pop('sd', None))

        self.mu = mu = tt.as_tensor_variable(mu)
        self.sd = sd = tt.as_tensor_variable(sd)

        super(NormalMixture, self).__init__(w, Normal.dist(mu, sd=sd, shape=comp_shape),
                                            *args, **kwargs)

    def _repr_latex_(self, name=None, dist=None):
        if dist is None:
            dist = self
        mu = dist.mu
        w = dist.w
        sd = dist.sd
        name = r'\text{%s}' % name
        return r'${} \sim \text{{NormalMixture}}(\mathit{{w}}={},~\mathit{{mu}}={},~\mathit{{sigma}}={})$'.format(name,
                                                get_variable_name(w),
                                                get_variable_name(mu),
                                                get_variable_name(sd))
