
from cryoem.model.isomixture cimport GeneralIsoMixtureNoise1D
from cryoem.model.distribution cimport Distribution, RealDistribution
from cryoem.model.isomixture cimport IsoMixture
from cryoem.model.rotation cimport Rotation3D
from cryoem.model.gaussian cimport IsoGaussianND, Gaussian1D
from cryoem.model.isomixture cimport IsoMixture1D, IsoMixture3D
from cryoem.model.isotransform cimport IsoTransform, IsoTransformNoise


cdef class IsoMixturePrior(Distribution):
    ## cdef IsoGaussianND _mean_prior
    cdef RealDistribution _mean_prior
    cdef Gamma _precision_prior
    cdef Dirichlet _weights_prior
    cpdef double log_likelihood(self, IsoMixture mixture)


cdef class IsoMixturePrior1D(IsoMixturePrior):
    ## def sample(self)
    pass


cdef class IsoMixturePrior2D(IsoMixturePrior):
    pass


cdef class IsoMixturePrior3D(IsoMixturePrior):
    pass


cdef class IsoMixtureConjugatePriorND(Distribution):
    cdef:
        NormalGammaND _mean_precision_prior
        Dirichlet _weights_prior


cdef class ScaledMixturePrior1D(Distribution):
    cdef:
        IsoMixturePrior1D _mixture_prior
        Gamma _scale_prior, _precision_prior

    ## cpdef sample(self)

cdef class GeneralIsoMixture1DPrior(IsoMixturePrior):
    pass

cdef class GeneralIsoMixtureNoise1DPrior(Distribution):
    cdef:
        Gaussian1D _mean_prior
        Gamma _component_precision_prior
        Gamma _weight_prior
        Gamma _noise_precision_prior
        int _num_components
    cpdef sample(self)
    cpdef energy_means(self, IsoMixture1D mixture)
    cpdef energy_gradient_means(self, IsoMixture1D mixture,
                                double[:] gradient)
    cpdef energy_weights(self, IsoMixture1D mixture)
    cpdef energy_gradient_weights(self, IsoMixture1D mixture,
                                  double[:] gradient)

cdef class GeneralIsoMixturePrior3D(Distribution):
    cdef:
        int _num_components
        IsoGaussianND _mean_prior
        Gamma _precision_prior, _weight_prior
    cpdef sample(self)

## cdef class GeneralIsoMixtureNoise3DPrior(Distribution):
##     cdef:
##         IsoGaussianND _mean_prior
##         Gamma _component_precision_prior
##         Gamma _weight_prior
##         Gamma _noise_precision_prior
##         int _num_components
##     cpdef sample(self)


cdef class Gamma(Distribution):
    cdef:
        int _constant
        float _alpha, _beta
    cpdef double log_likelihood(self, double x)

cdef class NormalGammaND(Distribution):
    cdef:
        double _variance_coefficient
        IsoGaussianND _mean_prior
        Gamma _precision_prior


cdef class Dirichlet(Distribution):
    cdef:
        double[:] _alpha
        int _dimension


cdef class TransformationPrior(Distribution):
    cdef RotationPrior _rotation_prior
    ## cdef TranslationPrior _translation_prior
    cdef Distribution _translation_prior


## cdef class FixedTransformationPrior(TransformationPrior):
##     pass


cdef class RotationPrior(Distribution):
    cdef Rotation3D _rotation


cdef class UniformRotationPrior(RotationPrior):
    pass

cdef class UniformRotation2DPrior(RotationPrior):
    pass

cdef class FixedRotationPrior(RotationPrior):
    pass


cdef class TranslationPrior(Distribution):
    pass


cdef class FixedTranslationPrior(TranslationPrior):
    cdef double[:] _translation



cdef class IsoTransformPrior(Distribution):
    cdef IsoMixturePrior _mixture_prior
    cdef list _transformation_priors
    cdef int _num_projections



cdef class IsoTransformNoisePrior(Distribution):
    cdef:
        GeneralIsoMixturePrior3D _mixture_prior
        list _transformation_priors
        Gamma _noise_prior
        int _num_projections
    cpdef log_likelihood(self, IsoTransformNoise model)
    cpdef energy(self, IsoTransformNoise model)
    cpdef energy_mixture(self, IsoMixture3D mixture)
    cpdef energy_means(self, IsoMixture3D mixture)
    cpdef energy_gradient_means(
        self, IsoMixture3D mixture, double[:] gradient)
    cpdef energy_weights(self, IsoMixture3D mixture)
    cpdef energy_gradient_weights(self, IsoMixture3D mixture,
                                  double[:] gradient)
    cpdef energy_component_precision(self, IsoMixture3D mixture)
    cpdef double energy_precisions(self, double[:] precisions)

cdef class IsoSubstructurePrior(Distribution):
    cdef:
        IsoMixture3D _mixture_fixed
        IsoMixturePrior3D _mixture_prior
        double _ratio_prior
        TransformationPrior _transformation_prior
