from cryoem.model.isomixture cimport IsoMixture, IsoMixture3D, IsoMixture2D
from cryoem.model.transformation cimport Transformation
from cryoem.data.grid cimport Grid2D
from cryoem.data.griddata cimport GridDataReal



cdef class IsoTransform:
    cdef IsoMixture _mixture
    cdef list _transformations


cdef class IsoTransformNoise:
    cdef:
        int _num_projections, _num_components, _grid_size, _unnormalised
        double[:] _precisions, _residuals2, _predicted, _residuals
        double[:,:] _point, _point2, _weight_accum, _precision_accum
        double[::1] _point3
        list _transformations, _inverse_transformations
        Grid2D _grid
        IsoMixture3D _mixture, _transformed
        IsoMixture2D _projected
        GridDataReal _griddata
    cpdef set_all_precisions(self, float value)
    cdef init_memory(self, data)
    cdef residual(self, double[:] a, double[:] b)
    cpdef double energy_single(self, data,
                               int projection_index, double std_range)
    cpdef energy(self, data, double std_range)
    cpdef energy_gradient_means(self, data, double[:] gradients,
                               double std_range)
    cpdef energy_gradient_weights(self, data, double[:] gradients,
                               double std_range)
    cpdef energy_gradient_weight(self, data, double std_range)
    cpdef energy_gradient_component_precision(self, data, double std_range)
    cpdef energy_gradient_noise_precision(self, data, double std_range)
    cpdef energy_gradient_weight_precision(
        self, data,
        double[:] gradients,
        double[:,:] gradients2,
        double std_range)
    cpdef optimal_weight(self, data, double std_range)
    cpdef energy_gradient_rotation(
        self, data, int projection_index, double[:, :] gradient,
        double std_range)
