from cryoem.data.grid cimport Grid
from cryoem.data.grid cimport Grid2D, Grid3D
from cryoem.model.transformation cimport Transformation3D

cdef class GridDataReal:
    cdef Grid _grid
    cdef double[:] _values
    cpdef GridDataReal add(self, GridDataReal other)
    cpdef GridDataReal add_squared(self, GridDataReal other)
    cpdef add_noise(self, double std)

cdef class GridDataInt:
    ## cdef Grid2D _grid
    cdef Grid _grid
    cdef long[::1] _values
    cdef long _values_sum
    cdef long[::1] _cumsum

cdef class GridDatasNested:
    cdef GridDataReal _griddata_real
    cdef GridDataInt _default
    cdef long _max_counts, _C
    cdef long[::1] _indices
    cpdef GridDataInt from_num_counts(self, int num_counts)
    cpdef int beta_to_counts(self, double beta)
    cpdef long[::1] beta_array_to_counts(self, double[::1] betas)
    cpdef GridDataInt from_beta(self, double beta)

cdef class GridDataList:
    cdef list _data
    cpdef get(self, index, beta)
    cpdef get_num_projections(self)
    cpdef as_list(self)

cdef class GridDataNestedList(GridDataList):
    cpdef get(self, index, beta)

cdef class GridDataProjector:
    cdef:
        int _shape, _N, _order
        double _spacing
        double[:] _evaluated3D, _alpha, _beta
        double[:, :] _transformed, _transformed_transpose
        GridDataReal _reference, _projection
        Grid2D _grid2D
        Grid3D _grid3D
    cpdef GridDataReal project(self, Transformation3D transformation)
