from cryoem.data.griddata cimport GridDataInt
from cryoem.model.isomixture cimport IsoMixture3D
from cryoem.model.prior cimport IsoMixturePrior3D
from cryoem.sampling.multinomial cimport IsoGrid3DMultinomial



cdef class IsoMixture3DSampler:
    cdef:
        int _num_components, _hard_size, _assignment_sums_sum
        double _std_range, _max_std, _max_variance
        long[::1] _assignment_sums
        double[:, ::1] _sampled_means
        GridDataInt _data
        IsoMixturePrior3D _prior
        IsoGrid3DMultinomial _multinomial
    cpdef init_memory(self, int K, double std)
    cpdef init_memory_check(self, IsoMixture3D mixture)
    cpdef sample(self, IsoMixture3D mixture)
    cdef sample_assignments(self, IsoMixture3D mixture)
    cdef sample_weights(self, IsoMixture3D mixture)
    cdef sample_means(self, IsoMixture3D mixture)
    cdef sample_variance(self, IsoMixture3D mixture)


cdef class IsoMixture3DCloudSampler:
    cdef:
        int _num_components, _num_data_points, _assignment_sums_sum
        bint _fix_weights, _fix_std
        long[:] _assignments, _assignment_sums
        double[:] _distances
        double[:, :] _data, _sampled_means, _soft_assignments
        # object _data_tree
        # IsoMixturePrior3D _prior
        object _prior
    cpdef init_memory(self, int K)
    cpdef init_memory_check(self, IsoMixture3D mixture)
    cpdef fix_weights(self, bint value)
    cpdef fix_std(self, bint value)
    cpdef sample(self, IsoMixture3D mixture)
    cdef sample_assignments(self, IsoMixture3D mixture)
    cdef sample_weights(self, IsoMixture3D mixture)
    cdef sample_means(self, IsoMixture3D mixture)
    cdef sample_variance(self, IsoMixture3D mixture)    
