from cryoem.data.griddata cimport GridDataInt
from cryoem.model.prior cimport IsoSubstructurePrior
from cryoem.sampling.multinomial cimport IsoGrid3DMultinomial
from cryoem.model.rotation cimport ExpTrace
from cryoem.model.isosubstructure cimport IsoSubstructure


cdef class IsoSubstructureSampler:
    cdef:
        int _num_components, _hard_size
        double _std_range, _max_variance, _max_std
        double _assignment_sums_sum, _assignment_sums_fixed
        double _assignment_sums_variable
        double[:] _assignment_sums, _alpha
        double[:, :] _sampled_means
        GridDataInt _data
        IsoSubstructurePrior _prior
        IsoGrid3DMultinomial _multinomial
        ExpTrace _exptrace
    cpdef init_memory(self, int K, double std, int K1)
    cpdef init_memory_check(self, IsoSubstructure model)
    cpdef sample(self, IsoSubstructure model)
    cdef sample_assignments(self, IsoSubstructure model)
    cdef sample_weights(self, IsoSubstructure model)
    cdef sample_means(self, IsoSubstructure model)
    cpdef sample_transformation(self, IsoSubstructure model)
    cdef sample_variance(self, IsoSubstructure model)
    cdef sample_rotation(self, IsoSubstructure model)
    cdef sample_translation(self, IsoSubstructure model)
