from cryoem.model.isomixture cimport IsoMixture2D, IsoMixture3D
from cryoem.model.transformation cimport Transformation3D
from cryoem.data.griddata cimport GridDatasNested, GridDataNestedList
from cryoem.data.griddata cimport GridDataInt, GridDataReal
from cryoem.data.grid cimport Grid2D


cdef class States:
    cdef:
        double[:, ::1] _rotations
        ## double[:] _betas
        IsoMixture3D _mixture #, _transformed
        ## IsoMixture2D _projected
        int _num_states #, _num_replicas
        ## Transformation3D _transformation
        StatesGridEnergies _state_energy
    ## cpdef init_memory(self, int K)
    ## cpdef set_betas(self, double[:] betas)
    ## cpdef get_energy(self)
    ## cpdef set_energy(self, energy)
    ## cpdef compute_energy(self, list griddatas)
    cpdef energy_at_beta_exact(self, double beta, GridDatasNested data,
                               double[::1] energy)
    ## cpdef energy_at_beta(self, double beta, double[::1] energy)
    cpdef get_direct_energy_calculator(self)
    cpdef get_grid_energy_calculator(self, Grid2D grid)


cdef class StatesDirectEnergies:
    cdef:
        States _states
        int _num_states
    cpdef energy_vector_in_place(self, GridDataInt griddata,
                                 double[::1] energies)
    cpdef energy_matrix_in_place(self, list griddatas,
                                 double[:, :] energies)
    cpdef energy_vector(self, GridDataInt griddata)
    cpdef energy_matrix(self, list griddatas)


cdef class StatesGridEnergies:
    cdef:
        int _num_states, _grid_size
        double[:, :] _energy
    cpdef energy_vector_in_place(self, GridDataInt griddata,
                                 double[::1] energies)
    cpdef energy_matrix_in_place(self, list griddatas,
                                 double[:, :] energies)
    cpdef energy_vector(self, GridDataInt griddata)
    cpdef energy_matrix(self, list griddatas)

    ## cpdef compute(self, GridDataInt griddata, double[::1] energy_out)


cdef class StatesGridEnergiesGaussian:
    cdef:
        int _num_states, _grid_size
        double[:] _precisions
        double[:, :] _predicted

    cpdef energy_vector_in_place(self, GridDataReal griddata,
                                 double[::1] energies,
                                 double beta_i)
    cpdef energy_vector(self, GridDataReal griddata, double beta_i)
    


cdef class DOSEstimator:
    cdef:
        double[:, :] _energy
        double[:] _free_energy
        int _num_replicas, _num_states, _verbose
        double _tol
    cpdef set_verbose(self, verbose)
    cdef double rel_error(self, double a, double b)
    cpdef run(self, int num_steps)
    cpdef double[:] get_free_energy(self)


cdef class DOS:
    cdef:
        States _states
        double[::1] _log_dos, _state_array, _log_p1, _log_p2
        double[::1] _energy_1, _energy_2
        int _num_states, _num_replicas
    ## cpdef double[:, :] sample_state(self, double beta)
    cpdef double swap_rate(self, double beta_1, double beta_2,
                           GridDatasNested data)
    cdef double log_normalisation_constant(self, double[::1] energy)
    cdef compute_log_probability_vector(self,
                                        double[::1] energy,
                                        double log_constant,
                                        double[::1] p_vector)



## cdef double log_sum_exp(double[::1] x)

cdef double sum_array(double[::1] x)
