"""For sampling isotropic mixture given 3D count data."""

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

from scipy.spatial import KDTree

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
from cryoem.model.isotransform cimport IsoTransform
from cryoem.model.transformation cimport Transformation, 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, IsoMixturePrior3D
from cryoem.data.grid cimport Grid2D, Grid3D
from cryoem.data.griddata cimport GridDataReal
from cryoem.sampling.multinomial cimport IsoGrid3DMultinomial
from cryoem.model.math cimport exp_fast



cdef class IsoMixture3DSampler:
    """Sample mixture given 3D count data."""

    def __cinit__(self, GridDataInt data, IsoMixturePrior3D 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()

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

        This is done at the beginning, and every time K changes or
        the component std becomes too large.

        Args:
            K: Number of components for new models.
            std: Current component std.
        """

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

        # The number of counts assigned to each component.
        self._assignment_sums = np.empty(K, dtype='int')

        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)

    cpdef init_memory_check(self, IsoMixture3D mixture):
        """Initialise memory if needed."""

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

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

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

        ## times = []
        self.init_memory_check(mixture)
        ## times.append(time.time())
        self.sample_assignments(mixture)
        ## times.append(time.time())
        self.sample_weights(mixture)
        ## times.append(time.time())
        self.sample_means(mixture)
        ## times.append(time.time())
        self.sample_variance(mixture)
        ## 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'))



    cdef sample_assignments(self, IsoMixture3D mixture):

        cdef:
            int k, c
            IsoGrid3DMultinomial m = self._multinomial

        self._hard_size = m.sample(self._data, mixture, 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]

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

        alpha = (np.array(self._assignment_sums) +
                 np.array(self._prior._weights_prior.alpha))

        mixture._weights = np.random.dirichlet(alpha)

    cdef sample_means(self, IsoMixture3D mixture):
        """Sample mixture model means."""

        cdef:
            int k, c, j, z
            double v, s, std
            Grid3D grid
            long[::1] index = np.empty(3, dtype='int')

        s = 1.0 / mixture._variance
        grid = self._data._grid

        # Compute the mean of the conditional on the mean.
        # For each component, it's the mean of the counts assigned to
        # that component.
        # Here we just compute the weighted sum of the grid points.
        # (Where the weights are the counts.)
        for k in range(self._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]))

        for k in range(self._num_components):

            # Compute n_k * s + r.
            precision_scalar = (self._assignment_sums[k] * s +
                                1.0 / self._prior._mean_prior.variance)

            std = 1.0 / sqrt(precision_scalar)

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

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

        cdef:
            int c, k, j, z
            double v, a, b1, b, dsq, pz
            long[::1] index = np.empty(3, dtype='int')
            Grid3D grid = self._data._grid


        # First hyperparameter.
        a = (self._assignment_sums_sum * 1.5 +
             self._prior._precision_prior.alpha)

        # Second hyperparameter.
        b1 = 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):
                pz = index[z] * grid._spacing[z] + grid._origin[z]
                dsq = (mixture._means[k, z] - pz) ** 2
                b1 += dsq * v

        b = 0.5 * b1 + self._prior._precision_prior.beta

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



cdef class IsoMixture3DCloudSampler:
    """Sample mixture given 3D point data."""

    def __cinit__(self, double[:, :] data, prior):
        """
        Args:
            data: (N, 3) double array.
            prior: Prior on the mixture hyperparameters.
                Either IsoMixturePrior3D or IsoMixtureConjugatePriorND.
        """

        self._data = data
        self._prior = prior

        self._num_components = 0

        self._num_data_points = len(data) # N
        self._assignments = np.zeros(
            self._num_data_points, dtype='int')

        self._fix_weights = False
        self._fix_std = False

    cpdef init_memory(self, int K):
        """Initialise memory that depends on K.

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

        N = self._num_data_points

        self._num_components = K
        self._sampled_means = np.empty((K, 3), dtype='float')
        self._distances = np.zeros(K, dtype='float')
        # The number of counts assigned to each component.
        self._assignment_sums = np.empty(K, dtype='int')

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

    cpdef init_memory_check(self, IsoMixture3D mixture):
        """Initialise memory if needed."""

        if mixture._num_components != self._num_components:
            self.init_memory(mixture._num_components)

    cpdef fix_weights(self, bint value):
        """If True, then weights will not be sampled."""
        self._fix_weights = value

    cpdef fix_std(self, bint value):
        """If True, then std will not be sampled."""
        self._fix_std = value

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

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

        # times = []
        self.init_memory_check(mixture)
        # times.append(time.time())
        self.sample_assignments(mixture)
        # times.append(time.time())
        if not self._fix_weights:
            self.sample_weights(mixture)
        # times.append(time.time())
        if not self._fix_std:
            self.sample_variance(mixture)
        # times.append(time.time())
        self.sample_means(mixture)
        # 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'))

        # print np.array(self._assignments)[:20]
        # print mixture.weights
        # print mixture.means[:10, 0]
        # print mixture.std



    cdef sample_assignments(self, IsoMixture3D mixture):

        cdef:
            int N, Ki, j, k, l, kk
            double c2, epsilon, t, acc, v, sum_i, w_k, log_w_k
            long[:] mean_indices
            double[:, :] data, soft_assignments

        # Compute soft assignments.
        N = self._num_data_points
        K = mixture._num_components
        data = self._data
        soft_assignments = self._soft_assignments
        epsilon = 1e-12

        c2 = -0.5 / mixture._variance

        # Using logs and normalisation. Use this if underflow problems
        # start occurring.
        # for i in range(N):
        #     log_w_k = log(mixture._weights[k])
        #     for k in range(K):
        #         v = 0.0
        #         for j in range(3):
        #             v += (data[i, j] - mixture._means[k, j]) ** 2
        #         v = log_w_k + c2*v
        #         soft_assignments[i, k] = v

        # # Normalise log-assignments.
        # # For each i the sum over k should be 1.
        # for i in range(N):
        #     max_value = soft_assignments[i, 0]
        #     for k in range(K):
        #         if soft_assignments[i, k] > max_value:
        #             max_value = soft_assignments[i, k]
        #     sum_exp = 0.0
        #     for k in range(K):
        #         sum_exp += exp(soft_assignments[i, k] - max_value)
        #     log_sum_exp = log(sum_exp) + max_value
        #     for k in range(K):
        #         soft_assignments[i, k] -= log_sum_exp
        #         # Convert from log-assignments to assignments.
        #         soft_assignments[i, k] = exp(
        #             soft_assignments[i, k])


        # Without normalisation:
        for i in range(N):
            w_k = mixture._weights[k]
            for k in range(K):
                v = 0.0
                for j in range(3):
                    v += (data[i, j] - mixture._means[k, j]) ** 2
                v = w_k * exp(c2 * v)
                soft_assignments[i, k] = v

        # Sample hard assignments.
        for i in range(N):
            sum_i = 0.0
            for k in range(K):
                sum_i += soft_assignments[i, k]
            # if sum_i < epsilon:
            #     self._assignments[i] = -1
            # else:
            # Sample from multinomial.
            t = gsl_ran_flat(r, 0, 1) * sum_i
            acc = 0.0
            l = 0
            while l < K and acc < t:
                acc += soft_assignments[i, l]
                l += 1
            self._assignments[i] = l-1
                

        # Compute the number of counts assigned to each component.
        for k in range(self._num_components):
            self._assignment_sums[k] = 0
        for i in range(self._num_data_points):
            kk = self._assignments[i]
            if kk >= 0: # nearby components exist
                self._assignment_sums[kk] += 1

        # Add up the total number of assignments.
        self._assignment_sums_sum = 0
        for k in range(self._num_components):
            self._assignment_sums_sum += self._assignment_sums[k]



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

        cdef:
            Dirichlet weights_prior

        alpha = (np.array(self._assignment_sums) +
                 np.array(self._prior.weights_prior.alpha))

        mixture._weights = np.random.dirichlet(alpha)

    cdef sample_means(self, IsoMixture3D mixture):
        """Sample mixture model means."""

        cdef:
            int k, c, j, z
            double v, s, std, rs, vc
            long[::1] index = np.empty(3, dtype='int')

        s = 1.0 / mixture._variance

        # Compute the mean of the conditional on the mean.
        # For each component, it's the mean of the points assigned to
        # that component.
        # Here we just compute the sum of the points.
        for k in range(self._num_components):
            for z in range(3):
                self._sampled_means[k, z] = 0.0
        for i in range(self._num_data_points):
            k = self._assignments[i]
            if k >= 0:
                for z in range(3):
                    self._sampled_means[k, z] += self._data[i, z]
                    
        if self._prior.is_conjugate():
            vc = self._prior.mean_precision_prior.variance_coefficient
            rs = s / vc
        else:
            rs = s

        for k in range(self._num_components):

            # Compute n_k * s + r.
            # precision_scalar = (
            #     self._assignment_sums[k] * s +
            #     1.0 / self._prior.mean_prior.variance)
            precision_scalar = self._assignment_sums[k] * s + rs

            std = 1.0 / sqrt(precision_scalar)

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

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

        cdef:
            int i, k, z
            double v, a, b1, b, dsq, vc, vr

        # First hyperparameter.
        if self._prior.is_conjugate():
            mp_prior = self._prior.mean_precision_prior
            a = ((self._assignment_sums_sum +
                  self._num_components) * 1.5 +
                 mp_prior.precision_prior.alpha)
        else:
            a = (self._assignment_sums_sum * 1.5 +
                 self._prior.precision_prior.alpha)
            

        # Second hyperparameter.
        b1 = 0.0
        for i in range(self._num_data_points):
            k = self._assignments[i]
            for z in range(3):
                dsq = (mixture._means[k, z] - self._data[i, z]) ** 2
                b1 += dsq
        if self._prior.is_conjugate():
            vc = self._prior.mean_precision_prior.variance_coefficient
            vr = 1.0 / vc
            for k in range(mixture._num_components):
                for z in range(3):
                    dsq = mixture._means[k, z]**2
                    b1 += vr * dsq

        if self._prior.is_conjugate():
            mp_prior = self._prior.mean_precision_prior
            b2 = mp_prior.precision_prior.beta
        else:
            b2 = self._prior.precision_prior.beta
            
        b = 0.5 * b1 + b2

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



