"""Sampling multinomial distributions."""

import numpy as np
cimport numpy as np
cimport cython
import time

from cryoem.data.grid cimport Grid2D, Grid3D, Grid1D
from cryoem.data.griddata cimport GridDataInt
from cryoem.model.isomixture cimport IsoMixture
from cryoem.model.isodiagonalmixture cimport IsoDiagonalMixture3D
from cryoem.sampling.quicksort cimport qsort
from cryoem.model.math cimport exp_fast

VERBOSE = False

from libc.math cimport exp, sqrt

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

from cython_gsl cimport *
gsl_rng_env_setup()
cdef:
    const gsl_rng_type * T
    gsl_rng * r
T = gsl_rng_default
r = gsl_rng_alloc (T)


@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class IsoGridMultinomial:
    """Sample hard assignments for isotropic mixtures on 2D grids."""

    def __cinit__(self):
        pass


    cpdef init_memory(self, data, max_std, std_range, num_components):
        """Reserve required memory.

        Run this to reassign memory if the number of components changes,
        or one of the other parameters.

        For the soft assignments, the same memory is reused for each
        projection. A first estimate of the maximum size is the number
        of pixels times the number of components. But this number can
        be significantly reduced if we know the maximum component std
        that will occur during the algorithm. The memory is re-assigned
        whenever the current component std exceeds this maximum.
        So the initial std should not be too high, otherwise memory will
        be wasted.

        For the hard assignments, memory is assigned for each projection.
        Size is the number of projections times the maximum
        number of counts per projection.
        """

        cdef:
            GridDataInt griddata
            Grid2D grid

        # Find the maximum number of counts per grid,
        # and maximum number of grid points.
        # 'max_counts' will be the upper bound for the
        # number of hard assignments per grid.
        max_counts = 0
        max_grid_size = 0
        griddata = data.get(0, 1.0) #data[0]
        grid = griddata._grid
        min_grid_spacing_0 = grid._spacing_0
        min_grid_spacing_1 = grid._spacing_1
        ## for griddata in data:
        for i in range(data.get_num_projections()):
            griddata = data.get(i, 1.0)
            if griddata._values_sum > max_counts:
                max_counts = griddata._values_sum
            grid = griddata._grid
            if grid._size > max_grid_size:
                max_grid_size = grid._size
            if grid._spacing_0 < min_grid_spacing_0:
                min_grid_spacing_0 = grid._spacing_0
            if grid._spacing_1 < min_grid_spacing_1:
                min_grid_spacing_1 = grid._spacing_1
            if min_grid_spacing_0 == 0:
                print 'grid spacing dropped to 0 at i = %d' % i

        # Determine maximum number of soft assignments per grid.
        ## print 'min grid spacing: ', min_grid_spacing_0
        max_local_range = max_std * std_range
        max_local_steps_0 = (int)(max_local_range / min_grid_spacing_0)
        max_local_steps_1 = (int)(max_local_range / min_grid_spacing_1)

        # Prepare data structures for storing non-zero entries
        # of the original (K, N) assignment matrix.
        max_soft_size = (2 * (max_local_steps_0 + 1) *
                         2 * (max_local_steps_1 + 1) * num_components)
        # This should should not be larger than the number of grid points:
        if max_soft_size > (max_grid_size * num_components):
            max_soft_size = max_grid_size * num_components

        self._soft_components = np.empty(max_soft_size, dtype='int')
        self._soft_grid_indices = np.empty(max_soft_size, dtype='int')
        self._soft_values = np.empty(max_soft_size, dtype='float')

        # Initialise for evaluation of mixture on grid:
        self._evaluated = np.zeros(max_grid_size, dtype='float')

        self._cumsums = np.zeros(grid._size, dtype='float')
        # For each grid point, the sum of the soft assignments
        # seen so far.
        self._cum_index = np.zeros(grid._size, dtype='int')

        self._random_values = np.zeros(max_counts, dtype='float')

        # Initialise arrays for hard assignments.
        P = data.get_num_projections() #len(data)
        N = max_counts
        self._hard_components = np.zeros((P, N), dtype='int')
        self._hard_grid_indices = np.zeros((P, N), dtype='int')
        self._hard_values = np.zeros((P, N), dtype='int')



    cpdef int sample(self,
                     GridDataInt griddata,
                     IsoMixture mixture,
                     float std_range,
                     int projection_index):
        """Samples from a multinomial for each grid point.

        Args:
            mixture: IsoMixture object, dimension 2D.
            std_range: Truncate Gaussians at 'std_range' * std.
                If this is at least 5, the loss in accuracy should
                be minimal.

        Returns:
            tuple (components, grid_indices, values),
            each a list of integers.
                components: list of component indices
                grid_indices: list of grid point indices
                values: sampled count for the corresponding
                    (component, grid index) pair
        """

        cdef:
            Grid2D grid = griddata._grid
            long[::1] values = griddata._values
            int values_sum = griddata._values_sum
            int K = mixture._num_components
            double local_range
            int local_steps_0, local_steps_1
            int mean_index_0, mean_index_1
            double PI = 3.141592653589793238462643383279
            double c1 = 1.0 / (2.0 * PI * mixture._variance)
            double c2 = -0.5 / mixture._variance
            double ck
            int count, i, j, k, grid_index, soft_index, soft_size,
            int hard_index, count_temp, hard_size
            int a0, a1, b0, b1, index0, index1
            double grid_point_0, grid_point_1, dist2, value
            double[::1] cumsums = self._cumsums
            double[::1] evaluated = self._evaluated
            long[::1] cum_index = self._cum_index
            long[::1] cum_counts = griddata._cumsum
            long[::1] soft_components = self._soft_components
            long[::1] soft_grid_indices = self._soft_grid_indices
            double[::1] soft_values = self._soft_values
            double[:,::1] means = mixture._means
            double[::1] weights = mixture._weights
            double[::1] random_values = self._random_values
            double means_k_0, means_k_1
            long[:,::1] hard_components = self._hard_components
            long[:,::1] hard_grid_indices = self._hard_grid_indices
            long[:,::1] hard_values = self._hard_values



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

        # Prepare data structures for storing non-zero entries
        # of the original (K, N) assignment matrix.
        # TODO: assign this memory just once
        ## soft_size = (2*local_steps_0 + 1) * (2*local_steps_1 + 1) * K
        soft_size = 2*(local_steps_0 + 1) * 2*(local_steps_1 + 1) * K
            # maximum number of sparse soft assignments
        if soft_size > (grid._size * K):
            soft_size = grid._size * K

        ## if VERBOSE:
        ##     print 'Soft assignment allocated sparsity = %1.3f' % (
        ##         (float)(soft_size) / (float)(grid._size * K))

        soft_index = 0

        # Initialise for evaluation of mixture on grid:
        for i in range(evaluated.shape[0]):
            evaluated[i] = 0.0

        ## time_start = time.time()

        for k in range(K):

            means_k_0 = means[k, 0]
            means_k_1 = means[k, 1]
            ck = 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)
                grid_point_0 = grid._origin_0 + grid._spacing_0 * index0
                for index1 in range(a1, b1):
                    j = index0 * grid._shape_1 + index1
                    if values[j] > 0:
                        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)
                        value = ck * exp_fast(c2*dist2)
                        # Evaluation of k'th component at j'th grid point.
                        evaluated[j] += value
                        soft_components[soft_index] = k
                        soft_grid_indices[soft_index] = j
                        soft_values[soft_index] = value
                        soft_index += 1

        ## print 'computing soft assignments in %1.4f' % (time.time() -
        ##                                                time_start)

        soft_size = soft_index
        ## if VERBOSE:
        ##     print 'Soft assignment true sparsity = %1.3f' % (
        ##         np.float(soft_size) / np.float(grid._size * K))

        # Now have sparse version of soft assignment matrix.

        # Prepare for sampling multinomials to get hard assignments.
        # Note that only if the maximum entry in the hard assignment
        # matrix is 1 do we have a total of 'values_sum' hard
        # assignments. If not (usual case) our arrays are too long.
        ## hard_index = 0 # remove this line? (it's repeated later)

        # Sample required random numbers, and sort.
        for i in range(random_values.shape[0]):
            random_values[i] = gsl_ran_flat(r, 0, 1)

        # Sort segments.
        for i in range(grid._size):
            count = values[i]
            if count >= 2:
                if i==0:
                    qsort(random_values, 0, cum_counts[i])
                    ## random_values[0:cum_counts[i]].sort()
                else:
                    qsort(random_values, cum_counts[i-1], cum_counts[i])
                    ## random_values[cum_counts[i-1]:cum_counts[i]].sort()

        # Sample from multinomials.
        soft_index = 0 # Runs through soft assignments.
        random_index = 0 # Runs through random values.
        grid_index = 0 # Runs through grid points.
        for i in range(cumsums.shape[0]):
            cumsums[i] = 0.0
        ## cumsums = np.zeros(grid._size, dtype='float')
            # For each grid point, the sum of the soft assignments
            # seen so far.
        cum_index[0] = 0
        for i in range(1, cum_index.shape[0]):
            cum_index[i] = cum_counts[i-1]
        ## cum_index = np.zeros(grid._size, dtype='int')
            # For each grid point, the index of a point in
            # 'random_values' within the range of random values
            # for that grid point, and such that the sum of
            # random values for that grid point to its left
            # is just less than or equal to the sum of soft values
            # seen so far for that grid point.

        hard_index = 0
        for soft_index in range(soft_size):
            grid_index = soft_grid_indices[soft_index]
            # Check if we need to generate samples for this grid point.
            if values[grid_index] == 0:
                continue
            cumsums[grid_index] += soft_values[soft_index]
            # Move cum_index[grid_index] along to count the number
            # of counts to assign here.
            count_temp = 0
            while (cum_index[grid_index] < cum_counts[grid_index]
                   and (random_values[cum_index[grid_index]]
                    * evaluated[grid_index] < cumsums[grid_index])):
                count_temp += 1
                cum_index[grid_index] += 1
            if count_temp > 0:
                hard_components[projection_index,
                                hard_index] = soft_components[
                    soft_index]
                hard_grid_indices[projection_index, hard_index] = grid_index
                hard_values[projection_index, hard_index] = count_temp
                hard_index += 1

        hard_size = hard_index
        ## if VERBOSE:
        ##     print 'Hard assignment sparsity = %1.6f' % (
        ##         (float)(hard_size) / (float)(grid._size * K))

        # Now have sparse version of hard assignment matrix.

        return hard_size





#@cython.boundscheck(False)
#@cython.cdivision(True)
#@cython.wraparound(False)
cdef class IsoGrid3DMultinomial:
    """Sample hard assignments for isotropic mixtures on 3D grids."""

    def __cinit__(self):
        pass


    cpdef init_memory(self, GridDataInt data,
                      max_std, std_range, num_components):
        """Reserve required memory.

        Run this to reassign memory if the number of components changes,
        or one of the other parameters.

        Regarding the maximum number of soft assignments:
        A first estimate of the maximum size is the number
        of pixels times the number of components. But this number can
        be significantly reduced if we know the maximum component std
        that will occur during the algorithm. The memory is re-assigned
        whenever the current component std exceeds this maximum.
        So the initial std should not be too high, otherwise memory will
        be wasted.

        The maximum number of hard assignments is equal to the
        total number of counts.
        """

        cdef:
            Grid3D grid

        # Determine maximum number of soft assignments per grid.
        grid = data._grid
        max_local_range = max_std * std_range
        max_local_steps_0 = (int)(max_local_range / grid._spacing_0)
        max_local_steps_1 = (int)(max_local_range / grid._spacing_1)
        max_local_steps_2 = (int)(max_local_range / grid._spacing_2)

        # Prepare data structures for storing non-zero entries
        # of the original (K, N) assignment matrix.
        max_soft_size = (2 * (max_local_steps_0 + 1) *
                         2 * (max_local_steps_1 + 1) *
                         2 * (max_local_steps_2 + 2) * num_components)
        # This should should not be larger than the number of grid points:
        if max_soft_size > (grid._size * num_components):
            max_soft_size = grid._size * num_components

        print('Assigning memory for a maximum of ' +
              '%d soft assignments.' % max_soft_size)

        self._soft_components = np.empty(max_soft_size, dtype='int')
        self._soft_grid_indices = np.empty(max_soft_size, dtype='int')
        self._soft_values = np.empty(max_soft_size, dtype='float')

        # Initialise for evaluation of mixture on grid:
        self._evaluated = np.zeros(grid._size, dtype='float')

        self._cumsums = np.zeros(grid._size, dtype='float')
        # For each grid point, the sum of the soft assignments
        # seen so far.
        self._cum_index = np.zeros(grid._size, dtype='int')

        # Initialise arrays for hard assignments.
        N = data._values_sum
        self._hard_components = np.zeros(N, dtype='int')
        self._hard_grid_indices = np.zeros(N, dtype='int')
        self._hard_values = np.zeros(N, dtype='int')

        self._random_values = np.zeros(N, dtype='float')

    cpdef int sample(self,
                     GridDataInt griddata,
                     IsoMixture3D mixture,
                     float std_range):
        """Samples from a multinomial for each grid point.

        Args:
            griddata: Counts on 3D grid.
            mixture: Current mixture model. Paramemters will be changed.
            std_range: Truncate Gaussians at 'std_range' * std.
                If this is at least 5, the loss in accuracy should
                be minimal.
        """

        cdef:
            Grid3D grid = griddata._grid
            long[::1] values = griddata._values
            int values_sum = griddata._values_sum
            int K = mixture._num_components
            double local_range
            int local_steps_0, local_steps_1, local_steps_2
            int mean_index_0, mean_index_1, mean_index_2
            double PI = 3.141592653589793238462643383279
            double c1 = 1.0 / ((2.0 * PI * mixture._variance) ** 1.5)
            double c2 = -0.5 / mixture._variance
            double ck
            int count, i, j, k, grid_index, soft_index, soft_size, j_level
            int hard_index, count_temp, hard_size
            int a0, a1, b0, b1, a2, b2, index0, index1, index2
            double grid_point_0, grid_point_1, grid_point_2, dist2, value
            double[::1] cumsums = self._cumsums
            double[::1] evaluated = self._evaluated
            long[::1] cum_index = self._cum_index
            long[::1] cum_counts = griddata._cumsum
            long[::1] soft_components = self._soft_components
            long[::1] soft_grid_indices = self._soft_grid_indices
            double[::1] soft_values = self._soft_values
            double[:,::1] means = mixture._means
            double[::1] weights = mixture._weights
            double[::1] random_values = self._random_values
            double means_k_0, means_k_1, means_k_2
            long[::1] hard_components = self._hard_components
            long[::1] hard_grid_indices = self._hard_grid_indices
            long[::1] hard_values = self._hard_values


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

        soft_index = 0

        # Initialise for evaluation of mixture on grid:
        for i in range(evaluated.shape[0]):
            evaluated[i] = 0.0

        for k in range(K):

            means_k_0 = means[k, 0]
            means_k_1 = means[k, 1]
            means_k_2 = means[k, 2]
            ck = 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)
            mean_index_2 = (int)((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 + 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)
                grid_point_0 = grid._origin_0 + grid._spacing_0 * index0
                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 + 2,
                                 grid._shape_2)
                    grid_point_1 = grid._origin_1 + grid._spacing_1 * index1
                    j_level = index0 * grid._shape_1 + index1
                    for index2 in range(a2, b2):
                        j = j_level * grid._shape_2 + index2
                        if values[j] > 0:
                            grid_point_2 = (grid._origin_2 +
                                            grid._spacing_2 * index2)
                            dist2 = ( (means_k_0 - grid_point_0)**2 +
                                      (means_k_1 - grid_point_1)**2 +
                                      (means_k_2 - grid_point_2)**2)
                            ## value = ck * exp_fast(c2*dist2)
                            value = ck * exp(c2*dist2)
                            # Evaluation of k'th component at j'th grid point.
                            evaluated[j] += value
                            soft_components[soft_index] = k
                            soft_grid_indices[soft_index] = j
                            soft_values[soft_index] = value
                            soft_index += 1

        soft_size = soft_index
        # Now have sparse version of soft assignment matrix.

        # Prepare for sampling multinomials to get hard assignments.
        # Note that only if the maximum entry in the hard assignment
        # matrix is 1 do we have a total of 'values_sum' hard
        # assignments. If not (usual case) our arrays are too long.

        ## hard_index = 0 # remove this line? (it's repeated later)

        # Sample required random numbers, and sort.
        for i in range(random_values.shape[0]):
            random_values[i] = gsl_ran_flat(r, 0, 1)

        # Sort segments.
        for i in range(grid._size):
            count = values[i]
            if count >= 2:
                if i==0:
                    qsort(random_values, 0, cum_counts[i])
                else:
                    qsort(random_values, cum_counts[i-1], cum_counts[i])

        # Sample from multinomials.
        soft_index = 0 # Runs through soft assignments.
        random_index = 0 # Runs through random values.
        grid_index = 0 # Runs through grid points.
        for i in range(cumsums.shape[0]):
            cumsums[i] = 0.0
            # For each grid point, the sum of the soft assignments
            # seen so far.
        cum_index[0] = 0
        for i in range(1, cum_index.shape[0]):
            cum_index[i] = cum_counts[i-1]
            # For each grid point, the index of a point in
            # 'random_values' within the range of random values
            # for that grid point, and such that the sum of
            # random values for that grid point to its left
            # is just less than or equal to the sum of soft values
            # seen so far for that grid point.

        hard_index = 0
        for soft_index in range(soft_size):
            grid_index = soft_grid_indices[soft_index]
            # Check if we need to generate samples for this grid point.
            if values[grid_index] == 0:
                continue
            cumsums[grid_index] += soft_values[soft_index]
            # Move cum_index[grid_index] along to count the number
            # of counts to assign here.
            count_temp = 0
            while (cum_index[grid_index] < cum_counts[grid_index]
                   and (random_values[cum_index[grid_index]]
                    * evaluated[grid_index] < cumsums[grid_index])):
                count_temp += 1
                cum_index[grid_index] += 1
            if count_temp > 0:
                hard_components[hard_index] = soft_components[soft_index]
                hard_grid_indices[hard_index] = grid_index
                hard_values[hard_index] = count_temp
                hard_index += 1

        hard_size = hard_index

        # Now have sparse version of hard assignment matrix.

        return hard_size


    cpdef int sample_diagonal(self,
                              GridDataInt griddata,
                              IsoDiagonalMixture3D mixture,
                              float std_range):
        """Samples from a multinomial for each grid point.

        The mixture components can have different sizes.

        Args:
            griddata: Counts on 3D grid.
            mixture: Current mixture model. Paramemters will be changed.
            std_range: Truncate Gaussians at 'std_range' * std.
                If this is at least 5, the loss in accuracy should
                be minimal.
        """

        cdef:
            Grid3D grid = griddata._grid
            long[::1] values = griddata._values
            int values_sum = griddata._values_sum
            int K = mixture._num_components
            double local_range
            int local_steps_0, local_steps_1, local_steps_2
            int mean_index_0, mean_index_1, mean_index_2
            double PI = 3.141592653589793238462643383279
            double c1, c2, ck
            int count, i, j, k, grid_index, soft_index, soft_size, j_level
            int hard_index, count_temp, hard_size
            int a0, a1, b0, b1, a2, b2, index0, index1, index2
            double grid_point_0, grid_point_1, grid_point_2, dist2, value
            double[::1] cumsums = self._cumsums
            double[::1] evaluated = self._evaluated
            long[::1] cum_index = self._cum_index
            long[::1] cum_counts = griddata._cumsum
            long[::1] soft_components = self._soft_components
            long[::1] soft_grid_indices = self._soft_grid_indices
            double[::1] soft_values = self._soft_values
            double[:,::1] means = mixture._means
            double[::1] weights = mixture._weights
            double[::1] random_values = self._random_values
            double means_k_0, means_k_1, means_k_2
            long[::1] hard_components = self._hard_components
            long[::1] hard_grid_indices = self._hard_grid_indices
            long[::1] hard_values = self._hard_values



        soft_index = 0

        # Initialise for evaluation of mixture on grid:
        for i in range(evaluated.shape[0]):
            evaluated[i] = 0.0

        for k in range(K):

            # Calculate the number of steps to move in each direction.
            local_range = sqrt(mixture._variances[k]) * std_range
            local_steps_0 = (int)(local_range / grid._spacing_0)
            local_steps_1 = (int)(local_range / grid._spacing_1)
            local_steps_2 = (int)(local_range / grid._spacing_2)

            means_k_0 = means[k, 0]
            means_k_1 = means[k, 1]
            means_k_2 = means[k, 2]

            c1 = 1.0 / ((2.0 * PI * mixture._variances[k]) ** 1.5)
            c2 = -0.5 / mixture._variances[k]
            ck = 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)
            mean_index_2 = (int)((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 + 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)
                grid_point_0 = grid._origin_0 + grid._spacing_0 * index0
                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 + 2,
                                 grid._shape_2)
                    grid_point_1 = grid._origin_1 + grid._spacing_1 * index1
                    j_level = index0 * grid._shape_1 + index1
                    for index2 in range(a2, b2):
                        j = j_level * grid._shape_2 + index2
                        if values[j] > 0:
                            grid_point_2 = (grid._origin_2 +
                                            grid._spacing_2 * index2)
                            dist2 = ( (means_k_0 - grid_point_0)**2 +
                                      (means_k_1 - grid_point_1)**2 +
                                      (means_k_2 - grid_point_2)**2)
                            ## value = ck * exp_fast(c2*dist2)
                            value = ck * exp(c2*dist2)
                            # Evaluation of k'th component at j'th grid point.
                            evaluated[j] += value
                            soft_components[soft_index] = k
                            soft_grid_indices[soft_index] = j
                            soft_values[soft_index] = value
                            soft_index += 1

        soft_size = soft_index
        # Now have sparse version of soft assignment matrix.

        # Prepare for sampling multinomials to get hard assignments.
        # Note that only if the maximum entry in the hard assignment
        # matrix is 1 do we have a total of 'values_sum' hard
        # assignments. If not (usual case) our arrays are too long.

        ## hard_index = 0 # remove this line? (it's repeated later)

        # Sample required random numbers, and sort.
        for i in range(random_values.shape[0]):
            random_values[i] = gsl_ran_flat(r, 0, 1)

        # Sort segments.
        for i in range(grid._size):
            count = values[i]
            if count >= 2:
                if i==0:
                    qsort(random_values, 0, cum_counts[i])
                else:
                    qsort(random_values, cum_counts[i-1], cum_counts[i])

        # Sample from multinomials.
        soft_index = 0 # Runs through soft assignments.
        random_index = 0 # Runs through random values.
        grid_index = 0 # Runs through grid points.
        for i in range(cumsums.shape[0]):
            cumsums[i] = 0.0
            # For each grid point, the sum of the soft assignments
            # seen so far.
        cum_index[0] = 0
        for i in range(1, cum_index.shape[0]):
            cum_index[i] = cum_counts[i-1]
            # For each grid point, the index of a point in
            # 'random_values' within the range of random values
            # for that grid point, and such that the sum of
            # random values for that grid point to its left
            # is just less than or equal to the sum of soft values
            # seen so far for that grid point.

        hard_index = 0
        for soft_index in range(soft_size):
            grid_index = soft_grid_indices[soft_index]
            # Check if we need to generate samples for this grid point.
            if values[grid_index] == 0:
                continue
            cumsums[grid_index] += soft_values[soft_index]
            # Move cum_index[grid_index] along to count the number
            # of counts to assign here.
            count_temp = 0
            while (cum_index[grid_index] < cum_counts[grid_index]
                   and (random_values[cum_index[grid_index]]
                    * evaluated[grid_index] < cumsums[grid_index])):
                count_temp += 1
                cum_index[grid_index] += 1
            if count_temp > 0:
                hard_components[hard_index] = soft_components[soft_index]
                hard_grid_indices[hard_index] = grid_index
                hard_values[hard_index] = count_temp
                hard_index += 1

        hard_size = hard_index

        # Now have sparse version of hard assignment matrix.

        return hard_size


cdef class IsoTransformGrid2DMultinomial:
    """Sampling multinomial on 1D grid, for 2D reconstruction.

    Samples hard assignments, for reconstruction.
    """

    def __cinit__(self):
        pass


    cpdef init_memory(self, data, max_std, std_range, num_components):
        """Reserve required memory.

        Run this to reassign memory if the number of components changes,
        or one of the other parameters.
        """

        cdef:
            GridDataInt griddata
            Grid1D grid

        # Find the maximum number of counts per grid,
        # and maximum number of grid points.
        # 'max_counts' will be the upper bound for the
        # number of hard assignments per grid.
        max_counts = 0
        max_grid_size = 0
        griddata = data[0]
        grid = griddata._grid
        min_grid_spacing_0 = grid._spacing[0]
        ## for griddata in data:
        for i in range(len(data)):
            griddata = data[i]
            if griddata._values_sum > max_counts:
                max_counts = griddata._values_sum
            grid = griddata._grid
            if grid._size > max_grid_size:
                max_grid_size = grid._size
            if grid._spacing[0] < min_grid_spacing_0:
                min_grid_spacing_0 = grid._spacing[0]
            if min_grid_spacing_0 == 0:
                print 'grid spacing dropped to 0 at i = %d' % i

        # Determine maximum number of soft assignments per grid.
        ## print 'min grid spacing: ', min_grid_spacing_0
        max_local_range = max_std * std_range
        max_local_steps_0 = (int)(max_local_range / min_grid_spacing_0)

        # Prepare data structures for storing non-zero entries
        # of the original (K, N) assignment matrix.
        max_soft_size = (2 * (max_local_steps_0 + 1) * num_components)
        # This should should not be larger than the number of grid points:
        if max_soft_size > (max_grid_size * num_components):
            max_soft_size = max_grid_size * num_components

        self._soft_components = np.empty(max_soft_size, dtype='int')
        self._soft_grid_indices = np.empty(max_soft_size, dtype='int')
        self._soft_values = np.empty(max_soft_size, dtype='float')

        # Initialise for evaluation of mixture on grid:
        self._evaluated = np.zeros(max_grid_size, dtype='float')

        self._cumsums = np.zeros(grid._size, dtype='float')
        # For each grid point, the sum of the soft assignments
        # seen so far.
        self._cum_index = np.zeros(grid._size, dtype='int')

        self._random_values = np.zeros(max_counts, dtype='float')

        # Initialise arrays for hard assignments.
        P = len(data)
        N = max_counts
        self._hard_components = np.zeros((P, N), dtype='int')
        self._hard_grid_indices = np.zeros((P, N), dtype='int')
        self._hard_values = np.zeros((P, N), dtype='int')



    cpdef int sample(self,
                     GridDataInt griddata,
                     IsoMixture1D mixture,
                     float std_range,
                     int projection_index):
        """Samples from a multinomial for each grid point.

        Args:
            mixture: IsoMixture object, dimension 1D.
            std_range: Truncate Gaussians at 'std_range' * std.
                If this is at least 5, the loss in accuracy should
                be minimal.

        Returns:
            tuple (components, grid_indices, values),
            each a list of integers.
                components: list of component indices
                grid_indices: list of grid point indices
                values: sampled count for the corresponding
                    (component, grid index) pair
        """

        cdef:
            Grid1D grid = griddata._grid
            long[::1] values = griddata._values
            int values_sum = griddata._values_sum
            int K = mixture._num_components
            double local_range
            int local_steps_0, local_steps_1
            int mean_index_0, mean_index_1
            double PI = 3.141592653589793238462643383279
            double c1 = 1.0 / sqrt(2.0 * PI * mixture._variance)
            double c2 = -0.5 / mixture._variance
            double ck
            int count, i, j, k, grid_index, soft_index, soft_size,
            int hard_index, count_temp, hard_size
            int a0, a1, b0, b1, index0, index1
            double grid_point_0, grid_point_1, dist2, value
            double[::1] cumsums = self._cumsums
            double[::1] evaluated = self._evaluated
            long[::1] cum_index = self._cum_index
            long[::1] cum_counts = griddata._cumsum
            long[::1] soft_components = self._soft_components
            long[::1] soft_grid_indices = self._soft_grid_indices
            double[::1] soft_values = self._soft_values
            double[:] means = mixture._means
            double[::1] weights = mixture._weights
            double[::1] random_values = self._random_values
            double means_k_0, means_k_1
            long[:,::1] hard_components = self._hard_components
            long[:,::1] hard_grid_indices = self._hard_grid_indices
            long[:,::1] hard_values = self._hard_values



        # Calculate the number of steps to move in each direction.
        local_range = sqrt(mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing[0])

        # Prepare data structures for storing non-zero entries
        # of the original (K, N) assignment matrix.
        # TODO: assign this memory just once
        ## soft_size = (2*local_steps_0 + 1) * (2*local_steps_1 + 1) * K
        soft_size = 2*(local_steps_0 + 1) * K
            # maximum number of sparse soft assignments
        if soft_size > (grid._size * K):
            soft_size = grid._size * K

        ## if VERBOSE:
        ##     print 'Soft assignment allocated sparsity = %1.3f' % (
        ##         (float)(soft_size) / (float)(grid._size * K))

        soft_index = 0

        # Initialise for evaluation of mixture on grid:
        for i in range(evaluated.shape[0]):
            evaluated[i] = 0.0

        ## time_start = time.time()

        for k in range(K):

            means_k_0 = means[k]
            ck = weights[k] * c1

            # Compute indices of grid point near component mean.
            mean_index_0 = (int)((means[k] - grid._origin[0])
                                  / grid._spacing[0])
            # 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):
                j = index0
                if values[j] > 0:
                    grid_point_0 = (grid._origin[0] +
                                    grid._spacing[0] * index0)
                    dist2 = (means_k_0 - grid_point_0)**2
                    value = ck * exp_fast(c2*dist2)
                    # Evaluation of k'th component at j'th grid point.
                    evaluated[j] += value
                    soft_components[soft_index] = k
                    soft_grid_indices[soft_index] = j
                    soft_values[soft_index] = value
                    soft_index += 1

        ## print 'computing soft assignments in %1.4f' % (time.time() -
        ##                                                time_start)

        soft_size = soft_index
        ## if VERBOSE:
        ##     print 'Soft assignment true sparsity = %1.3f' % (
        ##         np.float(soft_size) / np.float(grid._size * K))

        # Now have sparse version of soft assignment matrix.

        # Prepare for sampling multinomials to get hard assignments.
        # Note that only if the maximum entry in the hard assignment
        # matrix is 1 do we have a total of 'values_sum' hard
        # assignments. If not (usual case) our arrays are too long.
        ## hard_index = 0 # remove this line? (it's repeated later)

        # Sample required random numbers, and sort.
        for i in range(random_values.shape[0]):
            random_values[i] = gsl_ran_flat(r, 0, 1)

        # Sort segments.
        for i in range(grid._size):
            count = values[i]
            if count >= 2:
                if i==0:
                    qsort(random_values, 0, cum_counts[i])
                    ## random_values[0:cum_counts[i]].sort()
                else:
                    qsort(random_values, cum_counts[i-1], cum_counts[i])
                    ## random_values[cum_counts[i-1]:cum_counts[i]].sort()

        # Sample from multinomials.
        soft_index = 0 # Runs through soft assignments.
        random_index = 0 # Runs through random values.
        grid_index = 0 # Runs through grid points.
        for i in range(cumsums.shape[0]):
            cumsums[i] = 0.0
        ## cumsums = np.zeros(grid._size, dtype='float')
            # For each grid point, the sum of the soft assignments
            # seen so far.
        cum_index[0] = 0
        for i in range(1, cum_index.shape[0]):
            cum_index[i] = cum_counts[i-1]
        ## cum_index = np.zeros(grid._size, dtype='int')
            # For each grid point, the index of a point in
            # 'random_values' within the range of random values
            # for that grid point, and such that the sum of
            # random values for that grid point to its left
            # is just less than or equal to the sum of soft values
            # seen so far for that grid point.

        hard_index = 0
        for soft_index in range(soft_size):
            grid_index = soft_grid_indices[soft_index]
            # Check if we need to generate samples for this grid point.
            if values[grid_index] == 0:
                continue
            cumsums[grid_index] += soft_values[soft_index]
            # Move cum_index[grid_index] along to count the number
            # of counts to assign here.
            count_temp = 0
            while (cum_index[grid_index] < cum_counts[grid_index]
                   and (random_values[cum_index[grid_index]]
                    * evaluated[grid_index] < cumsums[grid_index])):
                count_temp += 1
                cum_index[grid_index] += 1
            if count_temp > 0:
                hard_components[projection_index,
                                hard_index] = soft_components[
                    soft_index]
                hard_grid_indices[projection_index, hard_index] = grid_index
                hard_values[projection_index, hard_index] = count_temp
                hard_index += 1

        hard_size = hard_index
        ## if VERBOSE:
        ##     print 'Hard assignment sparsity = %1.6f' % (
        ##         (float)(hard_size) / (float)(grid._size * K))

        # Now have sparse version of hard assignment matrix.

        return hard_size
