"""
Implements Monte Carlo inference for Gaussian Mixtures:
    Y[i·]|(Z, mu, sig2) ~ Normal(mu[Z[i·]·], sig2[Z[i·]·])
    Z[i·]|pi ~ Categorical(pi)
    where Z is latent.

Data
----
Y : np.ndarray in R^(nobs, nvar)

Latent
------
Z : np.ndarray in (0:ncomp-1)^nobs

Parameters
----------
pi : np.ndarray in Simplex_ncomp
mu : np.ndarray in R^(ncomp, nvar)
sig2 : np.ndarray in R+^(ncomp, nvar)

Hyperparameters
---------------
comp0 : tuple(np.ndarray in R+^ncomp,)
    component distribution multinomial hyperparameters (a0,)
emit0 : tuple (np.ndarray in R^(ncomp, nvar), 3 * np.ndarray in R+^(ncomp, nvar))
    emission distribution normal hyperparameters (m0, l0, v0, s0)
"""

import numpy as np
from ...lib.gibbs import em, gibbs
from ..conjugate import normal as nor
from ..semi_conjugate import cluster as clu


def maximize(niter, ncomp, Y, comp0, emit0):
    """Find the MAP estimator of the model with the EM algorithm.

    Parameters
    ----------
    niter : int in N+
        number of iterations to run
    ncomp : int in N+
        number of clusters
    Y, comp0, emit0 : see module docstring

    Returns
    -------
    tuple (4 * np.ndarray)
        parameter modes (Z, pi, mu, sig2)
    """

    init = np.random.dirichlet(np.ones(ncomp), size=Y.shape[0])

    return em(
        niter, (Y,), (comp0, emit0), init,
        eval_logobserved, update_latent, update_param)


def sample(ndraws, ncomp, Y, comp0, emit0):
    """Draw Gibbs samples from the parameter distribution given state and hyperparameters.

    Parameters
    ----------
    ndraws : int in N+
        number of draws to be sampled
    ncomp : int in N+
        number of clusters
    Y, comp0, emit0 : see module docstring

    Returns
    -------
    tuple (4 * np.ndarray)
        parameter draws (Z, pi, mu, sig2)
    """

    init = maximize(100, ncomp, Y, comp0, emit0)

    return gibbs(
        ndraws, (Y,), (comp0, emit0), (np.argmax(init[0], 1), *init[1:]),
        sample_latent, sample_param)


def eval_logobserved(Y, pi, mu, sig2):
    """Evaluate the observed likelihood.

    Parameters
    ----------
    Y, pi, mu, sig2 : see module docstring

    Returns
    -------
    float
        log observed likelihood
    """

    return clu.eval_logobserved(pi, eval_emit(Y, mu, sig2))


def update_latent(Y, pi, mu, sig2):
    """Compute the expectation of latent given data and parameters.

    Parameters
    ----------
    Y, pi, mu, sig2 : see module docstring

    Returns
    -------
    np.ndarray in Simplex_ncomp^nobs
        posterior means
    """

    return np.exp(clu.update_latent(pi, eval_emit(Y, mu, sig2)))


def update_param(Z, Y, comp0, emit0):
    """Compute mode the mode of parameters given data and latent.
    
    Parameters
    ----------
    Z, Y, comp0, emit0 : see module docstring
    
    Returns
    -------
    tuple (3 * np.ndarray)
        modes (pi, mu, sig2)
    """

    pi = clu.get_mode(Z.shape[1], Z, *comp0)[0]
    mu = np.empty((Z.shape[1], Y.shape[1]))
    sig2 = np.empty((Z.shape[1], Y.shape[1]))
    
    for k in range(Z.shape[1]):
        weights = np.broadcast_to(Z[:,(k,)], Y.shape)
        mu[k], sig2[k] = nor.get_mode(
            *nor.weighted_update(Y, weights, *[x[k] for x in emit0]))
    
    return (pi, mu, sig2)
    

def sample_latent(Y, pi, mu, sig2):
    """Sample from latent given data and parameters.

    Parameters
    ----------
    Y, pi, mu, sig2 : see module docstring

    Returns
    -------
    np.ndarray in (0:ncomp-1)^nobs
        sample
    """

    return clu.sample_latent(pi, eval_emit(Y, mu, sig2))


def sample_param(Z, Y, comp0, emit0):
    """Sample from parameters given data and latent.
    
    Parameters
    ----------
    Z, Y, comp0 : see module docstring
    
    Returns
    -------
    tuple (3 * np.ndarray)
        samples (pi, mu, sig2)
    """

    pi = clu.sample_param(1, Z.shape[1], Z, *comp0)[0]
    mu = np.empty((Z.shape[1], Y.shape[1]))
    sig2 = np.empty((Z.shape[1], Y.shape[1]))
    
    for k in range(Z.shape[1]):
        mu[(k,)], sig2[(k,)] = nor.sample_param(
            1, *nor.update(Y[Z == k], *[x[k] for x in emit0]))

    return (pi, mu, sig2)


def eval_loglik(Y, Z, pi, mu, sig2):
    """Evaluate the log likelihood given data and parameters.

    Parameters
    ----------
    Y, Z, pi, mu, sig2 : see module docstring

    Returns
    -------
    np.ndarray in R^nobs
        log density at each state and observation
    """

    return np.take(eval_emit(Y, mu, sig2), Z)


def eval_emit(Y, mu, sig2):
    """Evaluate the log likelihood given data and parameters for all states.

    Parameters
    ----------
    Y, mu, sig2 : see module docstring

    Returns
    -------
    np.ndarray in R^(nobs, ncomp)
        log density at each state and observation
    """

    return np.array([
        nor.eval_loglik(Y, mu[k], sig2[k])
        for k in range(mu.shape[0])]).T
