from cryoem.data.grid cimport Grid2D
from cryoem.data.griddata cimport GridDataInt, GridDataList, GridDataReal
from cryoem.model.prior cimport IsoTransformPrior
from cryoem.model.prior cimport IsoTransformNoisePrior
from cryoem.model.isotransform cimport IsoTransform, IsoTransformNoise
from cryoem.model.isomixture cimport IsoMixture3D, IsoMixture2D
from cryoem.sampling.multinomial cimport IsoGridMultinomial
from cryoem.model.rotation cimport ExpTrace



cdef class IsoTransformSampler:
    cdef:
        GridDataList _data
        IsoTransformPrior _prior
        double _rotation_beta, _counts_beta
        int _num_components
        int _num_projections
        double[:,:,::1] _sampled_means
        long[:,::1] _assignment_sums_all
        long[::1] _assignment_sums
        int _assignment_sums_sum
        long[::1] _true_counts
        double _std_range, _max_std, _max_variance
        ## long[:,::1] _hard_components
        ## long[:,::1] _hard_grid_indices
        ## long[:,::1] _hard_values
        IsoMixture3D _transformed
        IsoMixture2D _projected
        IsoGridMultinomial _multinomial
        ExpTrace _exptrace

    cpdef init_memory(self, int K, double std)
    cpdef init_memory_check(self, IsoTransform model)
    cpdef sample(self, IsoTransform model)
    cpdef sample_mixture_part(self, IsoTransform model)
    cpdef sample_transformations_part(self, IsoTransform model)
    cpdef sample_rotations_part(self, IsoTransform model)
    cpdef sample_assignments(self, IsoTransform model)
    cdef sample_missing_components(self, IsoTransform model)
    cdef sample_weights(self, IsoTransform model)
    cdef sample_means(self, IsoTransform model)
    cdef sample_variance(self, IsoTransform model)
    cdef sample_transformations(self, IsoTransform model)
    cdef sample_rotations(self, IsoTransform model)
    cdef sample_translations(self, IsoTransform model)


cdef class IsoTransformSymmetrySampler:
    cdef:
        int _num_components
        int _num_projections
        int _assignment_sums_sum
        int _num_symmetries
        double _rotation_beta, _counts_beta
        double[:,:,::1] _sampled_means
        long[:,::1] _assignment_sums_all
        long[::1] _assignment_sums
        long[::1] _true_counts
        double _std_range, _max_std, _max_variance
        list _symmetries
        GridDataList _data
        IsoTransformPrior _prior
        IsoMixture3D _transformed, _mixture_expanded
        IsoMixture2D _projected
        IsoGridMultinomial _multinomial
        ExpTrace _exptrace

    cpdef init_memory(self, int K, double std)
    cpdef init_memory_check(self, IsoTransform model)
    cpdef sample(self, IsoTransform model)
    cpdef sample_mixture_part(self, IsoTransform model)
    cpdef sample_transformations_part(self, IsoTransform model)
    cpdef sample_rotations_part(self, IsoTransform model)
    cpdef sample_assignments(self, IsoTransform model)
    cdef sample_missing_components(self, IsoTransform model)
    cdef sample_weights(self, IsoTransform model)
    cdef sample_means(self, IsoTransform model)
    cdef sample_variance(self, IsoTransform model)
    cdef sample_transformations(self, IsoTransform model)
    cdef sample_rotations(self, IsoTransform model)
    cdef sample_translations(self, IsoTransform model)



