"""Isotropic Gaussian mixture model.

All components have the same variance.
"""

import numpy as np
cimport numpy as np
cimport cython
import csv

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

from scipy.spatial.distance import cdist
from scipy.spatial import KDTree

from cryoem.model.distribution cimport RealDistribution
from cryoem.model.transformation cimport (
    Transformation, Transformation3D)
from cryoem.data.grid cimport Grid2D
from cryoem.data.griddata cimport GridDataReal, GridDatasNested
from cryoem.data.griddata cimport GridDataInt
from cryoem.model.math cimport exp_fast, exp_fast2
from cryoem.model.rotation cimport Rotation3D

from cryoem.model.math import log_sum_exp

cdef inline int int_max(int a, int b): return a if a >= b else b
cdef inline int int_min(int a, int b): return a if a <= b else b


cdef class IsoMixture(RealDistribution):
    """Isotropic Gaussian mixture model, constant covariance.

    Dimension should be at least 2.
    """

    def __cinit__(self, means, variance, weights):
        """Initialise mixture.

        Args:
            means: (K, D) array with means of all components
            variance: Single double value.
            weights: (K, ) array, summing to 1.0.
        """

        self._means = means
        self._variance = variance
        self._weights = weights
        self._num_components = len(means)
        self._use_fast_exp = 1

    @property
    def means(self):
        return np.array(self._means)

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

    @variance.setter
    def variance(self, value):
        self._variance = value

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

    @property
    def std(self):
        return np.sqrt(self._variance)

    def set_std(self, value):
        self._variance = value*value

    @property
    def weights(self):
        return np.array(self._weights)

    def set_weights_uniform(self):
        self.set_all_weights(1.0 / self._num_components)

    cpdef set_all_weights(self, float value) :
        """Set all weights to the same value."""
        cdef int k
        for k in range(self._num_components):
            self._weights[k] = value

    def set_means(self, means):
        self._means = means
        self._num_components = len(means)        

    cpdef set_use_fast_exp(self, value):
        if value:
            self._use_fast_exp = 1
        else:
            self._use_fast_exp = 0

    @property
    def use_fast_exp(self):
        return self._use_fast_exp

    @property
    def dimension(self):
        return self._means.shape[1]

    @property
    def num_components(self):
        return self._means.shape[0]

    @property
    def mixture(self):
        return self

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

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

        # Repeat means to obtain an (N, D) array.
        num_component_points = np.random.multinomial(num_samples,
                                                     self.weights)
        means = np.vstack([np.tile(self.means[i],
                                   (num_component_points[i], 1))
                           for i in range(self.num_components)])

        # Sample covariances.
        m = np.zeros(self.dimension)
        cov = self.variance * np.eye(self.dimension)
        shifts = np.random.multivariate_normal(m, cov, num_samples)

        samples = means + shifts

        return samples

    def log_likelihood(self, x):
        
        return self.log_likelihoods(x).sum()

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

        Dimension should be at least 2.

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

        Returns:
            (N, ) double array with log of value of mixture
            at every point.
        """

        # Compute squared distance from each mean to each data point.
        if len(x.shape)==1: # N = 1
            dsq = cdist(self.means, np.reshape(x, [1, len(x)]),
                        'sqeuclidean') # (K, N)
        else:
            dsq = cdist(self.means, x, 'sqeuclidean') # (K, N)

        # Convert these to log-probabilities.
        logp = -0.5 * dsq / self.variance
        logp -= (0.5 * np.log(2.0 * np.pi * self.variance) *
                 self.dimension)

        # Add effect of mixture weights.
        logp += np.log(self.weights)[:, np.newaxis]

        # Sum over components.
        values = log_sum_exp(logp, axis=0)

        return values

    def copy(self):
        """Returns a deep copy of this mixture."""

        means2 = np.array(self._means).copy()
        weights2 = np.array(self._weights).copy()

        return self.__class__(means2, self._variance, weights2)

    def rotate(self, rotation):
        """Returns a transformed copy by applying a rotation.

        Args:
            rotation: Rotation3D object.

        The translation is assumed to be zero.
        """

        means_rotated = rotation(self.means)

        return self.__class__(means_rotated, self.variance,
                              self.weights)

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

        Args:
            transformation: Transformation to be applied to this mixture.
            target: IsoMixture object of right dimensions to contain result.

        Returns:
            target, with new values. Note that the weights are not
            copied.
        """

        means_transformed = transformation(self.means)

        return self.__class__(means_transformed, self.variance,
                              self.weights)

    def inverse_transform(self, transformation):
        """Returns an inverse transformed copy.

        Args:
            transformation: Transformation of which the inverse is
                to be applied to this mixture.
            target: IsoMixture object of right dimensions to contain result.

        Returns:
            target, with new values. Note that the weights are not
            copied.
        """

        means_transformed = transformation.inverse_transform(self.means)

        return self.__class__(means_transformed, self.variance,
                                  self.weights)

    def ellipses(self):
        """
        Return list of parameters representing contours of model.
        """
        scaling = 3.2
        angle = 0.0
        width = np.sqrt(self.variance) * scaling
        height = np.sqrt(self.variance) * scaling

        parameters = [{'mean': mean,
                      'width': width,
                      'height': height,
                      'angle': angle} for mean in self.means]
        return parameters


    def expand(self, symmetries):
        """
        Returns a new IsoGaussianMixture consisting of several transformed
        copies of the original.

        @param symmetries: symmetry group
        @type symmetries: list of 3x3 matrices
        """

        means_list = []
        for s in symmetries:
            means_transformed = np.dot(self.means, s.T) # (K, D)
            means_list.append(means_transformed)
        means_expanded = np.vstack(means_list) # (M*K, D)

        M = len(symmetries)

        weights_expanded = np.tile(self.weights, (1, M)).flatten() / M

        return IsoMixture(means_expanded, self.variance, weights_expanded)

    def submixture(self, indices):
        """Return mixture containing components with the given indices."""

        means = np.array(self._means)[indices, :]
        weights = np.array(self._weights)[indices]
        weights /= weights.sum()

        return self.__class__(means, self._variance, weights)

    def center_of_mass(self):
        """Return center of mass of mixture."""

        m = np.array(self._means)
        w = np.array(self._weights)

        return np.sum(m * w[:, np.newaxis], axis=0)

    def save(self, parent_group, name='mixture'):
        """Save to group in HDF5 database.

        Assume that the group parent_group[name] doesn't exist.

        Args:
            parent_group: Group within which to create the mixture group.
            name: Name of the new mixture group.
        """

        group = parent_group.create_group(name)
        dataset_means = group.create_dataset('means',
                                             (<object>self._means).shape,
                                             np.double)
        dataset_means[...] = self._means
        dataset_weights = group.create_dataset('weights',
                                             (<object>self._weights).shape,
                                             np.double)
        dataset_weights[...] = self._weights
        group.attrs['variance'] = self._variance
        # group.attrs['use_fast_exp'] = self._use_fast_exp


    @classmethod
    def load(cls, parent_group, name='mixture'):
        """Loads IsoMixture from group in HDF5 database."""

        group = parent_group[name]
        means = group['means'][:]
        weights = group['weights'][:]
        variance = group.attrs['variance']
        mixture = cls(means, variance, weights)

        return mixture


cdef class IsoMixture1D:
    """1D isotropic mixture model.

    Don't want to extend 'IsoMixture', because it stores mixture means
    as 2D array.
    """

    def __cinit__(self, means, variance, weights):
        """
        Args:
            means: (K, ) array with means of components.
            variance: Single double value.
            weights: (K, ) array with weights. Sums to 1.0.
        """

        self._means = means
        self._variance = variance
        self._weights = weights
        self._num_components = len(weights)

    @property
    def means(self):
        return self._means

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

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

    @property
    def std(self):
        return np.sqrt(self._variance)

    @property
    def precision(self):
        return 1.0 / self._variance

    @property
    def weights(self):
        return self._weights


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

        Args:
            grid: Grid1D object.

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

        Note: This could be made faster by exploiting sparseness,
        as has been done in the 2D and 3D cases.
        Could also be made faster by doing the evaluation in-place.
        """

        cdef:
            int N, i, k
            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)

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

        return GridDataReal(grid, values)


    cpdef GridDataInt sample_grid(self, double scale, Grid1D grid):
        """Sample Poisson for each grid point from scaled mixture.

        This mixture is multiplied by the scaling factor and evaluated
        at each grid point. These values are the rates for Poisson
        distributions from which counts are sampled.

        Args:
            scale: Positive real value by which to multiply mixture.
            grid: Grid on which to evaluate mixture.

        Return:
            GridDataInt with sampled counts for each grid point.
        """

        cdef:
            GridDataReal evaluated

        evaluated = self.evaluate_grid(grid)
        values = evaluated._values

        counts = np.zeros(grid.shape, dtype='int')

        for i in range(len(counts)):

            counts[i] = np.random.poisson(values[i] * scale)

        return GridDataInt(grid, counts)

    cpdef copy(self):
        """Returns a deep copy of this mixture."""

        return IsoMixture1D(np.array(self._means).copy(),
                            self._variance,
                            np.array(self._weights).copy())


cdef class GeneralIsoMixtureNoise1D:
    """Weighted sum of isotropic components with noise.

    The weights are unnormalised, and could be negative. ('general')
    The noise is assumed to be added after evaluating the mixture.
    ('noise')
    """

    def __cinit__(self, IsoMixture1D mixture, double precision):
        """
        Args:
            mixture: IsoMixture1D objects, with arbitrary weights.
            precision: Precision on Gaussian noise.
        """
        self._mixture = mixture
        self._precision = precision


    @property
    def mixture(self):
        return self._mixture

    @property
    def precision(self):
        return self._precision

    def set_precision(self, value):
        self._precision = value

    cpdef GridDataReal sample_grid(self, Grid1D grid):
        """Sample Gaussian for each grid point from mixture.

        The mixture is evaluated on the grid, and Gaussian noise is added.

        Args:
            grid: Grid on which to evaluate mixture.

        Return:
            GridDataReal with noisy data for each grid point.
        """

        cdef:
            GridDataReal evaluated
            int i, N

        evaluated = self._mixture.evaluate_grid(grid)
        values = evaluated._values

        N = grid.size
        std = np.sqrt(1.0 / self._precision)
        rands = np.random.randn(N)

        for i in range(N):
            values[i] += rands[i] * std

        return evaluated


    cpdef double energy(self, GridDataReal data):
        """Return negative log-likelihood of given data.

        This involves evaluating the mixture on the same grid,
        computing the residuals, and using the noise precision.
        """

        cdef:
            GridDataReal evaluated
            double t12, t3
            int N

        evaluated = self._mixture.evaluate_grid(data._grid)
        diff = np.array(data._values) - np.array(evaluated._values)
        t3 = 0.5 * self._precision * np.sum(diff**2)
        N = data.size
        t12 = 0.5 * N * (log(2 * M_PI) - log(self._precision))

        return t12 + t3

    cpdef energy_gradient_means(self, GridDataReal data,
                                double[:] gradients):
        """Return gradient of energy w.r.t. means.

        Args:
            data: Data on 1D grid, Gaussian noise.
            gradients: Double array, length = num components
                Will contain resulting gradients.

        Return:
            Double array of length num_components. Gradient of
            negative log-likelihood w.r.t. each mean.
            Just the array 'gradients'.

        Note: Not making using of sparsity. Should evaluate
        mixture only near its mean.
        """

        cdef:
            int N, k, i
            double origin, spacing, c1, c2, x, t1, t2, t3, t4, means_k
            double[:] evaluated_values, data_values
            Grid1D grid
            GridDataReal evaluated

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

        evaluated = self._mixture.evaluate_grid(data._grid)
        evaluated_values = evaluated._values
        data_values = data._values

        K = self._mixture._num_components
        ## gradients = np.zeros((K, ), dtype='double')
        for k in range(K):
            gradients[k] = 0.0
            means_k = self._mixture._means[k]
            x = origin
            t1 = (self._precision * self._mixture._weights[k] /
                  self._mixture._variance)
            for i in range(N):
                t2 = (data_values[i] - evaluated_values[i])
                t3 = means_k - x
                t4 = c1 * exp(c2 * t3**2)
                gradients[k] += t1 * t2 * t3 * t4
                x += spacing

        return gradients

    cpdef energy_gradient_weights(self, GridDataReal data,
                                  double[:] gradients):
        """Return gradient of energy w.r.t. weights.

        Args:
            data: Data on 1D grid, Gaussian noise.
            gradients: Storing result.

        Return:
            Double array of length num_components. Gradient of
            negative log-likelihood w.r.t. each weight.

        Note: Not making using of sparsity. Should evaluate
        mixture only near its mean.
        """

        cdef:
            int N, k, i
            double origin, spacing, c1, c2, x, t1, t2, t3, t4
            Grid1D grid
            GridDataReal evaluated

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

        evaluated = self._mixture.evaluate_grid(data._grid)

        K = self._mixture._num_components
        ## gradients = np.zeros((K, ), dtype='double')
        for k in range(K):
            gradients[k] = 0.0
            x = origin
            t1 = self._precision
            for i in range(N):
                t2 = (data._values[i] - evaluated._values[i])
                t3 = self._mixture._means[k] - x
                t4 = c1 * exp(c2 * t3**2)
                gradients[k] -= t1 * t2 * t4
                x += spacing

        return gradients


@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class IsoMixture2D(IsoMixture):
    """2D isotropic mixture model."""

    cpdef evaluate_grid(self, Grid2D grid):
        """Returns evaluation at each grid point.

        Args:
            grid: Grid2D object

        Returns:
            GridDataReal object, containing value at each grid point.
        """

        cdef:
            int i, j, k
            double gridx, gridy, d2
            double c1 = 1.0 / (2.0 * M_PI * self._variance)
            double c2 = -0.5 / self._variance
            double[:] values
            int grid_size = grid._shape_0 * grid._shape_1

        values = np.zeros(grid_size)

        for k in range(self._num_components):
            for i in range(grid._shape_0):
                for j in range(grid._shape_1):
                    gridx = grid._origin_0 + grid._spacing_0 * i
                    gridy = grid._origin_1 + grid._spacing_1 * j
                    d2 = ((self._means[k, 0] - gridx)**2 +
                          (self._means[k, 1] - gridy)**2)
                    values[i * grid._shape_1 + j] += (
                        self._weights[k] * c1 * exp(c2*d2))

        return GridDataReal(grid, values)

    cpdef evaluate_grid_in_place(self, Grid2D grid, double std_range,
                                 GridDataReal evaluated,
                                 int unnormalised):
        """Evaluate at every grid point, store result in 'evaluated'.

        Makes use of grid structure to obtain significant speedup,
        especially if the Gaussian components are small relative
        to the grid.

        Uses 'exp_fast' instead of 'exp', which is much faster,
        but also slightly inaccurate.
        Note: currently changing exp version manually here
        when switching between Gaussian and Poisson noise.

        Args:
            grid: Grid object
            std_range: Multiple of standard deviation to evaluate.
            evaluated: Where result will be stored.
            unnormalised: Default should be 0.
        """

        cdef:
            int i, j, k, a0, a1, b0, b1
            double gridx, gridy, dist2
            double c1 = 1.0 / (2.0 * M_PI * self._variance)
            double c2 = -0.5 / self._variance
            double[:] values
            int grid_size = grid._shape_0 * grid._shape_1
            double local_range
            int local_steps_0, local_steps_1
            int mean_index_0, mean_index_1, index0, index1
            double grid_point_0, grid_point_1
            double means_k_0, means_k_1, ck

        # print 'use_fast_exp =', self._use_fast_exp

        # Use unnormalised Gaussians if desired.
        if unnormalised == 1:
            c1 = 1.0
            ## print 'Using unnormalised Gaussians.'

        # Initialise for evaluation of mixture on grid:
        ## values = np.zeros(grid_size, dtype='float')
        values = evaluated._values
        for j in range(grid_size):
            values[j] = 0.0

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        for k in range(self._num_components):

            means_k_0 = self._means[k, 0]
            means_k_1 = self._means[k, 1]
            ck = self._weights[k] * c1

            # Compute indices of grid point near component mean.
            mean_index_0 = (int)((means_k_0 - grid._origin_0)
                                  / grid._spacing_0)
            mean_index_1 = (int)((means_k_1 - grid._origin_1)
                                  / grid._spacing_1)
            # Compute the indices of the grid points to consider.
            a0 = int_max(mean_index_0 - local_steps_0, 0)
            b0 = int_min(mean_index_0 + local_steps_0 + 2,
                         grid._shape_0)
            for index0 in range(a0, b0):
                a1 = int_max(mean_index_1 - local_steps_1, 0)
                b1 = int_min(mean_index_1 + local_steps_1 + 2,
                             grid._shape_1)
                for index1 in range(a1, b1):
                    j = index0 * grid._shape_1 + index1
                    grid_point_0 = (grid._origin_0 +
                                    grid._spacing_0 * index0)
                    grid_point_1 = (grid._origin_1 +
                                    grid._spacing_1 * index1)
                    dist2 = ( (means_k_0 - grid_point_0)**2
                              + (means_k_1 - grid_point_1)**2 )
                    if self._use_fast_exp == 1:
                        values[j] += ck * exp_fast2(c2 * dist2)
                    else:
                        values[j] += ck * exp(c2 * dist2)

    cpdef GridDataReal evaluate_grid_sparse(self,
                                            Grid2D grid,
                                            double std_range):
        """Returns evaluation at each grid point.

        Makes use of grid structure to obtain significant speedup,
        especially if the Gaussian components are small relative
        to the grid.

        Uses 'exp_fast' instead of 'exp', which is much faster,
        but also slightly inaccurate.

        Args:
            grid: Grid object
            std_range: multiple of standard deviation to evaluate.

        Returns:
            GridDataReal object with the mixture evaluated on the grid
            as values.
        """

        cdef:
            GridDataReal evaluated

        # Initialise for evaluation of mixture on grid:
        values = np.zeros(grid.size, dtype='float')
        evaluated = GridDataReal(grid, values)

        self.evaluate_grid_in_place(grid, std_range, evaluated, 0)

        return evaluated

    def log_likelihood_grid(self, griddata, std_range=5.0):
        """Compute log-likelihood of griddata given this model.

        Args:
            griddata: GridDataReal or GridDataInt object, 2D.

        Returns:
            The log-likelihood, double.

        Uses sparse evaluation of mixture.

        If griddata.values is sparse, it's much faster to use
        the masked version below.
        """

        v = np.array(self.evaluate_grid_sparse(griddata.grid,
                                               std_range).values)

        v[v < 1e-320] = 1e-320 # log(1e-320) != -inf

        return np.sum(np.log(v) * griddata.values)

    cpdef double log_likelihood_grid_sparse(self, GridDataInt griddata,
                                     double std_range):
        """Compute log-likelihood of griddata given this model.

        Args:
            griddata: GridDataReal or GridDataInt object, 2D.

        Returns:
            The log-likelihood, double.

        Uses sparse evaluation of mixture.

        Evaluates mixture only at non-zero grid points.

        Assumes that every grid point is near to some mixture component.
        If a grid point is not covered by the mixture, it will
        contribute a fixed (large) negative value to the log-likelihood.
        This is not numerically correct, but the extra computation
        needed for the correct result is not justified in our case.
        """

        cdef:
            int i, j, k, a0, a1, b0, b1
            double gridx, gridy, dist2
            double c1 = 1.0 / (2.0 * M_PI * self._variance)
            double c2 = -0.5 / self._variance
            double[::1] values
            Grid2D grid = griddata._grid
            long[::1] griddata_values = griddata._values
            int grid_size = grid._shape_0 * grid._shape_1
            double local_range
            int local_steps_0, local_steps_1
            int mean_index_0, mean_index_1, index0, index1
            double grid_point_0, grid_point_1
            double means_k_0, means_k_1, ck, ll, v


        # Initialise for evaluation of mixture on grid:
        values = np.zeros(grid_size, dtype='float')

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        for k in range(self._num_components):

            means_k_0 = self._means[k, 0]
            means_k_1 = self._means[k, 1]
            ck = self._weights[k] * c1

            # Compute indices of grid point near component mean.
            mean_index_0 = (int)((means_k_0 - grid._origin_0)
                                  / grid._spacing_0)
            mean_index_1 = (int)((means_k_1 - grid._origin_1)
                                  / grid._spacing_1)
            # Compute the indices of the grid points to consider.
            a0 = int_max(mean_index_0 - local_steps_0, 0)
            b0 = int_min(mean_index_0 + local_steps_0 + 2, grid._shape_0)
            for index0 in range(a0, b0):
                a1 = int_max(mean_index_1 - local_steps_1, 0)
                b1 = int_min(mean_index_1 + local_steps_1 + 2,grid._shape_1)
                for index1 in range(a1, b1):
                    j = index0 * grid._shape_1 + index1
                    if griddata_values[j] > 0:
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index1)
                        dist2 = ( (means_k_0 - grid_point_0)**2
                                  + (means_k_1 - grid_point_1)**2 )
                        ## values[j] += ck * exp_fast(c2 * dist2)
                        values[j] += ck * exp(c2 * dist2)

        # Calculate the log_likelihood.
        ll = 0.0
        for j in range(grid_size):
            if griddata_values[j] > 0:
                v = values[j]
                if v < 1e-320:
                    v = 1e-320
                ll += griddata_values[j] * log(v)

        return ll

    cpdef log_likelihoods_nested(self, GridDatasNested nested,
                                 double[::1] betas,
                                 double[::1] result):
        """Compute log-likelihood of mixture on every grid in family."""

        cdef:
            double std_range, ll
            int count_index, nested_index, grid_index, i, n
            double[::1] log_values
            long[::1] counts
            GridDataReal evaluated


        std_range = 4.0
        evaluated = self.evaluate_grid_sparse(nested.grid, std_range)

        values = np.array(evaluated._values)
        values[values < 1e-320] = 1e-320
        log_values = np.log(values)

        counts = nested.beta_array_to_counts(betas)

        print 'Counts corresponding to betas:'
        print np.array(counts)

        n = len(betas)

        # Algorithm adds contributions to the log-likelihood in the order
        # that the counts are added. At every point where the number
        # of counts corresponds to one of the specified betas,
        # the running total log-likelihood is assigned to that beta.
        ll = 0.0
        count_index = 0
        nested_index = 0
        while count_index < n:
            while nested_index < counts[count_index]:
                grid_index = nested._indices[nested_index]
                ll += log_values[grid_index]
                nested_index += 1
            result[count_index] = ll
            count_index += 1

    cpdef _transform(self, Transformation2D transformation,
                     IsoMixture2D target):
        """Performs transformation, result in 'target'.

        Args:
            transformation: Transformation to be applied to this mixture.
            target: IsoMixture3D object to contain result.
        """

        transformation.transform(self._means, target._means)
        target._variance = self._variance
        target._weights[:] = self._weights[:]


    def project(self):
        """
        Marginalise out the last coordinate, i.e. project to z=0 hyperplane.

        @note means and weights are not copied, so don't change them
        """
        means_new = self.means[:,:-1] # last coordinate was in last column
        means_new = np.array(self.means[:, 0]).flatten()

        return IsoMixture1D(means_new.copy(), self.variance, self.weights)

    cpdef _project(self, IsoMixture1D target):
        """
        Marginalise out the last coordinate, i.e. project to z=0 hyperplane.
        """

        cdef:
            int K, k

        K = target._num_components
        for k in range(K):
            target._means[k] = self._means[k, 0]
        ## target._means[:] = self._means[:, :-1]
        target._variance = self._variance
        target._weights[:] = self._weights[:]

    def circles(self, std_ratio=2.0, **kwargs):

        from matplotlib.patches import Circle

        radius = std_ratio * np.sqrt(self._variance)
        m = self._means
        if 'fill' not in kwargs:
            kwargs['fill'] = False
        for k in range(self._num_components):
            yield Circle((m[k,0], m[k,1]), radius, **kwargs)

    def add_to_axes(self, ax, **kwargs):

        for circle in self.circles(**kwargs):
            ax.add_artist(circle)


cdef class IsoMixture3D(IsoMixture):
    """3D isotropic mixture model."""


    cpdef _transform(self, Transformation3D transformation,
                     IsoMixture3D target):
        """Performs transformation, result in 'target'.

        Args:
            transformation: Transformation to be applied to this mixture.
            target: IsoMixture3D object to contain result.
        """

        transformation.transform(self._means, target._means)
        target._variance = self._variance
        target._weights[:] = self._weights[:]

    def project(self):
        """
        Marginalise out the last coordinate, i.e. project to z=0 hyperplane.

        @note means and weights are not copied, so don't change them
        """
        means_new = self.means[:,:-1] # last coordinate was in last column

        return IsoMixture2D(means_new.copy(), self.variance, self.weights)

    cpdef _project(self, IsoMixture2D target):
        """
        Marginalise out the last coordinate, i.e. project to z=0 hyperplane.
        """

        target._means[:] = self._means[:, :-1]
        target._variance = self._variance
        target._weights[:] = self._weights[:]


    cpdef evaluate_grid(self, Grid3D grid):
        """Returns evaluation at each grid point.

        Args:
            grid: Grid3D object

        Returns:
            (N, ) array with value of mixture at every point.
        """

        cdef:
            int k, index_0, index_1, index_2
            double gridx, gridy, gridz, dist2
            double c1 = 1.0 / (np.sqrt(2.0 * M_PI * self._variance))**3
            double c2 = -0.5 / self._variance
            double[:] values
            int grid_size = grid._shape_0 * grid._shape_1 * grid._shape_2

        values = np.zeros(grid_size)

        for k in range(self._num_components):
            for index_0 in range(grid._shape_0):
                for index_1 in range(grid._shape_1):
                    for index_2 in range(grid._shape_2):
                        gridx = grid._origin_0 + grid._spacing_0 * index_0
                        gridy = grid._origin_1 + grid._spacing_1 * index_1
                        gridz = grid._origin_2 + grid._spacing_2 * index_2
                        dist2 = ((self._means[k, 0] - gridx)**2 +
                                 (self._means[k, 1] - gridy)**2 +
                                 (self._means[k, 2] - gridz)**2)
                        values[(index_0 * grid._shape_1 + index_1) *
                               grid._shape_2 + index_2] += (
                            self._weights[k] * c1 * exp(c2*dist2))

        return GridDataReal(grid, values)

    cpdef evaluate_grid_sparse(self, Grid3D grid, double std_range):
        """Returns evaluation at each grid point as GridDataReal.

        Makes use of grid structure to obtain significant speedup,
        especially if the Gaussian components are small relative
        to the grid.

        Args:
            grid: Grid object
            std_range: multiple of standard deviation to evaluate.

        Returns:
            GridDataReal object with value of mixture at each gridpoint
        """

        cdef:
            int i, j, k, a0, a1, b0, b1, c0, c1
            double gridx, gridy, gridz, dist2
            double const1 = 1.0 / (np.sqrt(2.0 * M_PI * self._variance))**3
            double const2 = -0.5 / self._variance
            double[:] values
            int grid_size = grid._shape_0 * grid._shape_1 * grid._shape_2
            double local_range
            int local_steps_0, local_steps_1, local_steps_2
            int mean_index_0, mean_index_1, mean_index_2
            int index0, index1, index2
            double grid_point_0, grid_point_1, grid_point_2

        # Initialise for evaluation of mixture on grid:
        values = np.zeros(grid_size, dtype='float')

        # Calculate the number of steps to move in each direction.
        local_range = self.std * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_0 += 1
        local_steps_1 = (int)(local_range / grid._spacing_1)
        local_steps_1 += 1
        local_steps_2 = (int)(local_range / grid._spacing_2)
        local_steps_2 += 1

        for k in range(self._num_components):
            # Compute indices of grid point near component mean.
            mean_index_0 = (int)((self._means[k, 0] - grid._origin_0)
                                  / grid._spacing_0)
            mean_index_1 = (int)((self._means[k, 1] - grid._origin_1)
                                  / grid._spacing_1)
            mean_index_2 = (int)((self._means[k, 2] - grid._origin_2)
                                  / grid._spacing_2)
            # Compute the indices of the grid points to consider.
            a0 = int_max(mean_index_0 - local_steps_0, 0)
            b0 = int_min(mean_index_0 + local_steps_0 + 1, grid._shape_0)
            for index0 in range(a0, b0):
                a1 = int_max(mean_index_1 - local_steps_1, 0)
                b1 = int_min(mean_index_1 + local_steps_1 + 1,grid._shape_1)
                for index1 in range(a1, b1):
                    a2 = int_max(mean_index_2 - local_steps_2, 0)
                    b2 = int_min(mean_index_2 + local_steps_2 + 1,
                                 grid._shape_2)
                    for index2 in range(a2, b2):
                        j = ((index0 * grid._shape_1 + index1) *
                             grid._shape_2 + index2)
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index1)
                        grid_point_2 = (grid._origin_2 +
                                        grid._spacing_2 * index2)
                        dist2 = ( (self._means[k, 0] - grid_point_0)**2 +
                                  (self._means[k, 1] - grid_point_1)**2 +
                                  (self._means[k, 2] - grid_point_2)**2 )
                        values[j] += (self._weights[k] * const1 *
                                      exp(const2 * dist2))

        return GridDataReal(grid, values)

    cpdef evaluate_max_grid(self, Grid3D grid, double std_range):
        """Returns maximum at each grid point.

        Makes use of grid structure to obtain significant speedup,
        especially if the Gaussian components are small relative
        to the grid.

        Args:
            grid: Grid object
            std_range: multiple of standard deviation to evaluate.

        Returns:
            (N, ) array with value of mixture at every point.
        """

        cdef:
            int i, j, k, a0, a1, b0, b1, c0, c1
            double gridx, gridy, gridz, dist2, v
            double const1 = 1.0 / (np.sqrt(2.0 * M_PI * self._variance))**3
            double const2 = -0.5 / self._variance
            double[:] values
            int grid_size = grid._shape_0 * grid._shape_1 * grid._shape_2
            double local_range
            int local_steps_0, local_steps_1, local_steps_2
            int mean_index_0, mean_index_1, mean_index_2
            int index0, index1, index2
            double grid_point_0, grid_point_1, grid_point_2

        # Initialise for evaluation of mixture on grid:
        values = np.zeros(grid_size, dtype='float')

        # Calculate the number of steps to move in each direction.
        local_range = self.std * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_0 += 1
        local_steps_1 = (int)(local_range / grid._spacing_1)
        local_steps_1 += 1
        local_steps_2 = (int)(local_range / grid._spacing_2)
        local_steps_2 += 1

        for k in range(self._num_components):
            # Compute indices of grid point near component mean.
            mean_index_0 = (int)((self._means[k, 0] - grid._origin_0)
                                  / grid._spacing_0)
            mean_index_1 = (int)((self._means[k, 1] - grid._origin_1)
                                  / grid._spacing_1)
            mean_index_2 = (int)((self._means[k, 2] - grid._origin_2)
                                  / grid._spacing_2)
            # Compute the indices of the grid points to consider.
            a0 = int_max(mean_index_0 - local_steps_0, 0)
            b0 = int_min(mean_index_0 + local_steps_0 + 1, grid._shape_0)
            for index0 in range(a0, b0):
                a1 = int_max(mean_index_1 - local_steps_1, 0)
                b1 = int_min(mean_index_1 + local_steps_1 + 1,grid._shape_1)
                for index1 in range(a1, b1):
                    a2 = int_max(mean_index_2 - local_steps_2, 0)
                    b2 = int_min(mean_index_2 + local_steps_2 + 1,
                                 grid._shape_2)
                    for index2 in range(a2, b2):
                        j = ((index0 * grid._shape_1 + index1) *
                             grid._shape_2 + index2)
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index1)
                        grid_point_2 = (grid._origin_2 +
                                        grid._spacing_2 * index2)
                        dist2 = ( (self._means[k, 0] - grid_point_0)**2 +
                                  (self._means[k, 1] - grid_point_1)**2 +
                                  (self._means[k, 2] - grid_point_2)**2 )
                        ## v = self._weights[k] * const1 * exp(const2 * dist2)
                        ## v = const1 * exp(const2 * dist2)
                        v = exp(-dist2 / 100.0)
                        if v > values[j]:
                            values[j] = v

        return GridDataReal(grid, values)

    cpdef log_likelihood_sparse(self, double[:, :] x):
        
        return np.array(self.log_likelihoods_sparse(x)).sum()

    cpdef log_likelihoods_sparse(self, double[:, :] x):
        """Compute log-likelihood of point cloud.

        Args:
            x: (N, 3) array
        """

        cdef:
            int i, j, k, l
            double std_range, c1, c2, dsq, epsilon
            double[:] likelihoods

        std_range = 5.0
        c1 = 1.0 / (sqrt(2.0 * M_PI * self._variance))**3
        c2 = -0.5 / self._variance
        epsilon = 1e-9

        data_tree = KDTree(np.array(x))
        means_tree = KDTree(np.array(self._means))
        rr = std_range * self.std
        neighbour_lists = data_tree.query_ball_tree(means_tree, rr)

        likelihoods = np.zeros(len(x))
        for i in range(len(x)):
            mean_indices = neighbour_lists[i]
            for l in range(len(mean_indices)):
                k = mean_indices[l] # component index
                dsq = 0.0
                for j in range(3):
                    dsq += (x[i, j] - self._means[k, j])**2
                likelihoods[i] += c1 * self._weights[k] * exp(c2 * dsq)
            if likelihoods[i] < epsilon:
                likelihoods[i] = epsilon


        return np.log(likelihoods)
        
    def log_likelihood_grid(self, griddata, std_range=5.0):
        """Compute log-likelihood of griddata given this model.

        Args:
            griddata: GridDataReal or GridDataInt object, 2D.

        Returns:
            The log-likelihood, double.

        Uses sparse evaluation of mixture.

        If griddata.values is sparse, it's much faster to use
        the masked version below.
        """

        v = np.array(self.evaluate_grid_sparse(griddata.grid,
                                               std_range).values)

        v[v < 1e-320] = 1e-320 # log(1e-320) != -inf

        return np.sum(np.log(v) * griddata.values)


    def get_centered_mixture(self):
        """Return IsoMixture3D with center of mass at origin."""

        center_of_mass = self.center_of_mass()
        m = np.array(self._means) - center_of_mass

        return IsoMixture3D(m, self._variance, np.array(self._weights))

    def save_csv(self, filename):
        """Save mixture parameters as csv file."""

        with open(filename, 'wb') as csvfile:
            writer = csv.writer(csvfile, delimiter=',')
            writer.writerow([np.sqrt(self._variance)])
            means = np.array(self._means)
            weights = np.array(self._weights)
            for k in range(self.num_components):
                writer.writerow(np.append(means[k], weights[k]))

    @classmethod
    def load_csv(cls, filename):
        """Load mixture parameters from csv file."""

        means = []
        weights = []

        with open(filename, 'rb') as csvfile:
            reader = csv.reader(csvfile, delimiter=',')
            first = True
            for row in reader:
                if first:
                    variance = float(row[0])**2
                    first = False
                else:
                    x = np.array([float(a) for a in row])
                    means.append(x[:3])
                    weights.append(x[-1])

        means = np.vstack(means)
        weights = np.hstack(weights)

        return cls(means, variance, weights)


    def save_markers(self, filename, radius, colour,
                     model_name='mixture_model'):
        """Save as .cmm file for loading in chimera.

        Only component means are used.

        Args:
            radius: double, in Angstrom.
            colour: [r, g, b], each between 0 and 1.
        """

        from cryoem.data.chimera_utils import MarkerSet

        means = np.array(self._means)

        mset = MarkerSet(model_name)
        for k in range(len(means)):
            mset.add_sphere(means[k], radius, colour)

        mset.write(filename)


    cpdef _expand(self, IsoMixture3D target, list symmetries):
        """Apply symmetries and construct combined mixture.

        First K components are the original ones
        (assuming symmetries[0] is the identity).
        Then follow K components which have been transformed by symmetries[1].
        And so on.

        Note that the target should have M*K components,
        where M is the order of the symmetry group.

        Args:
            target: IsoMixture3D object to contain result.
            symmetries: List of Rotation3D objects. Symmetry group.
        """

        cdef:
            int M, K, index, m, k, i, j, l
            double[:, :] x, y, matrix
            Rotation3D rotation
            
        M = len(symmetries) # Order of symmetry group.
        K = self._num_components

        # Transform means by symmetry rotations.
        x = self._means
        y = target._means
        i = 0
        for m in range(M):
            rotation = symmetries[m]
            matrix = rotation._matrix
            for k in range(K):
                # Transform the k'th mean to produce the i'th mean.
                # km = m*K + k
                # i == m*K + k
                # This is just a matrix multiplication,
                # taken from cryoem.model.rotation.Rotation.rotate(...).
                for j in range(3):
                    y[i, j] = 0.0
                    for l in range(3):
                        y[i, j] += x[k, l] * matrix[j, l]
                i += 1
        
        # Copy variance.
        target._variance = self._variance

        # Copy weights.
        i = 0
        for m in range(M):
            for k in range(K):
                # index = m*K + k
                target._weights[i] = self._weights[k]
                i += 1

        # print 'Expanding:'
        # print np.array(x)
        # print np.array(y)
        # print np.array(self._weights)
        # print np.array(target._weights)
            
    cpdef expand(self, symmetries):
        """Apply symmetries and construct combined mixture.

        Expanded mixture is created, so this is not efficient
        if done multiple times.
        """

        M = len(symmetries)
        K = self._num_components

        means = np.zeros((M*K, 3), dtype='float')
        weights = np.zeros((M*K,), dtype='float')

        mixture = IsoMixture3D(means, self.variance, weights)

        self._expand(mixture, symmetries)

        return mixture

    cpdef _get_means(self, double[:] x):
        """Copy current means to 'x'.

        Note that x should be initialised and have length 3*K.
        """

        cdef:
            int K, k, l

        K = self._num_components
        for k in range(K):
            for l in range(3):
                x[3*k+l] = self._means[k, l]

    cpdef _set_means(self, double[:] x):
        """Set current means."""

        cdef:
            int K, k, l

        K = self._num_components
        for k in range(K):
            for l in range(3):
                self._means[k, l] = x[3*k+l]
        
