"""For sampling isotropic mixtures together with transformations.

"""


import numpy as np
cimport numpy as np
import time

from libc.math cimport exp, sqrt, fabs

cimport cython

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)
seed = np.random.randint(2**16)
# seed = 1
gsl_rng_set(r, seed)

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 cryoem.sampling cimport multinomial
from cryoem.model.rotation cimport ExpTrace, Rotation3D
from cryoem.model.rotation cimport exponential_coordinates_to_rotation
from cryoem.model.isotransform cimport IsoTransform
from cryoem.model.transformation cimport Transformation
from cryoem.model.transformation cimport Transformation3D
from cryoem.model.isomixture cimport IsoMixture3D, IsoMixture2D
from cryoem.model.prior cimport Dirichlet, TransformationPrior, Gamma
from cryoem.model.prior cimport FixedRotationPrior, IsoTransformPrior
from cryoem.model.prior cimport IsoTransformNoisePrior
from cryoem.data.grid cimport Grid2D
from cryoem.data.griddata cimport GridDataReal



@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class IsoTransformSampler:
    """Sample mixture and transformations given count data."""

    def __cinit__(self, GridDataList data, IsoTransformPrior prior,
                  double std_range):
        """
        Args:
            data: GridDataList or GridDataNestedList object.
            prior: IsoTransformPrior object.
            std_range: Components will be evaluated at grid points
                within this multiple of their std. Example: 5.0.
        """

        self._data = data
        self._prior = prior
        self._std_range = std_range

        self._rotation_beta = 1.0 # For annealing rotations.
        self._counts_beta = 1.0 # For annealing by reducing counts.

        self._num_components = 0 # Number of mixture components.
        self._num_projections = data.get_num_projections() #len(data)

        # Keep track of number of counts assigned for each direction.
        # This is equal to or slightly less than the number
        # of counts given by griddata for each direction.
        self._true_counts = np.zeros(self._num_projections, dtype='int')

        self._max_std = 0.0

        self._multinomial = IsoGridMultinomial()

        self._exptrace = ExpTrace()

    @property
    def data(self):
        return self._data

    @property
    def prior(self):
        return self._prior

    @property
    def rotation_beta(self):
        return self._rotation_beta

    def set_rotation_beta(self, value):
        self._rotation_beta = value

    @property
    def counts_beta(self):
        return self._counts_beta

    def set_counts_beta(self, value):
        self._counts_beta = value

    @property
    def std_range(self):
        return self._std_range

    cpdef init_memory(self, int K, double std):
        """Initialise large arrays.

        This is done at the beginning, and every time K changes,
        so that we don't re-assign the memory after every step.

        Args:
            K: int, Number of components for new models.
        """

        cdef:
            int P

        P = self._num_projections
        self._num_components = K
        self._sampled_means = np.empty([P, K, 3], dtype='float')

        # Initialise assignment sums. For each projection and each
        # component, this is the number of counts assigned to that
        # component.
        self._assignment_sums_all = np.empty([P, K], dtype='int')

        self._assignment_sums = np.empty(K, dtype='int')

        self._transformed = IsoMixture3D(np.zeros([K, 3], dtype='float'),
                                         1.0, np.zeros(K, dtype='float'))
        self._projected = IsoMixture2D(np.zeros([K, 2], dtype='float'),
                                       1.0, np.zeros(K, dtype='float'))

        self._max_std = 2.5 * std
        self._max_variance = self._max_std ** 2
        self._multinomial.init_memory(self._data, self._max_std,
                                      self._std_range,
                                      self._num_components)

    cpdef init_memory_check(self, IsoTransform model):
        """Initialise memory if needed."""

        if (model._mixture._num_components != self._num_components or
            model._mixture._variance > self._max_variance):
            self.init_memory(model._mixture._num_components,
                             sqrt(model._mixture._variance))


    cpdef sample(self, IsoTransform model):
        """Generates a new sample given the current one.

        Uses Gibbs sampling.

        Args:
            model: IsoTransform object.
                Changed in-place.
        """

        ## times = []
        self.init_memory_check(model)
        ## times.append(time.time())
        self.sample_assignments(model)
        ## times.append(time.time())
        self.sample_missing_components(model)
        ## times.append(time.time())
        self.sample_weights(model)
        ## times.append(time.time())
        self.sample_means(model)
        ## times.append(time.time())
        self.sample_variance(model)
        ## times.append(time.time())
        self.sample_transformations(model)
        ## times.append(time.time())

        ## times = np.array(times)
        ## l = len(times)
        ## diffs = times[1:] - times[:(l-1)]
        ## p = diffs / np.sum(diffs)
        ## print((100*p).astype('i'))


    cpdef sample_mixture_part(self, IsoTransform model):
        """Samples everything but the transformations."""

        self.init_memory_check(model)
        self.sample_assignments(model)
        self.sample_missing_components(model)
        self.sample_weights(model)
        self.sample_means(model)
        self.sample_variance(model)

    cpdef sample_transformations_part(self, IsoTransform model):
        """Samples everything but the mixture."""

        self.init_memory_check(model)
        self.sample_assignments(model)
        self.sample_missing_components(model)
        self.sample_transformations(model)

    cpdef sample_rotations_part(self, IsoTransform model):
        """Samples everything but the mixture and the translations."""

        self.init_memory_check(model)
        self.sample_assignments(model)
        self.sample_missing_components(model)
        self.sample_rotations(model)

    cpdef sample_assignments(self, IsoTransform model):
        """Sample assignments for each image."""

        cdef:
            IsoMixture3D transformed
            int i, k, c
            Transformation3D transformation
            IsoMixture3D mixture


        # Sample hard assignments for each direction.
        for i in range(self._num_projections):
            transformation = model._transformations[i]
            mixture = model._mixture
            mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            ## time_start = time.time()
            self._true_counts[i] = self._multinomial.sample(
                self._data.get(i, self._counts_beta),
                ## self._data[i], # pass self._counts_beta as well
                self._projected,
                self._std_range,
                i)
            ## if i == 0:
            ##     print 'sampling from multinomial in %1.4f' % (time.time() -
            ##                                                   time_start)
            ##     print '(should be around 0.037 seconds)'

            # Debug:
            ## a1 = np.sum(self._data[i].values)
            ## a2 = np.sum(self._multinomial._hard_values[
            ##             i, :self._true_counts[i]])
            ## if a1 != a2:
            ##     print 'Projection %d: %d / %d unassigned counts' % (
            ##         i, a1 - a2, a1)

            # Compute the total number of counts assigned to
            # each component for this direction.
            for k in range(self._num_components):
                self._assignment_sums_all[i, k] = 0
            for c in range(self._true_counts[i]):
                self._assignment_sums_all[i, \
                    self._multinomial._hard_components[i, c]] += \
                        self._multinomial._hard_values[i, c]


        # Compute the total number of counts for each component
        # over all directions.
        for k in range(self._num_components):
            self._assignment_sums[k] = 0
        self._assignment_sums_sum = 0
        for k in range(self._num_components):
            for i in range(self._num_projections):
                self._assignment_sums[k] += self._assignment_sums_all[i, k]
            self._assignment_sums_sum += self._assignment_sums[k]
        # (K, )

    cdef sample_missing_components(self, IsoTransform model):
        """Sample missing z-component for each count.

        In reality we're computing the sufficient statistics
        for allowing the means and rotations to be sampled later.
        It's not necessary to sample every missing z-component.
        """

        cdef:
            double[:,::1] rands
            Transformation transformation
            int i, k, c, v, j
            float px, py, std2
            GridDataInt griddata
            Grid2D grid
            IsoMixture3D mixture

        rands = np.random.randn(self._num_projections,
                                self._num_components)

        for i in range(self._num_projections):
            transformation = model._transformations[i]
            mixture = model._mixture
            mixture._transform(transformation, self._transformed)
            griddata = self._data.get(i, self._counts_beta)
            ## griddata = self._data[i] # self._counts_beta
            grid = griddata._grid

            # Sample the missing components.
            for k in range(self._num_components):
                self._sampled_means[i, k, 2] = (
                    self._transformed._means[k, 2] *
                    self._assignment_sums_all[i, k])
                std2 = (model._mixture._variance *
                        self._assignment_sums_all[i, k])
                self._sampled_means[i, k, 2] += (
                    rands[i, k] * sqrt(std2))

            # Compute the observed, deterministic part.
            # This is just a weighted sum of the grid points.
            for k in range(self._num_components):
                self._sampled_means[i, k, 0] = 0
                self._sampled_means[i, k, 1] = 0
            for c in range(self._true_counts[i]):
                k = self._multinomial._hard_components[i, c]
                v = self._multinomial._hard_values[i, c]
                j = self._multinomial._hard_grid_indices[i, c]
                px = ((j / grid._shape_1) * grid._spacing_0 +
                      grid._origin_0)
                py = ((j % grid._shape_1) * grid._spacing_1 +
                      grid._origin_1)
                self._sampled_means[i, k, 0] += v * px
                self._sampled_means[i, k, 1] += v * py

    cdef sample_weights(self, IsoTransform model):
        """Sample mixture model weights."""

        cdef:
            double[::1] alpha

        alpha = (np.array(self._assignment_sums) +
                 np.array(self._prior._mixture_prior._weights_prior.alpha))
        ## print 'min alpha: ', np.sort(alpha)[:5]
        model._mixture._weights = np.random.dirichlet(alpha)

    cdef sample_means(self, IsoTransform model):
        """Sample means of each component of mixture model."""

        cdef:
            float s, precision_scalar, precision_scalar_initial, mean_x
            double diff
            int k, i, x, y
            double[::1] precision_mean
            double[:,::1] rands
            double[:,:,::1] rds
            double[:] trans
            Transformation t

        s = 1.0 / model._mixture._variance

        precision_mean = np.zeros(3, dtype='float')

        rands = np.random.randn(self._num_components, 3)

        # Transpose and scale rotations, component not needed.
        rds = np.empty((self._num_projections, 3, 3), dtype='float')
        for i in range(self._num_projections):
            t = model._transformations[i]
            for x in range(3):
                for y in range(3):
                    rds[i, x, y] = t._rotation._matrix[y, x] * s


        precision_scalar_initial = (1.0 / self._prior._mixture_prior.
                                    _mean_prior.variance)


        for k in range(self._num_components):

            for x in range(3):
                precision_mean[x] = 0.0

            precision_scalar = precision_scalar_initial

            for i in range(self._num_projections):
                t = model._transformations[i]
                trans = t._translation
                for x in range(3):
                    for y in range(3):
                        diff = (self._sampled_means[i, k, y] -
                                self._assignment_sums_all[i, k] * trans[y])
                        precision_mean[x] += (rds[i, x, y] * diff)
                precision_scalar += self._assignment_sums_all[i, k] * s

            # Now precision_scalar = n_k * s + r.
            ## if k <= 5:
            ##     counts = ((precision_scalar - precision_scalar_initial) / s)
            ##     print('k = %d, std = %0.3f, mean_0 = %0.3f, counts = %1.1f' % (
            ##         k, 1.0 / sqrt(precision_scalar),
            ##         precision_mean[0] / precision_scalar, counts))

            for x in range(3):
                mean_x = precision_mean[x] / precision_scalar
                model._mixture._means[k, x] = (rands[k, x] /
                                               sqrt(precision_scalar) +
                                               mean_x)

    cdef sample_variance(self, IsoTransform model):
        """Sample the variance of the mixture model."""

        cdef:
            int i, c, k, v, j, nik
            Transformation transformation
            IsoMixture3D transformed
            double a, px, py, dsq, b1, b2, b
            GridDataInt griddata
            Grid2D grid
            IsoMixture3D mixture

        # First hyperparameter.
        ## print 'a = %1.1f * 1.5 + %1.1f' % (
        ##     self._assignment_sums_sum,
        ##     self._prior._mixture_prior._precision_prior._alpha)
        a = (self._assignment_sums_sum * 1.5 +
             self._prior._mixture_prior._precision_prior._alpha)

        # Second hyperparameter.
        # Part from observed components.
        b1 = 0.0

        for i, transformation in enumerate(model._transformations):
            mixture = model._mixture
            mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)
            griddata = self._data.get(i, self._counts_beta)
            ## griddata = self._data[i] # self._counts_beta
            grid = griddata._grid
            # Calculate distance between 2D grid points and projected means.
            for c in range(self._true_counts[i]):
                k = self._multinomial._hard_components[i, c]
                v = self._multinomial._hard_values[i, c]
                j = self._multinomial._hard_grid_indices[i, c]
                px = ((j / grid._shape_1) * grid._spacing_0 +
                      grid._origin_0)
                py = ((j % grid._shape_1) * grid._spacing_1 +
                      grid._origin_1)
                dsq = ((self._projected._means[k, 0] - px)**2 +
                       (self._projected._means[k, 1] - py)**2)
                b1 += dsq * v
        # Part from missing components.
        # Sum of sampled variances.
        b2 = 0.0
        for i in range(self._num_projections):
            for k in range(self._num_components):
                nik = self._assignment_sums_all[i, k]
                if nik > 0:
                    ## b2 += np.random.chisquare(nik)
                    b2 += gsl_ran_chisq(r, nik)
        b2 *= model._mixture._variance

        ## print 'b = 0.5 * (%1.1f + %1.1f) + %1.1f' % (
        ##     b1, b2, self._prior._mixture_prior._precision_prior._beta)
        b = (0.5 * (b1 + b2) +
             self._prior._mixture_prior._precision_prior._beta)

        ## model._mixture._variance = 1.0 / np.random.gamma(a, 1.0 / b)
        model._mixture._variance = 1.0 / gsl_ran_gamma(r, a, 1.0 / b)

    cdef sample_transformations(self, IsoTransform model):
        """Sample rotations and translations for each direction."""

        self.sample_rotations(model)
        self.sample_translations(model)

    cdef sample_rotations(self, IsoTransform model):
        """Sample rotations for each direction."""

        cdef:
            Transformation t
            double coefficient, s, diff
            int i, ix, iy, k
            double[:,::1] A
            double[:] trans

        A = np.zeros((3, 3))

        for i in range(self._num_projections):
            if isinstance(self._prior._transformation_priors[i].
                          rotation_prior, FixedRotationPrior):
                continue

            coefficient = self._rotation_beta / model._mixture._variance

            # The for-loops below implement the following matrix
            # multiplication:
            ## temp = np.dot(transpose(self._sampled_means[i]),
            ##               model.mixture.means) # (3, 3)
            ## A = self._rotation_beta * temp / model.mixture.variance

            t = model._transformations[i]
            trans = t._translation

            for ix in range(3):
                for iy in range(3):
                    s = 0
                    for k in range(model._mixture._num_components):
                        diff = (self._sampled_means[i, k, ix] -
                                self._assignment_sums_all[i, k] * trans[ix])
                        s += (diff * model._mixture._means[k, iy])
                        ## s += (self._sampled_means[i, k, ix] *
                        ##       model._mixture._means[k, iy])
                    A[ix, iy] = s * coefficient

            # Sample from exp(tr(A^T B)).
            ## exp_trace = ExpTrace(A)
            ## exp_trace = ExpTrace()
            B = self._exptrace.sample(A)
            ## t = model.transformations[i]
            t._rotation._matrix = B



    cdef sample_translations(self, IsoTransform model):
        """Sample translations for each direction."""

        cdef:
            Transformation t
            double[::1] diff_sum
            double[:,::1] rands
            double[:,:] r
            int i, k, x, y
            double rtd_x, n_i, precision_scalar


        rands = np.random.randn(self._num_projections, 3)
        diff_sum = np.zeros(3, dtype='float')

        for i in range(self._num_projections):


            t = model._transformations[i]
            r = t._rotation._matrix

            n_i = 0.0

            for x in range(3):
                diff_sum[x] = 0.0

            for k in range(self._num_components):

                for x in range(3): # row
                    rtd_x = 0.0
                    for y in range(3): # column
                        rtd_x += r[x, y] * model._mixture._means[k, y]
                            # R_i * \mu_k, row x

                    diff_sum[x] += (
                        self._sampled_means[i, k, x] -
                        self._assignment_sums_all[i, k] * rtd_x)

                n_i += self._assignment_sums_all[i, k]

            precision_scalar = n_i / model._mixture._variance

            for x in range(3):
                t._translation[x] = (rands[i, x] / sqrt(precision_scalar) +
                                     diff_sum[x] / n_i)



# @cython.boundscheck(False)
# @cython.cdivision(True)
# @cython.wraparound(False)
cdef class IsoTransformSymmetrySampler:
    """Sample mixture and transformations given count data.

    Subject to non-trivial symmetry constraints.

    The model passed as parameter to many of the methods consists
    only of the asymmetric components. Together with all the
    transformations of course.
    """

    def __cinit__(self, GridDataList data, IsoTransformPrior prior,
                  double std_range, list symmetries):
        """
        Args:
            data: GridDataList or GridDataNestedList object.
            prior: IsoTransformPrior object.
            std_range: Components will be evaluated at grid points
                within this multiple of their std. Example: 5.0.
            symmetries: list of Rotation3D objects. Symmetry group.
        """

        self._data = data
        self._prior = prior
        self._std_range = std_range
        self._symmetries = symmetries

        self._rotation_beta = 1.0 # For annealing rotations.
        self._counts_beta = 1.0 # For annealing by reducing counts.


        self._num_components = 0 # Number of asymmetric mixture components.
        self._num_projections = data.get_num_projections() #len(data)
        self._num_symmetries = len(symmetries) # Order of symmetry group.

        # Keep track of number of counts assigned for each direction.
        # This is equal to or slightly less than the number
        # of counts given by griddata for each direction.
        self._true_counts = np.zeros(self._num_projections, dtype='int')

        self._max_std = 0.0

        self._multinomial = IsoGridMultinomial()

        self._exptrace = ExpTrace()

    @property
    def data(self):
        return self._data

    @property
    def prior(self):
        return self._prior

    @property
    def rotation_beta(self):
        return self._rotation_beta

    def set_rotation_beta(self, value):
        self._rotation_beta = value

    @property
    def counts_beta(self):
        return self._counts_beta

    def set_counts_beta(self, value):
        self._counts_beta = value

    @property
    def std_range(self):
        return self._std_range

    cpdef init_memory(self, int K, double std):
        """Initialise large arrays.

        This is done at the beginning, and every time K changes,
        so that we don't re-assign the memory after every step.

        Args:
            K: int, Number of components for new models.
        """

        cdef:
            int P

        P = self._num_projections
        self._num_components = K
        KM = self._num_components * self._num_symmetries
        self._sampled_means = np.empty([P, KM, 3], dtype='float')


        # Initialise assignment sums. For each projection and each
        # component, this is the number of counts assigned to that
        # component.
        self._assignment_sums_all = np.empty([P, KM], dtype='int')

        self._assignment_sums = np.empty(KM, dtype='int')

        self._mixture_expanded = IsoMixture3D(
            np.zeros([KM, 3], dtype='float'),
            1.0, np.zeros(KM, dtype='float'))
        self._transformed = IsoMixture3D(np.zeros([KM, 3], dtype='float'),
                                         1.0, np.zeros(KM, dtype='float'))
        self._projected = IsoMixture2D(np.zeros([KM, 2], dtype='float'),
                                       1.0, np.zeros(KM, dtype='float'))

        self._max_std = 2.5 * std
        self._max_variance = self._max_std ** 2
        self._multinomial.init_memory(self._data, self._max_std,
                                      self._std_range,
                                      KM)

    cpdef init_memory_check(self, IsoTransform model):
        """Initialise memory if needed."""

        if (model._mixture._num_components != self._num_components or
            model._mixture._variance > self._max_variance):
            self.init_memory(model._mixture._num_components,
                             sqrt(model._mixture._variance))


    cpdef sample(self, IsoTransform model):
        """Generates a new sample given the current one.

        Uses Gibbs sampling.

        Args:
            model: IsoTransform object.
                Changed in-place.
        """

        ## times = []
        self.init_memory_check(model)
        ## times.append(time.time())
        self.sample_assignments(model)
        ## times.append(time.time())
        self.sample_missing_components(model)
        ## times.append(time.time())
        self.sample_weights(model)
        ## times.append(time.time())
        self.sample_means(model)
        ## times.append(time.time())
        self.sample_variance(model)
        ## times.append(time.time())
        self.sample_transformations(model)
        ## times.append(time.time())

        ## times = np.array(times)
        ## l = len(times)
        ## diffs = times[1:] - times[:(l-1)]
        ## p = diffs / np.sum(diffs)
        ## print((100*p).astype('i'))


    cpdef sample_mixture_part(self, IsoTransform model):
        """Samples everything but the transformations."""

        self.init_memory_check(model)
        self.sample_assignments(model)
        self.sample_missing_components(model)
        self.sample_weights(model)
        self.sample_means(model)
        self.sample_variance(model)

    cpdef sample_transformations_part(self, IsoTransform model):
        """Samples everything but the mixture."""

        self.init_memory_check(model)
        self.sample_assignments(model)
        self.sample_missing_components(model)
        self.sample_transformations(model)


    cpdef sample_rotations_part(self, IsoTransform model):
        """Samples everything but the mixture and the translations."""

        self.init_memory_check(model)
        self.sample_assignments(model)
        self.sample_missing_components(model)
        self.sample_rotations(model)

    cpdef sample_assignments(self, IsoTransform model):
        """Sample assignments for each image."""

        cdef:
            IsoMixture3D transformed
            int i, k, c, KM
            Transformation3D transformation
            IsoMixture3D mixture

        mixture = model._mixture
        mixture._expand(self._mixture_expanded, self._symmetries)
        KM = self._num_components * self._num_symmetries

        # Sample hard assignments for each direction.
        for i in range(self._num_projections):

            transformation = model._transformations[i]
            # mixture = model._mixture
            self._mixture_expanded._transform(
                transformation, self._transformed)
            self._transformed._project(self._projected)

            ## time_start = time.time()
            self._true_counts[i] = self._multinomial.sample(
                self._data.get(i, self._counts_beta),
                ## self._data[i], # pass self._counts_beta as well
                self._projected,
                self._std_range,
                i)
            ## if i == 0:
            ##     print 'sampling from multinomial in %1.4f' % (time.time() -
            ##                                                   time_start)
            ##     print '(should be around 0.037 seconds)'

            # Debug:
            ## a1 = np.sum(self._data[i].values)
            ## a2 = np.sum(self._multinomial._hard_values[
            ##             i, :self._true_counts[i]])
            ## if a1 != a2:
            ##     print 'Projection %d: %d / %d unassigned counts' % (
            ##         i, a1 - a2, a1)

            # Compute the total number of counts assigned to
            # each component for this direction.
            # Includes all extended components (symmetry copies).
            for k in range(KM):
                self._assignment_sums_all[i, k] = 0
            for c in range(self._true_counts[i]):
                self._assignment_sums_all[i, \
                    self._multinomial._hard_components[i, c]] += \
                        self._multinomial._hard_values[i, c]
            # self._assignment_sums_all: (P, KM)


        # Compute the total number of counts for each component
        # over all directions.
        for k in range(KM):
            self._assignment_sums[k] = 0
        self._assignment_sums_sum = 0
        for k in range(KM):
            for i in range(self._num_projections):
                self._assignment_sums[k] += self._assignment_sums_all[i, k]
            self._assignment_sums_sum += self._assignment_sums[k]
        # self._assignment_sums: (KM, )

    cdef sample_missing_components(self, IsoTransform model):
        """Sample missing z-component for each count.

        In reality we're computing the sufficient statistics
        for allowing the means and rotations to be sampled later.
        It's not necessary to sample every missing z-component.
        """

        cdef:
            double[:,::1] rands
            Transformation transformation
            int i, k, c, v, j, km, KM
            float px, py, std2
            GridDataInt griddata
            Grid2D grid
            IsoMixture3D mixture

        mixture = model._mixture
        mixture._expand(self._mixture_expanded, self._symmetries)
        KM = self._num_components * self._num_symmetries

        rands = np.random.randn(self._num_projections, KM)

        for i in range(self._num_projections):
            transformation = model._transformations[i]
            # mixture = model._mixture
            self._mixture_expanded._transform(
                transformation, self._transformed)
            griddata = self._data.get(i, self._counts_beta)
            grid = griddata._grid

            # Sample the missing components.
            for km in range(KM):
                self._sampled_means[i, km, 2] = (
                    self._transformed._means[km, 2] *
                    self._assignment_sums_all[i, km])
                std2 = (model._mixture._variance *
                        self._assignment_sums_all[i, km])
                self._sampled_means[i, km, 2] += (
                    rands[i, km] * sqrt(std2))

            # Compute the observed, deterministic part.
            # This is just a weighted sum of the grid points.
            for km in range(KM):
                self._sampled_means[i, km, 0] = 0
                self._sampled_means[i, km, 1] = 0
            for c in range(self._true_counts[i]):
                km = self._multinomial._hard_components[i, c]
                v = self._multinomial._hard_values[i, c]
                j = self._multinomial._hard_grid_indices[i, c]
                px = ((j / grid._shape_1) * grid._spacing_0 +
                      grid._origin_0)
                py = ((j % grid._shape_1) * grid._spacing_1 +
                      grid._origin_1)
                self._sampled_means[i, km, 0] += v * px
                self._sampled_means[i, km, 1] += v * py

    cdef sample_weights(self, IsoTransform model):
        """Sample mixture model weights."""

        cdef:
            double[::1] alpha

        # Contract assignment sums.
        # I.e. combine assignments for all copies of a given component
        # under the symmetry group.
        K = self._num_components
        M = self._num_symmetries
        assignment_sums_contracted = np.reshape(
            np.array(self._assignment_sums), (M, K)).sum(axis=0)

        # print 'weights:'
        # print assignment_sums_contracted

        alpha = (assignment_sums_contracted +
                 np.array(self._prior._mixture_prior._weights_prior.alpha))
        # print 'min alpha: ', np.sort(alpha)[:5]
        model._mixture._weights = np.random.dirichlet(alpha)
        # print np.array(model._mixture._weights)

    cdef sample_means(self, IsoTransform model):
        """Sample means of each component of mixture model."""

        cdef:
            int k, i, x, y, K, M, KM, P, im, km
            float s, precision_scalar, precision_scalar_initial, mean_x
            double diff
            double[::1] precision_mean
            double[:,::1] rands
            double[:,:,::1] rds
            double[:] trans
            Transformation t
            Rotation3D r, symm, rsymm

        K = self._num_components
        M = self._num_symmetries
        KM = K*M
        P = self._num_projections

        s = 1.0 / model._mixture._variance

        precision_mean = np.zeros(3, dtype='float')

        rands = np.random.randn(K, 3)

        # Transpose and scale rotations, component not needed.
        # Also include symmetries.
        rds = np.empty((P*M, 3, 3), dtype='float')
        for i in range(P):
            t = model._transformations[i]
            r = t._rotation
            for m in range(M):
                im = i*M + m
                symm = self._symmetries[m]
                rsymm = r.compose(symm).invert()
                for x in range(3):
                    for y in range(3):
                        rds[im, x, y] = rsymm._matrix[x, y] * s

        precision_scalar_initial = (1.0 / self._prior._mixture_prior.
                                    _mean_prior.variance)

        for k in range(K):

            for x in range(3):
                precision_mean[x] = 0.0

            # Set to 'r', will be 'n_k * s + r' after the for loop.
            precision_scalar = precision_scalar_initial

            for i in range(P):
                t = model._transformations[i]
                trans = t._translation
                for m in range(M):
                    km = m*K + k
                    im = i*M + m
                    for x in range(3):
                        for y in range(3):
                            diff = (self._sampled_means[i, km, y] -
                                    self._assignment_sums_all[i, km] *
                                    trans[y])
                            precision_mean[x] += (rds[im, x, y] * diff)
                    precision_scalar += self._assignment_sums_all[i, km] * s

            # Now precision_scalar = n_k * s + r.
            ## if k <= 5:
            ##     counts = ((precision_scalar - precision_scalar_initial) / s)
            ##     print('k = %d, std = %0.3f, mean_0 = %0.3f, counts = %1.1f' % (
            ##         k, 1.0 / sqrt(precision_scalar),
            ##         precision_mean[0] / precision_scalar, counts))

            for x in range(3):
                mean_x = precision_mean[x] / precision_scalar
                model._mixture._means[k, x] = (rands[k, x] /
                                               sqrt(precision_scalar) +
                                               mean_x)

        # print 'means:'
        # print np.array(model._mixture._means)

    cdef sample_variance(self, IsoTransform model):
        """Sample the variance of the mixture model."""

        cdef:
            int i, c, k, v, j, nik, km
            Transformation transformation
            IsoMixture3D transformed
            double a, px, py, dsq, b1, b2, b
            GridDataInt griddata
            Grid2D grid
            IsoMixture3D mixture

        mixture = model._mixture
        mixture._expand(self._mixture_expanded, self._symmetries)
        KM = self._num_components * self._num_symmetries


        # First hyperparameter.
        # print 'a = %1.1f * 1.5 + %1.1f' % (
        #     self._assignment_sums_sum,
        #     self._prior._mixture_prior._precision_prior._alpha)
        a = (self._assignment_sums_sum * 1.5 +
             self._prior._mixture_prior._precision_prior._alpha)

        # Second hyperparameter.
        # Part from observed components.
        b1 = 0.0

        for i, transformation in enumerate(model._transformations):
            self._mixture_expanded._transform(
                transformation, self._transformed)
            self._transformed._project(self._projected)
            griddata = self._data.get(i, self._counts_beta)
            ## griddata = self._data[i] # self._counts_beta
            grid = griddata._grid
            # Calculate distance between 2D grid points and projected means.
            for c in range(self._true_counts[i]):
                km = self._multinomial._hard_components[i, c]
                v = self._multinomial._hard_values[i, c]
                j = self._multinomial._hard_grid_indices[i, c]
                px = ((j / grid._shape_1) * grid._spacing_0 +
                      grid._origin_0)
                py = ((j % grid._shape_1) * grid._spacing_1 +
                      grid._origin_1)
                dsq = ((self._projected._means[km, 0] - px)**2 +
                       (self._projected._means[km, 1] - py)**2)
                b1 += dsq * v
        # Part from missing components.
        # Sum of sampled variances.
        b2 = 0.0
        for i in range(self._num_projections):
            for km in range(KM):
                nik = self._assignment_sums_all[i, km]
                if nik > 0:
                    ## b2 += np.random.chisquare(nik)
                    b2 += gsl_ran_chisq(r, nik)
        b2 *= model._mixture._variance

        ## print 'b = 0.5 * (%1.1f + %1.1f) + %1.1f' % (
        ##     b1, b2, self._prior._mixture_prior._precision_prior._beta)
        b = (0.5 * (b1 + b2) +
             self._prior._mixture_prior._precision_prior._beta)

        ## model._mixture._variance = 1.0 / np.random.gamma(a, 1.0 / b)
        model._mixture._variance = 1.0 / gsl_ran_gamma(r, a, 1.0 / b)

    cdef sample_transformations(self, IsoTransform model):
        """Sample rotations and translations for each direction."""

        self.sample_rotations(model)
        self.sample_translations(model)

    cdef sample_rotations(self, IsoTransform model):
        """Sample rotations for each direction."""

        cdef:
            Transformation t
            double coefficient, s, diff
            int i, ix, iy, k, KM, km, m
            double[:,::1] A
            double[:] trans
            IsoMixture3D mixture

        mixture = model._mixture
        mixture._expand(self._mixture_expanded, self._symmetries)
        KM = self._num_components * self._num_symmetries

        A = np.zeros((3, 3))

        for i in range(self._num_projections):
            if isinstance(self._prior._transformation_priors[i].
                          rotation_prior, FixedRotationPrior):
                continue

            coefficient = self._rotation_beta / model._mixture._variance

            # The for-loops below implement the following matrix
            # multiplication:
            ## temp = np.dot(transpose(self._sampled_means[i]),
            ##               model.mixture.means) # (3, 3)
            ## A = self._rotation_beta * temp / model.mixture.variance

            t = model._transformations[i]
            trans = t._translation

            for ix in range(3):
                for iy in range(3):
                    s = 0
                    for km in range(KM):
                        diff = (self._sampled_means[i, km, ix] -
                                self._assignment_sums_all[i, km] * trans[ix])
                        s += (diff * self._mixture_expanded._means[km, iy])
                    A[ix, iy] = s * coefficient

            # Sample from exp(tr(A^T B)).
            ## exp_trace = ExpTrace(A)
            ## exp_trace = ExpTrace()
            B = self._exptrace.sample(A)
            ## t = model.transformations[i]
            t._rotation._matrix = B



    cdef sample_translations(self, IsoTransform model):
        """Sample translations for each direction."""

        cdef:
            int i, k, x, y, KM, km
            double rtd_x, n_i, precision_scalar
            double[::1] diff_sum
            double[:,::1] rands
            double[:,:] r
            Transformation t
            IsoMixture3D mixture

        mixture = model._mixture
        mixture._expand(self._mixture_expanded, self._symmetries)
        KM = self._num_components * self._num_symmetries

        rands = np.random.randn(self._num_projections, 3)
        diff_sum = np.zeros(3, dtype='float')

        for i in range(self._num_projections):


            t = model._transformations[i]
            r = t._rotation._matrix

            n_i = 0.0

            for x in range(3):
                diff_sum[x] = 0.0

            for km in range(KM):

                for x in range(3): # row
                    rtd_x = 0.0
                    for y in range(3): # column
                        rtd_x += (r[x, y] * 
                                  self._mixture_expanded._means[km, y])
                            # R_i * \mu_k, row x

                    diff_sum[x] += (
                        self._sampled_means[i, km, x] -
                        self._assignment_sums_all[i, km] * rtd_x)

                n_i += self._assignment_sums_all[i, km]

            precision_scalar = n_i / model._mixture._variance

            for x in range(3):
                t._translation[x] = (rands[i, x] / sqrt(precision_scalar) +
                                     diff_sum[x] / n_i)







