"""Sampling IsoSubstructure objects from 3D grid data."""

import numpy as np
cimport numpy as np

import time

from libc.math cimport exp, sqrt

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)
gsl_rng_set(r, seed)

from cryoem.model.rotation cimport ExpTrace
from cryoem.model.transformation cimport Transformation, Transformation3D
from cryoem.model.isomixture cimport IsoMixture3D
from cryoem.model.prior cimport Dirichlet, Gamma
from cryoem.model.prior cimport IsoMixturePrior3D, IsoSubstructurePrior
from cryoem.data.grid cimport Grid3D
from cryoem.data.griddata cimport GridDataReal, GridDataInt
from cryoem.sampling.multinomial cimport IsoGrid3DMultinomial
from cryoem.model.isosubstructure cimport IsoSubstructure


cdef class IsoSubstructureSampler:
    """Sample mixtures with fixed substructure given 3D count data."""

    def __cinit__(self, GridDataInt data, IsoSubstructurePrior prior,
                  double std_range):

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

        self._num_components = 0
        self._max_variance = 0.0
        self._multinomial = IsoGrid3DMultinomial()
        self._exptrace = ExpTrace()

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

        At beginning, and every time K changes or component std
        becomes too large.

        Args:
            K: Total number of components.
            std: Maximum std of all components.
        """

        print 'Init memory.'

        self._num_components = K
        self._sampled_means = np.empty((K, 3), dtype='float')

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

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

        self._alpha = np.empty(K1 + 1, dtype='float')

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

        K = model._num_components
        K1 = model._num_variable_components
        var0 = model._mixture_fixed._variance
        var1 = model._mixture_variable._variance
        var_max = max(var0, var1)

        if K != self._num_components or var_max > self._max_variance:
            self.init_memory(K, sqrt(var_max), K1)

    cpdef sample(self, IsoSubstructure model):
        """Generates new model given current one using Gibbs sampling.

        Note that 'model' is changed in-place.
        """

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

    cdef sample_assignments(self, IsoSubstructure model):

        cdef:
            long[::1] index = np.empty(3, dtype='int')
            IsoGrid3DMultinomial m = self._multinomial
            Grid3D grid

        model.update_mixture_full()
        self._hard_size = m.sample_diagonal(
            self._data, model._mixture_full, self._std_range)

        # Compute the number of counts assigned to each component.
        for k in range(self._num_components):
            self._assignment_sums[k] = 0
        for c in range(self._hard_size):
            self._assignment_sums[m._hard_components[c]] += m._hard_values[c]

        # Add up the total number of assignments, i.e. the total counts.
        self._assignment_sums_sum = 0
        for k in range(self._num_components):
            self._assignment_sums_sum += self._assignment_sums[k]
        # Also compute the total counts for the submixtures.
        self._assignment_sums_fixed = 0
        for k in range(model._num_fixed_components):
            self._assignment_sums_fixed += self._assignment_sums[k]
        self._assignment_sums_variable = 0
        for k in range(model._num_fixed_components, model._num_components):
            self._assignment_sums_variable += self._assignment_sums[k]

        a_ratio = (float(self._assignment_sums_fixed) /
                   self._assignment_sums_sum)
        ## print 'Assignment sums fixed ratio = %1.4f' % a_ratio

        # Compute weighted sum of grid points for each component.
        # Will be used for:
        # sampling the means (last K1 components),
        # rotation and translation (first K0 components).
        # The weight is the number of counts from that grid point
        # assigned to the component.
        # Note: self._sampled_means.shape = (K, 3), not (K1, 3) or (K0, 3)
        grid = self._data._grid
        for k in range(model._num_components):
            for z in range(3):
                self._sampled_means[k, z] = 0.0
        for c in range(self._hard_size):
            k = self._multinomial._hard_components[c]
            v = self._multinomial._hard_values[c]
            j = self._multinomial._hard_grid_indices[c]
            # Convert linear index to 3D grid index.
            index[0] = (j / grid._shape_2) / grid._shape_1
            index[1] = (j / grid._shape_2) % grid._shape_1
            index[2] = j % grid._shape_2
            for z in range(3):
                self._sampled_means[k, z] += (
                    v * (index[z] * grid._spacing[z] + grid._origin[z]))

    cdef sample_weights(self, IsoSubstructure model):
        """Sample weights of variable mixture, and ratio of subsubstructure.
        """

        cdef:
            int k, K0
            IsoMixture3D m1

        # Prepare parameters for Dirichlet distribution.
        self._alpha[0] = (
            self._assignment_sums_fixed + self._prior._ratio_prior)
        K0 = model._num_fixed_components
        alpha_1 = self._prior._mixture_prior._weights_prior.alpha
        for k in range(model._num_variable_components):
            self._alpha[k+1] = self._assignment_sums[K0 + k] + alpha_1[k]

        # Sample from Dirichlet distribution.
        weights = np.random.dirichlet(np.array(self._alpha))

        model._ratio = weights[0]
        m1 = model._mixture_variable
        for k in range(model._num_variable_components):
            m1._weights[k] = (1 - model._ratio) * weights[k+1]

    cdef sample_means(self, IsoSubstructure model):
        """Sample means of variable mixture."""


        m1 = model._mixture_variable
        K0 = model._num_fixed_components

        s = 1.0 / m1._variance

        for k in range(m1._num_components):
            # Compute n_k * s + r.
            precision_scalar = (
                self._assignment_sums[k + K0] * s +
                1.0 / self._prior._mixture_prior._mean_prior.variance)

            std = 1.0 / sqrt(precision_scalar)

            for z in range(3):
                m = self._sampled_means[k + K0, z] * s / precision_scalar
                m1._means[k, z] = m + gsl_ran_gaussian(r, std)

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

        cdef:
            long[::1] index = np.empty(3, dtype='int')
            Grid3D grid = self._data._grid

        m1 = model._mixture_variable
        K0 = model._num_fixed_components

        # First hyperparameter.
        a = (self._assignment_sums_variable * 1.5 +
             self._prior._mixture_prior._precision_prior._alpha)

        # Second hyperparameter.
        b1 = 0.0

        for c in range(self._hard_size):
            k = self._multinomial._hard_components[c] - K0
            if k < 0: # Component belongs to fixed mixture.
                continue
            v = self._multinomial._hard_values[c]
            j = self._multinomial._hard_grid_indices[c]
            # Convert linear index to 3D grid index.
            index[0] = (j / grid._shape_2) / grid._shape_1
            index[1] = (j / grid._shape_2) % grid._shape_1
            index[2] = j % grid._shape_2
            for z in range(3):
                pz = index[z] * grid._spacing[z] + grid._origin[z]
                dsq = (m1._means[k, z] - pz) ** 2
                b1 += dsq * v

        b = 0.5 * b1 + self._prior._mixture_prior._precision_prior._beta

        m1._variance = 1.0 / gsl_ran_gamma(r, a, 1.0 / b)

    cpdef sample_transformation(self, IsoSubstructure model):

        self.sample_rotation(model)
        self.sample_translation(model)

    cdef sample_rotation(self, IsoSubstructure model):
        """Sample rotation of substructure given current translation."""

        A = np.zeros((3, 3))
        K0 = model._num_fixed_components
        t = model._substructure_transformation._translation
        m0 = model._mixture_fixed

        for ix in range(3):
            for iy in range(3):
                value = 0.0
                for k in range(K0):
                    diff = (self._sampled_means[k, ix] -
                            self._assignment_sums[k] * t[ix])
                    value += diff * m0._means[k, iy]
                A[ix, iy] = value / m0._variance

        B = self._exptrace.sample(A)
        model._substructure_transformation._rotation._matrix = B

    cdef sample_translation(self, IsoSubstructure model):
        """Sample translation of substructure given current rotation."""

        cdef:
            double[:] t, diff_sum
            double[:, :] rot
            IsoMixture3D m0

        K0 = model._num_fixed_components
        ts = model._substructure_transformation
        rot = ts._rotation._matrix
        t = ts._translation
        m0 = model._mixture_fixed

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

        for k in range(K0):
            for ix in range(3): # row
                # Compute R * \mu_k.
                rm_x = 0.0
                for iy in range(3): # column
                    rm_x += rot[ix, iy] * m0._means[k, iy]
                diff_sum[ix] += (
                    self._sampled_means[k, ix] -
                    self._assignment_sums[k] * rm_x)

        precision_scalar = (
            self._assignment_sums_fixed / m0._variance +
            1.0/self._prior._transformation_prior._translation_prior.variance)
        std = 1.0 / sqrt(precision_scalar)
        mean = np.array(diff_sum) / (precision_scalar * m0._variance)

        for ix in range(3):
            t[ix] = mean[ix] + gsl_ran_gaussian(r, std)


