from cryoem.data.griddata cimport GridDataInt
from cryoem.model.isomixture cimport IsoMixture, IsoMixture3D, IsoMixture1D
from cryoem.model.isodiagonalmixture cimport IsoDiagonalMixture3D

cdef class IsoGridMultinomial:
    cdef:
        long[::1] _soft_components, _soft_grid_indices
        double[::1] _soft_values
        double[::1] _evaluated
        double[::1] _cumsums
        long[::1] _cum_index
        double[::1] _random_values
        long[:,::1] _hard_components
        long[:,::1] _hard_grid_indices
        long[:,::1] _hard_values
    cpdef init_memory(self, data, max_std, std_range, num_components)
    cpdef int sample(self,
                     GridDataInt griddata,
                     IsoMixture mixture,
                     float std_range,
                     int projection_index)


cdef class IsoGrid3DMultinomial:
    cdef:
        long[::1] _soft_components, _soft_grid_indices
        double[::1] _soft_values
        double[::1] _evaluated
        double[::1] _cumsums
        long[::1] _cum_index
        double[::1] _random_values
        long[::1] _hard_components
        long[::1] _hard_grid_indices
        long[::1] _hard_values
    cpdef init_memory(self, GridDataInt data,
                      max_std, std_range, num_components)
    cpdef int sample(self,
                     GridDataInt griddata,
                     IsoMixture3D mixture,
                     float std_range)
    cpdef int sample_diagonal(self,
                              GridDataInt griddata,
                              IsoDiagonalMixture3D mixture,
                              float std_range)


cdef class IsoTransformGrid2DMultinomial:
    cdef:
        long[::1] _soft_components, _soft_grid_indices
        double[::1] _soft_values
        double[::1] _evaluated
        double[::1] _cumsums
        long[::1] _cum_index
        double[::1] _random_values
        long[:,::1] _hard_components
        long[:,::1] _hard_grid_indices
        long[:,::1] _hard_values
    cpdef init_memory(self, data, max_std, std_range, num_components)
    cpdef int sample(self,
                     GridDataInt griddata,
                     IsoMixture1D mixture,
                     float std_range,
                     int projection_index)

