"""For sampling isotropic ND mixtures given ND point cloud or counts.

Should work at least for 2D and 3D mixtures.
"""

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.model.isomixture cimport IsoMixture
from cryoem.model.prior cimport Dirichlet, Gamma


cdef class IsoMixtureCloudSampler:
    """Sample mixture given ND point data."""

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

        self._data = data
        self._prior = prior

        self._num_components = 0

        self._num_data_points = len(data) # N
        self._dimension = data.shape[1]
        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
        D = self._dimension

        self._num_components = K
        self._sampled_means = np.empty((K, D), 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, IsoMixture 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, IsoMixture 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, IsoMixture 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
        D = self._dimension
        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(D):
                    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, IsoMixture 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, IsoMixture mixture):
        """Sample mixture model means."""

        cdef:
            int k, c, j, z, D
            double v, s, std, rs, vc

        s = 1.0 / mixture._variance
        D = self._dimension

        # 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(D):
                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(D):
                    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(D):
                m = self._sampled_means[k, z] * s / precision_scalar
                mixture._means[k, z] = m + gsl_ran_gaussian(r, std)

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

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

        D = self._dimension

        # 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(D):
                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(D):
                    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)




