#! /usr/bin/python
#
# Copyrighted David Cournapeau
# Last Change: Thu Jul 12 04:00 PM 2007 J
"""This module implements various basic functions related to multivariate
gaussian, such as pdf estimation, confidence interval/ellipsoids, etc..."""
from __future__ import absolute_import

__docformat__ = "restructuredtext"

import numpy as N
import numpy.linalg as lin

# from numpy.random import randn
from scipy.stats import chi2
from . import misc

# Error classes
class DenError(Exception):
    """Base class for exceptions in this module.
    
    Attributes:
        expression -- input expression in which the error occurred
        message -- explanation of the error"""

    def __init__(self, message):
        self.message = message
        Exception.__init__(self)

    def __str__(self):
        return self.message


# The following function do all the fancy stuff to check that parameters
# are Ok, and call the right implementation if args are OK.
def gauss_den(x, mu, va, log=False):
    """Compute multivariate Gaussian density at points x for 
    mean mu and variance va.
    
    :Parameters:
        x : ndarray
            points where to estimate the pdf.  each row of the array is one
            point of d dimension
        mu : ndarray
            mean of the pdf. Should have same dimension d than points in x.
        va : ndarray
            variance of the pdf. If va has d elements, va is interpreted as the
            diagonal elements of the actual covariance matrix. Otherwise,
            should be a dxd matrix (and positive definite).
        log : boolean
            if True, returns the log-pdf instead of the pdf.

    :Returns:
        pdf : ndarray
            Returns a rank 1 array of the pdf at points x.

    Note
    ----
        Vector are row vectors, except va which can be a matrix
        (row vector variance for diagonal variance)."""

    lmu = N.atleast_2d(mu)
    lva = N.atleast_2d(va)
    lx = N.atleast_2d(x)

    # =======================#
    # Checking parameters   #
    # =======================#
    if len(N.shape(lmu)) != 2:
        raise DenError("mu is not rank 2")

    if len(N.shape(lva)) != 2:
        raise DenError("va is not rank 2")

    if len(N.shape(lx)) != 2:
        raise DenError("x is not rank 2")

    d = N.shape(lx)[1]
    (dm0, dm1) = N.shape(lmu)
    (dv0, dv1) = N.shape(lva)

    # Check x and mu same dimension
    if dm0 != 1:
        msg = "mean must be a row vector!"
        raise DenError(msg)
    if dm1 != d:
        msg = "x and mu not same dim"
        raise DenError(msg)
    # Check va and mu same size
    if dv1 != d:
        msg = "mu and va not same dim"
        raise DenError(msg)
    if dv0 != 1 and dv0 != d:
        msg = "va not square"
        raise DenError(msg)

    # ===============#
    # Computation   #
    # ===============#
    if d == 1:
        # scalar case
        return _scalar_gauss_den(lx[:, 0], lmu[0, 0], lva[0, 0], log)
    elif dv0 == 1:
        # Diagonal matrix case
        return _diag_gauss_den(lx, lmu, lva, log)
    elif dv1 == dv0:
        # full case
        return _full_gauss_den(lx, lmu, lva, log)
    else:
        raise DenError("variance mode not recognized, this is a bug")


# Those 3 functions do almost all the actual computation
def _scalar_gauss_den(x, mu, va, log):
    """ This function is the actual implementation
    of gaussian pdf in scalar case. It assumes all args
    are conformant, so it should not be used directly
    
    Call gauss_den instead"""
    d = mu.size
    inva = 1 / va
    fac = (2 * N.pi) ** (-d / 2.0) * N.sqrt(inva)
    inva *= -0.5
    y = ((x - mu) ** 2) * inva
    if not log:
        y = fac * N.exp(y)
    else:
        y += N.log(fac)

    return y


def _diag_gauss_den(x, mu, va, log):
    """ This function is the actual implementation
    of gaussian pdf in scalar case. It assumes all args
    are conformant, so it should not be used directly
    
    Call gauss_den instead"""
    # Diagonal matrix case
    d = mu.size
    # n   = x.shape[0]
    if not log:
        inva = 1 / va[0]
        fac = (2 * N.pi) ** (-d / 2.0) * N.prod(N.sqrt(inva))
        inva *= -0.5
        x = x - mu
        x **= 2
        y = fac * N.exp(N.dot(x, inva))
    else:
        # XXX optimize log case as non log case above
        y = _scalar_gauss_den(x[:, 0], mu[0, 0], va[0, 0], log)
        for i in range(1, d):
            y += _scalar_gauss_den(x[:, i], mu[0, i], va[0, i], log)
    return y


def _full_gauss_den(x, mu, va, log):
    """ This function is the actual implementation
    of gaussian pdf in full matrix case. 
    
    It assumes all args are conformant, so it should 
    not be used directly Call gauss_den instead
    
    Does not check if va is definite positive (on inversible 
    for that matter), so the inverse computation and/or determinant
    would throw an exception."""
    d = mu.size
    inva = lin.inv(va)
    fac = 1 / N.sqrt((2 * N.pi) ** d * N.fabs(lin.det(va)))

    # we are using a trick with sum to "emulate"
    # the matrix multiplication inva * x without any explicit loop
    # y   = -0.5 * N.sum(N.dot((x-mu), inva) * (x-mu), 1)
    y = (
        -0.5
        * N.dot(N.dot((x - mu), inva) * (x - mu), N.ones((mu.size, 1), x.dtype))[:, 0]
    )

    if not log:
        y = fac * N.exp(y)
    else:
        y = y + N.log(fac)

    return y


# To get coordinatea of a confidence ellipse from multi-variate gaussian pdf
def gauss_ell(
    mu, va, dim=misc.DEF_VIS_DIM, npoints=misc.DEF_ELL_NP, level=misc.DEF_LEVEL
):
    """Given a mean and covariance for multi-variate
    gaussian, returns the coordinates of the confidense ellipsoid.
    
    Compute npoints coordinates for the ellipse of confidence of given level
    (all points will be inside the ellipsoides with a probability equal to
    level).
    
    :Parameters:
        mu : ndarray
            mean of the pdf
        va : ndarray
            variance of the pdf
        dim : sequence
            sequences of two integers which represent the dimensions where to
            project the ellipsoid.
        npoints: int
            number of points to generate for the ellipse.
        level : float
            level of confidence (between 0 and 1).

    :Returns:
        Returns the coordinate x and y of the ellipse."""
    if level >= 1 or level <= 0:
        raise ValueError("level should be a scale strictly between 0 and 1." "")

    mu = N.atleast_1d(mu)
    va = N.atleast_1d(va)
    d = N.shape(mu)[0]
    c = N.array(dim)

    if N.any(c < 0) or N.any(c >= d):
        raise ValueError(
            "dim elements should be >= 0 and < %d (dimension" " of the variance)" % d
        )
    if N.size(mu) == N.size(va):
        mode = "diag"
    else:
        if N.ndim(va) == 2:
            if N.shape(va)[0] == N.shape(va)[1]:
                mode = "full"
            else:
                raise DenError("variance not square")
        else:
            raise DenError("mean and variance are not dim conformant")

    # When X is a sample from multivariante N(mu, sigma), (X-mu)Sigma^-1(X-mu)
    # follows a Chi2(d) law. Here, we only take 2 dimension, so Chi2 with 2
    # degree of freedom (See Wasserman. This is easy to see with characteristic
    # functions)
    chi22d = chi2(2)
    mahal = N.sqrt(chi22d.ppf(level))

    # Generates a circle of npoints
    theta = N.linspace(0, 2 * N.pi, npoints)
    circle = mahal * N.array([N.cos(theta), N.sin(theta)])

    # Get the dimension which we are interested in:
    mu = mu[c]
    if mode == "diag":
        va = va[c]
        elps = N.outer(mu, N.ones(npoints))
        elps += N.dot(N.diag(N.sqrt(va)), circle)
    elif mode == "full":
        va = va[c, :][:, c]
        # Method: compute the cholesky decomp of each cov matrix, that is
        # compute cova such as va = cova * cova'
        # WARN: scipy is different than matlab here, as scipy computes a lower
        # triangular cholesky decomp:
        #   - va = cova * cova' (scipy)
        #   - va = cova' * cova (matlab)
        # So take care when comparing results with matlab !
        cova = lin.cholesky(va)
        elps = N.outer(mu, N.ones(npoints))
        elps += N.dot(cova, circle)
    else:
        raise ValueError("var mode not recognized")

    return elps[0, :], elps[1, :]


def logsumexp(x):
    """Compute log(sum(exp(x), 1)) while avoiding underflow.
    
    :Parameters:
        x : ndarray
            data in log domain to sum"""
    axis = 1
    mc = N.max(x, axis)
    return mc + N.log(N.sum(N.exp(x - mc[:, N.newaxis]), axis))


def multiple_gauss_den(data, mu, va, log=False):
    """Helper function to generate several Gaussian
    pdf (different parameters) at the same points

    :Parameters:
        data : ndarray
            points where to estimate the pdfs (n,d).
        mu : ndarray
            mean of the pdf, of shape (k,d). One row of dimension d per
            different component, the number of rows k being the number of
            component
        va : ndarray
            variance of the pdf. One row per different component for diagonal
            covariance (k, d), or d rows per component for full matrix pdf
            (k*d,d).
        log : boolean
            if True, returns the log-pdf instead of the pdf.

    :Returns:
        Returns a (n, k) array, each column i being the pdf of the ith mean and
        ith variance."""
    mu = N.atleast_2d(mu)
    va = N.atleast_2d(va)

    k = N.shape(mu)[0]
    n = N.shape(data)[0]
    d = N.shape(mu)[1]

    y = N.zeros((k, n))
    if N.size(mu) == N.size(va):
        for i in range(k):
            y[i] = gauss_den(data, mu[i, :], va[i, :], log)
        return y.T
    else:
        for i in range(k):
            y[i] = gauss_den(data, mu[i, :], va[d * i : d * i + d, :], log)
        return y.T


if __name__ == "__main__":
    pass
