from cryoem.model.distribution cimport RealDistribution
from cryoem.data.grid cimport Grid2D, Grid3D, Grid1D
from cryoem.data.griddata cimport GridDataReal, GridDatasNested, GridDataInt
from cryoem.model.transformation cimport Transformation, Transformation3D
from cryoem.model.transformation cimport Transformation2D
from cryoem.model.isomixture cimport IsoMixture3D, IsoMixture2D

cdef class IsoMixture(RealDistribution):
    cdef:
        double[:,::1] _means
        double _variance
        double[::1] _weights
        int _num_components
        int _use_fast_exp
    cpdef set_all_weights(self, float value)
    cpdef set_use_fast_exp(self, value)
    
cdef class IsoMixture1D:
    cdef:
        double[::1] _means
        double _variance
        double[::1] _weights
        int _num_components

    cpdef evaluate_grid(self, Grid1D grid)
    cpdef GridDataInt sample_grid(self, double scale, Grid1D grid)
    cpdef copy(self)

cdef class GeneralIsoMixtureNoise1D:
    cdef:
        IsoMixture1D _mixture
        double _precision
    cpdef GridDataReal sample_grid(self, Grid1D grid)
    cpdef double energy(self, GridDataReal data)
    cpdef energy_gradient_means(self, GridDataReal data,
                                double[:] gradients)
    cpdef energy_gradient_weights(self, GridDataReal data,
                                  double[:] gradients)


cdef class IsoMixture2D(IsoMixture):
    cpdef evaluate_grid(self, Grid2D grid)
    cpdef evaluate_grid_in_place(self, Grid2D grid, double std_range,
                                 GridDataReal evaluated, int unnormalised)
    cpdef GridDataReal evaluate_grid_sparse(self,
                                            Grid2D grid, double std_range)
    cpdef double log_likelihood_grid_sparse(self, GridDataInt griddata,
                                            double std_range)
    cpdef log_likelihoods_nested(self, GridDatasNested nested,
                                 double[::1] betas,
                                 double[::1] result)
    cpdef _transform(self, Transformation2D transformation,
                     IsoMixture2D target)
    cpdef _project(self, IsoMixture1D target)


cdef class IsoMixture3D(IsoMixture):
    cpdef _transform(self, Transformation3D transformation,
                     IsoMixture3D target)
    cpdef _project(self, IsoMixture2D target)
    cpdef evaluate_grid(self, Grid3D grid)
    cpdef evaluate_grid_sparse(self, Grid3D grid, double std_range)
    cpdef evaluate_max_grid(self, Grid3D grid, double std_range)
    cpdef log_likelihood_sparse(self, double[:, :] x)
    cpdef log_likelihoods_sparse(self, double[:, :] x)    
    cpdef _expand(self, IsoMixture3D target, list symmetries)
    cpdef expand(self, symmetries)
    cpdef _get_means(self, double[:] x)
    cpdef _set_means(self, double[:] x)    
