"""
Multivariate Gaussian distribution.
"""

import numpy as np
cimport numpy as np

import math

from libc.math cimport exp, log, sqrt
from libc.math cimport M_PI



from cryoem.data.griddata cimport GridDataReal



cdef class Gaussian(RealDistribution):
    """Gaussian distribution."""

    pass


cdef class Gaussian1D(Gaussian):
    """Univariate Gaussian distribution."""

    def __init__(self, mean, variance):
        """Initialise with mean and variance.

        Args:
            mean: The mean as a double.
            covariance: The variance as a double.
        """

        self._mean = mean
        self._variance = variance
        self._std = np.sqrt(variance)

    @property
    def mean(self):
        return self._mean

    @property
    def variance(self):
        return self._variance

    @property
    def std(self):
        return self._std

    @property
    def dimension(self):
        return 1

    def sample(self, num_samples=1):
        """Returns 'num_samples' samples from this distribution.

        Returns:
            An (N, ) array with N samples from the distribution.
        """

        return np.random.randn(num_samples) * self._std + self._mean

    def log_likelihood(self, x):
        """Returns sum of log of evaluation at each point in 'x'.

        Args:
            x: (N,) array of N points at which to evaluate Gaussian

        Returns:
            Sum of logs of values of Gaussian at every point.
        """

        t1 = np.sum((x - self._mean)**2) / (2.0 * self._variance)
        t2 = 0.5 * len(x) * np.log(2.0 * np.pi * self._variance)

        return -(t1 + t2)


    def log_likelihoods(self, x):
        """Returns log of evaluation at each point in 'x'.

        Args:
            x: (N, ) array of N points at which to evaluate Gaussian

        Returns:
            (N, ) array with log of value of Gaussian at every point
        """

        x_shifted = x - self._mean

        cost = (x_shifted * x_shifted) / self._variance

        constant = np.log(2*np.pi) + 2.0 * np.log(self._std)

        values = -0.5 * (constant + cost)

        return values

    cpdef evaluate_grid(self, Grid1D grid):
        """Return array with value of Gaussian at each grid point.

        Args:
            grid: Grid1D object.

        Return:
            GridDataReal object with value of Gaussian at each point.

        """

        cdef:
            int N, i
            double origin, spacing, x, d2, c1, c2
            double[:] values

        N = grid._shape[0] # num grid points
        origin = grid._origin[0]
        spacing = grid._spacing[0]
        c1 = 1.0 / sqrt(2.0 * M_PI * self._variance)
        c2 = -0.5 / self._variance

        values = np.zeros(N)

        x = origin
        for i in range(N):
            d2 = (self._mean - x)**2
            values[i] += c1 * exp(c2 * d2)
            x += spacing

        return GridDataReal(grid, values)


    def kl_divergence(self, g):
        """Returns KL divergence from this Gaussian to 'g'.

        Args:
            g: Second Gaussian in KL divergence definition.
                Gaussian1D object.
        """

        t1 = np.log(g.std / self._std)
        t2 = ( (self._variance + (self._mean - g.mean)**2) /
               (2.0 * g.variance) )
        return t1 + t2 - 0.5

    def get_quadrature(self, degree=5):
        """Returns points and weights for Gauss-Hermite quadrature.

        Args:
            degree: The degree of the quadrature rule, higher is
                more accurate, but slower.

        Returns:
            points, weights
            Both 1D arrays of the same length.
        """

        from commons.numerical.polynomial import hermgauss

        # Get the points and weights for integrating w.r.t. exp(-x^2).
        p, w = hermgauss(degree)
        # First convert to a standard Gaussian.
        p *= np.sqrt(2)
        w *= (1.0 / np.sqrt(np.pi))
        # Then convert to this specific Gaussian.
        p = p * self._std + self._mean

        return p, w

    def hellinger_distance(self, other):
        """Returns the Hellinger distance from this dist. to another.

        Args:
            other: Gaussian1D object.

        Using definition from Wikipedia, for some other conventions
        it is necessary to multiply this method's result by sqrt(2).
        """
        mu1 = self._mean
        mu2 = other.mean
        s1 = self._variance
        s2 = other.variance
        ss = s1*s1 + s2*s2
        d2 = 1.0 - np.sqrt(2.0*s1*s2/ss)*np.exp(-0.25*(mu1-mu2)**2/ss)

        return np.sqrt(d2)


cdef class GaussianND(Gaussian):
    """Multivariate Gaussian distribution."""

    def __init__(self, mean, covariance):
        """Initialise with mean vector and covariance matrix.

        Args:
            mean: The mean vector as a 1D array.
            covariance: The covariance matrix as a 2D array.
        """

        self._mean = mean
        self._covariance = covariance

    @property
    def mean(self):
        return self._mean

    @property
    def covariance(self):
        return self._covariance

    @property
    def dimension(self):
        return len(self._mean)

    def sample(self, num_samples=1):
        """Returns 'num_samples' samples from this distribution.

        Returns:
            A (N, D) array with N samples from the distribution.
        """

        return np.random.multivariate_normal(self._mean,
                                             self._covariance,
                                             num_samples)

    def log_likelihoods(self, x):
        """Returns log of evaluation at each point in 'x'.

        Args:
            x: (N, D) array of N points at which to evaluate Gaussian

        Returns:
            (N, ) array with log of value of Gaussian at every point
        """

        # Note: the cholesky decomposition sometimes returns
        # a LinAlgError: n-th leading minor no positive definite
        # Possible solution: add code to catch that error and
        # revert to the old way if it occurs.

        from scipy.linalg import cholesky
        from numpy.core.umath_tests import inner1d

        x_shifted = x - self._mean

        L = cholesky(self._covariance, lower=True)
        Q = np.linalg.solve(L, x_shifted.T) # (D, N)

        cost = inner1d(Q.T, Q.T)

        constant = self.dimension * np.log(2*np.pi)
        constant += 2.0 * np.sum(np.log(np.diag(L)))

        values = -0.5 * (constant + cost)

        return values

    @classmethod
    def from_data(cls, x):
        """Returns Gaussian with parameters estimated from 'x'.

        Args:
            x: (N, D) array of data points.

        Returns:
            GaussianND object with ML parameters.
        """

        N, D = x.shape
        mean = np.sum(x, axis=0) / N
        x_shifted = x - mean
        partial_cov = np.zeros([D, D])
        for i in range(N):
            vector = x_shifted[i]
            partial_cov += np.outer(vector, vector)
        covariance = partial_cov / N

        return cls(mean, covariance)

    def transform(self, transformation):
        """
        Returns transformed copy of model.

        Changes model parameters to describe points generated by 
        original one,
        then first rotated, and then translated.

        @param rotation: rotation matrix
        @type rotation: (d, d) array
        @param translation: translation vector
        @type translation: (d, ) array
        """

        r = transformation.rotation
        t = transformation.translation

        mean_transformed = np.dot(r, self._mean[:, None]).flatten() + t
        covariance_transformed = np.dot(
            r, np.dot(self._covariance, r.T))

        return Gaussian(mean_transformed, covariance_transformed)

    def project(self):
        """Returns projected Gaussian of one lower dimension.

        Projects to the z=0 hyperplane, i.e. by marginalising
        out the last component.

        Returns:
            GaussianND object of one lower dimension if original
            dimension was at least 3, otherwise Gaussian1D object.
        """

        mu_a = self._mean[:-1]
        A = self._covariance[:-1, :-1]
        if self.dimension >= 3:
            return GaussianND(mu_a, A)
        else:
            return Gaussian1D(np.double(mu_a), np.double(A))

    def _partition_parameters(self):
        """Necessary for computing conditionals."""

        mask = np.array([True] * self.dimension)
        mask[-1] = False

        mu_a = np.compress(mask, self._mean)
        mu_b = np.compress(1 - mask, self._mean)

        A = np.compress(
            mask, np.compress(mask, self._covariance, 1), 0)
        B = np.compress(1 - mask,
                        np.compress(1 - mask, self._covariance, 1), 0)
        C = np.compress(
            mask, np.compress(1 - mask, self._covariance, 1), 0)

        return mu_a, mu_b, A, B, C

    def conditionals_weighted_sum(self, xo, weights):
        """Returns weighted sum of conditionals.

        First conditions on the first D-1 components for each of the
        data points in 'xo' to obtain several ND distributions.
        Then compute the weighted sum of these distributions,
        i.e. the weighted sum of the means and covariances
        respectively.

        Args:
            xo: (N, D-1) array with observed components.
                If D=2, then 'xo' has shape (N, ).
            weights: (N, ) array with data point weights,
                such as assignments.

        Returns:
            sum_means, sum_second_moments
            Where 'sum_means' is a (D, ) array, and 'sum_second_moments'
            is a (D, D) array.
        """

        mu_a, mu_b, A, B, C = self._partition_parameters();

        # Calculate conditional mean.
        if self.dimension == 2:
            xo = xo[:, np.newaxis]
        xm = mu_b + np.dot(C.T, np.linalg.solve(A, (xo - mu_a).T)).T

        # Replace last column by conditional means.
        means = np.hstack([xo, xm])

        # Calculate weighted sum of means.
        sum_means = np.dot(means.T, weights)

        # Calculate conditional covariance matrices.
        xmm_central = B - np.dot(C.T, np.linalg.solve(A, C))
        # Central 2nd moments, just a single number.

        scaled_means = means * weights[:, np.newaxis]
        sum_second_moments = np.dot(scaled_means.T, means)

        xm2 = xm * xm
        xmm_non_central = xmm_central + xm2

        sum_second_moments[-1, -1] = np.inner(
            xmm_non_central.flatten(),
            weights)

        return sum_means, sum_second_moments

    def ellipse(self):
        """
        TODO: Move to Gaussian2D subclass.

        Calculate representative ellipse for 2D Gaussian.

        @return: parameters of ellipse
        """
        u, s, _v = np.linalg.svd(self._covariance)
        angle = math.atan2(u[1, 0], u[0, 0]) * 180.0 / np.pi
        scaling = 3.2     # arbitrary value to make it look good
        width = math.sqrt(s[0]) * scaling
        height = math.sqrt(s[1]) * scaling
        parameters = {'mean': self._mean,
                      'width': width,
                      'height': height,
                      'angle': angle}
        return parameters

    def kl_divergence(self, g):
        """Returns KL divergence from this Gaussian to 'g'.

        Args:
            g: Second Gaussian in KL divergence definition.
                GaussianND object.
        """

        from numpy.linalg import solve, det

        t1 = np.log(det(g.covariance) / det(self._covariance))
        t2 = np.trace(solve(g.covariance, self._covariance))
        mean_dif = self._mean - g.mean
        t3 = np.dot(
            mean_dif.T, solve(g.covariance, mean_dif)).flatten()

        return 0.5 * (t1 + t2 - self.dimension + t3)


cdef class IsoGaussianND(Gaussian):
    """Multivariate Gaussian distribution with isotropic covariance."""

    def __init__(self, mean, variance):
        """Initialise with mean vector and covariance matrix.

        Args:
            mean: The mean vector as a 1D array.
            variance: The variance which is the same in all directions.
                The resulting covariance matrix is 'variance' * eye(D)
        """

        self._mean = mean
        self._variance = variance

    @property
    def mean(self):
        return self._mean

    @property
    def variance(self):
        return self._variance

    def set_variance(self, value):
        self._variance = value

    @property
    def dimension(self):
        return len(self._mean)

    def sample(self, num_samples=1):
        """Returns 'num_samples' samples from this distribution.

        Returns:
            A (N, D) array with N samples from the distribution.
        """

        return (np.random.randn(num_samples, self.dimension) *
                np.sqrt(self._variance))

    def log_likelihood(self, x):
        """Returns sum of log of evaluation at each point in 'x'.

        Args:
            x: (N, D) array of N points at which to evaluate Gaussian

        Returns:
            Sum of logs of values of Gaussian at every point.
        """

        t1 = np.sum((x - self._mean)**2) / (2.0 * self._variance)
        t2 = 0.5 * len(x) * self.dimension * np.log(2.0 * np.pi *
                                                    self._variance)

        return -(t1 + t2)

    def log_likelihoods(self, x):
        """Returns log of evaluation at each point in 'x'.

        Args:
            x: (N, D) array of N points at which to evaluate Gaussian

        Returns:
            (N, ) array with log of value of Gaussian at every point
        """

        t1 = (np.sum((x - self._mean)**2, axis=1) /
              (2.0 * self._variance))
        t2 = 0.5 * self.dimension * np.log(
            2.0 * np.pi * self._variance)

        return -(t1 + t2)



