from cryoem.data.griddata cimport GridDataNestedList
from cryoem.model.isotransform cimport IsoTransform, IsoTransformNoise

cdef class BruteForceRotationSampler:
    cdef:
        GridDataNestedList _data
        int _num_projections, _sample_size, _use_symmetries
        double _counts_beta
        double[:, ::1] _rotations
        list _symmetries


    cpdef set_symmetries(self, list symmetries)
    cpdef initialise_rotations(self, int sample_size)
    cpdef update_rotations(self, IsoTransform model, select_index)
    cpdef int sample_index(self, double[::1] energies, double[::1] weights)
    cpdef int optimise_index(self, double[::1] energies,
                             double[::1] weights)
    cpdef sample(self, IsoTransform model)
    cpdef optimise(self, IsoTransform model)
    cdef compute_weights(self, double[::1] energies, double[::1] weights)


cdef class BruteForceGaussianRotationSampler:
    cdef:
        int _num_projections, _sample_size, _use_symmetries
        int _unnormalised
        list _data, _symmetries
        double _annealing_beta
        double[:, ::1] _rotations
    cpdef set_symmetries(self, list symmetries)
    cpdef update_rotations(
        self, IsoTransformNoise model, select_index)    
    cpdef int sample_index(self, double[::1] energies,
                           double[::1] weights)
    cpdef int optimise_index(self, double[::1] energies,
                             double[::1] weights)
    cpdef sample(self, IsoTransformNoise model)
    cpdef optimise(self, IsoTransformNoise model)
    cdef compute_weights(self, double[::1] energies,
                         double[::1] weights)    
