"""
Distributions used as priors.
"""

import numpy as np
cimport numpy as np

from libc.math cimport exp, sqrt, log

from scipy.special import gammaln

from cryoem.model.isomixture cimport IsoMixture, IsoMixture3D, IsoMixture1D
from cryoem.model.isomixture cimport IsoMixture2D
from cryoem.model.isomixture cimport GeneralIsoMixtureNoise1D
from cryoem.model.scaledmixture import ScaledIsoMixture1D
from cryoem.model.isotransform cimport IsoTransform, IsoTransformNoise
from cryoem.model.distribution import Distribution
from cryoem.model.transformation import Transformation, Transformation3D
from cryoem.model.rotation import Rotation3D, Rotation2D
from cryoem.model.gaussian cimport IsoGaussianND, Gaussian, Gaussian1D
from cryoem.model.isosubstructure cimport IsoSubstructure


cdef class IsoMixturePrior(Distribution):
    """Prior for isotropic Gaussian mixture model.

    Consists of separate priors on the mean (a single Gaussian or
    a mixture), the variance/precision (a gamma distribution on the
    precision) and the weights (Dirichlet distribution).
    """

    def __cinit__(self, RealDistribution mean_prior,
                  precision_prior, weights_prior):
        """Initialise priors on different parts of mixture model.

        Args:
            mean_prior: Distribution over R^n. Typically an
                IsoGaussianND or IsoMixture object. This is the prior
                over a single mean.
            precision_prior: Distribution over R, typically gamma.
            weights_prior: Distribution over the positive real.
                Typically a Dirichlet distribution.
                Could also be a Gamma distribution for unnormalised weights,
                in which case it is the prior over a single weight.
        """

        self._mean_prior = mean_prior
        self._precision_prior = precision_prior
        self._weights_prior = weights_prior

    @property
    def mean_prior(self):
        return self._mean_prior

    @property
    def precision_prior(self):
        return self._precision_prior

    @property
    def weights_prior(self):
        return self._weights_prior

    @property
    def num_components(self):
        return self._weights_prior.dimension

    @property
    def dimension(self):
        return self._mean_prior.dimension

    def is_conjugate(self):
        return False

    def sample(self):
        """Samples an isotropic Gaussian mixture model from this prior.

        Returns:
            IsoMixture object.
        """

        raise NotImplementedError

    cpdef double log_likelihood(self, IsoMixture mixture):
        """Returns the log-likelihood of 'mixture'.

        Args:
            mixture: IsoMixture object.

        Returns:
            The log-likelihood of 'mixture' as a double.
        """

        ll_means = self.mean_prior.log_likelihood(np.array(mixture.means))
        precision = 1.0 / mixture.variance
        ll_precision = self.precision_prior.log_likelihood(precision)
        ll_weights = self.weights_prior.log_likelihood(mixture.weights)

        return ll_means + ll_precision + ll_weights


cdef class IsoMixturePrior1D(IsoMixturePrior):
    """Prior for 1D isotropic mixture models."""

    def sample(self):
        """Samples an isotropic Gaussian mixture model from this prior.

        Returns:
            IsoMixture1D object.
        """

        means = self.mean_prior.sample(self.num_components)
        precision = self.precision_prior.sample()
        variance = 1.0 / precision
        weights = self.weights_prior.sample()

        return IsoMixture1D(means, variance, weights)


cdef class IsoMixturePrior2D(IsoMixturePrior):
    """Prior for 2D isotropic mixture models."""

    def sample(self):
        """Samples an isotropic Gaussian mixture model from this prior.

        Returns:
            IsoMixture2D object.
        """

        means = self.mean_prior.sample(self.num_components)
        precision = self.precision_prior.sample()
        variance = 1.0 / precision
        weights = self.weights_prior.sample()

        return IsoMixture2D(means, variance, weights)


cdef class IsoMixturePrior3D(IsoMixturePrior):
    """Prior for 3D isotropic mixture models."""

    def sample(self):
        """Samples an isotropic Gaussian mixture model from this prior.

        Returns:
            IsoMixture object.
        """

        means = self.mean_prior.sample(self.num_components)
        precision = self.precision_prior.sample()
        variance = 1.0 / precision
        weights = self.weights_prior.sample()

        return IsoMixture3D(means, variance, weights)


cdef class IsoMixtureConjugatePriorND(Distribution):
    """Prior for IsoMixture1D, IsoMixture2D and IsoMixture3D.

    Conjugate prior for isotropic mixture.
    Uses a Normal-Gamma distribution for the means and precision,
    and a Dirichlet distribution on the weights.

    Alternative to IsoMixturePrior1D, -2D and -3D.
    """

    def __cinit__(self, NormalGammaND mean_precision_prior,
                  Dirichlet weights_prior):
        """Initialise priors on different parts of mixture model.

        Args:
            mean_precision_prior: NormalGammaND object.
            weights_prior: Dirichlet distribution.
        """

        self._mean_precision_prior = mean_precision_prior
        self._weights_prior = weights_prior

    @property
    def mean_precision_prior(self):
        return self._mean_precision_prior

    @property
    def weights_prior(self):
        return self._weights_prior

    @property
    def num_components(self):
        return self._weights_prior.dimension

    @property
    def dimension(self):
        return self._mean_precision_prior.dimension

    def is_conjugate(self):
        return True

    def sample(self):
        """Samples an isotropic Gaussian mixture model from this prior.

        Returns:
            IsoMixture1D, -2D or -3D object, according to dimension.
        """

        means, precision = self._mean_precision_prior.sample(
            self.num_components)
        variance = 1.0 / precision
        weights = self.weights_prior.sample()

        d = self.dimension
        classes = [IsoMixture1D, IsoMixture2D, IsoMixture3D]
        if d > 3 or d < 1:
            raise ValueError(
                'Dimension should be between 1 and 3 (not %d)' % d)

        return classes[d-1](means, variance, weights)

    def log_likelihood(self, mixture):
        """Returns the log-likelihood of 'mixture'.

        Args:
            mixture: IsoMixture object.

        Returns:
            The log-likelihood of 'mixture' as a double.
        """

        precision = 1.0 / mixture.variance
        ll_means_precision = self._mean_precision_prior.log_likelihood(
            mixture.means, precision)
        ll_weights = self.weights_prior.log_likelihood(mixture.weights)

        return ll_means_precision + ll_weights


cdef class ScaledMixturePrior1D(Distribution):
    """Prior for scaled isotropic mixture."""

    def __cinit__(self, IsoMixturePrior1D mixture_prior,
                  scale_prior, precision_prior):
        """
        Args:
            mixture_prior: IsoMixturePrior1D object.
            scale_prior: Gamma distribution.
            precision_prior: Also Gamma distribution, Gaussian noise.
        """
        self._mixture_prior = mixture_prior
        self._scale_prior = scale_prior
        self._precision_prior = precision_prior

    def sample(self):
        """Samples a scaled mixture from the prior."""

        cdef:
            IsoMixturePrior1D mixture_prior

        mixture = self._mixture_prior.sample()
        scale = self._scale_prior.sample()
        precision = self._precision_prior.sample()

        return ScaledIsoMixture1D(mixture, scale, 1.0 / precision)


cdef class GeneralIsoMixture1DPrior(IsoMixturePrior):
    """Prior for general, i.e. unnormalised isotropic GMMS in 1D."""

    def sample(self):
        """Sample an IsoMixture1D object from the prior."""

        means = self.mean_prior.sample(self.num_components)
        precision = self.precision_prior.sample()
        variance = 1.0 / precision
        weights = self.weights_prior.sample(self.num_components)

        return IsoMixture1D(means, variance, weights)



cdef class GeneralIsoMixtureNoise1DPrior(Distribution):
    """Prior for IsoMixture1D, with arbitrary weights and Gaussian noise.
    """

    def __cinit__(self, mean_prior,
                  component_precision_prior, weight_prior,
                  noise_precision_prior,
                  num_components):
        """
        Args:
            mean_prior: Gaussian1D, over a single mean.
            component_precision_prior: Gamma distribution.
            weight_prior: Gamma distribution over single weight.
                Could also be a Gaussian distribution (1D).
            noise_precision_prior: Gamma distribution for noise precision.
            num_components: K, the number of components in the mixture.
        """
        self._mean_prior = mean_prior
        self._component_precision_prior = component_precision_prior
        self._weight_prior = weight_prior
        self._noise_precision_prior = noise_precision_prior
        self._num_components = num_components


    cpdef sample(self):
        """Samples a general isotropic GMM with noise level.

        Returns:
            GeneralIsoMixtureNoise1D object.
        """

        means = self._mean_prior.sample(self._num_components)
        component_precision = self._component_precision_prior.sample()
        weights = self._weight_prior.sample(self._num_components)
        noise_precision = self._noise_precision_prior.sample()
        mixture = IsoMixture1D(means, 1.0  / component_precision, weights)

        return GeneralIsoMixtureNoise1D(mixture, noise_precision)

    cpdef energy_means(self, IsoMixture1D mixture):
        """Return negative log-likelihood of means under this prior."""

        cdef:
            double c, t

        c = 0.5 / self._mean_prior._variance

        t = np.sum(np.array(mixture._means)**2)

        return c * t

    cpdef energy_gradient_means(self, IsoMixture1D mixture,
                                double[:] gradient):
        """Return gradient of energy w.r.t. means.

        Args:
            gradient: Where resulting gradient is stored.
        """

        cdef:
            int k, K
            double precision

        K = mixture._num_components
        precision = 1.0 / self._mean_prior._variance
        for k in range(K):
            gradient[k] = mixture._means[k] * precision

        return gradient

        ## return np.array(mixture._means) / self._mean_prior._variance

    cpdef energy_weights(self, IsoMixture1D mixture):
        """Return negative log-likelihood of weights under this prior."""

        cdef:
            Gamma weight_prior


        weight_prior = self._weight_prior

        weights = np.array(mixture._weights)

        # 15.4.2014, added minus sign.
        return (-(weight_prior._alpha - 1) * np.log(weights).sum() +
                weight_prior._beta * weights.sum())

    cpdef energy_gradient_weights(self, IsoMixture1D mixture,
                                  double[:] gradient):
        """Return gradient of energy w.r.t. weights."""

        cdef:
            Gamma weight_prior
            int k, K
            double alpha, beta

        K = mixture._num_components

        weight_prior = self._weight_prior
        alpha = weight_prior._alpha
        beta = weight_prior._beta

        for k in range(K):
            # 15.4.2014, added minus sign.
            gradient[k] = (1.0 - alpha) / mixture._weights[k] + beta

        return gradient


cdef class GeneralIsoMixturePrior3D(Distribution):
    """Prior for 3D mixture with unnormalised weights.

    Doesn't include projections or Gaussian noise.
    """

    def __cinit__(self, mean_prior, precision_prior,
                  weight_prior, num_components):
        """
        Args:
            mean_prior: IsoGaussianND, over a single mean.
            precision_prior: Gamma distribution.
            weight_prior: Gamma distribution over single weight.
            num_components: K, the number of components in the mixture.
        """

        self._mean_prior = mean_prior
        self._precision_prior = precision_prior
        self._weight_prior = weight_prior
        self._num_components = num_components

    cpdef sample(self):
        """Sample a IsoMixture3D with unnormalised weights."""

        means = self._mean_prior.sample(self._num_components)
        precision = self._precision_prior.sample()
        weights = self._weight_prior.sample(self._num_components)
        mixture = IsoMixture3D(means, 1.0  / precision, weights)

        return mixture


cdef class Gamma(Distribution):
    """Gamma distribution.

    Used as a prior over the precision of a Gaussian distribution.
    """

    def __cinit__(self, alpha, beta):
        """Initialise gamma distribution.

        Args:
            alpha: The shape as a double.
            beta: The rate as a double. Note that this is the inverse
                of the scale, denoted by theta.
        """

        self._alpha = alpha
        self._beta = beta

        # Set to true for sampling higher-dimensional vectors with the same
        # value.
        self._constant = 0

    @property
    def alpha(self):
        return self._alpha

    @property
    def beta(self):
        return self._beta

    def set_constant(self, value=True):
        if value:
            self._constant = 1
        else:
            self._constant = 0

    def sample(self, int num_samples=1):
        """Draw num_samples from distribution."""

        shape = self._alpha
        scale = 1.0 / self._beta

        if self._constant == 1:
            return np.array([np.random.gamma(shape, scale)] * num_samples)
        else:
            return np.random.gamma(shape, scale, num_samples)

    cpdef double log_likelihood(self, double x):
        """Returns the log likelihood of 'x'.

        Args:
            x: Single value at which to evaluate the distribution.

        Returns:
            The log-likelihood of 'x' as a double.
        """

        cdef:
            double constant, variable

        constant = self._alpha * log(self._beta) - gammaln(self._alpha)
        variable = (self._alpha - 1.0) * log(x) - self._beta * x

        return constant + variable


cdef class NormalGammaND(Distribution):
    """Conjugate prior for Gaussian with unknown mean and variance.

    Product of Gaussian and Gamma distributions, with the normal
    distribution depending on the Gamma distribution.
    """

    def __cinit__(self, mean, variance_coefficient, alpha, beta):
        """
        Args:
            mean: Mean vector. Usually shape is (2, ) or (3, ).
            variance_coefficient: (1/r) in latex notes.
            alpha: The shape as a double.
            beta: The rate as a double. Note that this is the inverse
                of the scale, denoted by theta.
        """

        self._variance_coefficient = variance_coefficient
        self._mean_prior = IsoGaussianND(mean, 1.0)
        self._precision_prior = Gamma(alpha, beta)

    @property
    def dimension(self):
        return self._mean_prior.dimension

    @property
    def variance_coefficient(self):
        return self._variance_coefficient

    @property
    def mean_prior(self):
        return self._mean_prior

    @property
    def precision_prior(self):
        return self._precision_prior

    def sample(self, num_samples=1):
        """Return pair (positions, precision).

        Returns:
            positions: (N, D) array. 
            precision: scalar.
        """

        precision = self._precision_prior.sample()
        self._mean_prior.set_variance(
            self._variance_coefficient / precision)
        positions = self._mean_prior.sample(num_samples)

        return (positions, precision)

    def log_likelihood(self, positions, precision):
        """Returns log-likelihood.

        Args:
            positions: (N, D) array.
            precision: scalar.
        """

        self._mean_prior.set_variance(
            self._variance_coefficient / precision)
        lla = self._mean_prior.log_likelihood(positions)
        llb = self._precision_prior.log_likelihood(precision)
        
        return lla + llb
        

cdef class Dirichlet(Distribution):
    """Dirichlet distribution.

    Used as a prior over the weights of a mixture model.
    """

    def __cinit__(self, alpha):
        """Initialise Dirichlet distribution.

        Args:
            alpha: (K, ) array of positive doubles.
        """

        self._alpha = alpha
        self._dimension = len(alpha)

    @property
    def alpha(self):
        return self._alpha

    @property
    def dimension(self):
        return self._dimension

    def sample(self):
        """Draw a single sample from the distribution.

        Returns:
            (K, ) array of non-negative reals that add up to one.
        """

        return np.random.dirichlet(self._alpha)

    def log_likelihood(self, x):
        """Returns the log likelihood of 'x'.

        Args:
            x: (K, ) array for evaluating the distribution.
                Typically the weights of a mixture.

        Returns:
            The log-likelihood of 'x' as a double.
        """

        constant = (gammaln(np.sum(self._alpha)) -
                    np.sum(gammaln(self.alpha)))
        variable = np.sum((np.array(self._alpha) - 1.0) * np.log(x))

        return constant + variable


cdef class IsoTransformPrior(Distribution):
    """Prior over IsoTransform object.

    Created for 3D case, should also work for 2D.
    """

    def __cinit__(self, mixture_prior, transformation_priors):
        """Initialise with two separate priors.

        To fix the transformations, use a Dirac delta distribution
        as the transformation prior.

        The same prior is used for all transformations.

        Args:
            mixture_prior: IsoMixturePrior object.
            transformation_prior: list of TransformationPrior objects.
        """

        self._mixture_prior = mixture_prior
        self._transformation_priors = transformation_priors

    @property
    def mixture_prior(self):
        return self._mixture_prior

    @property
    def transformation_priors(self):
        return self._transformation_priors

    @property
    def num_projections(self):
        return len(self._transformation_priors)

    def sample(self):
        """Draw a single sample from the prior."""

        mixture = self._mixture_prior.sample()
        transformations = [
            tp.sample() for tp in self._transformation_priors]

        return IsoTransform(mixture, transformations)

    def log_likelihood(self, model):
        """Returns the log-likelihood of 'model' under this prior.

        Args:
            model: IsoMixtureTransformations object.

        Returns:
            Log-likelihood of 'model'.
        """

        mixture_ll = self.mixture_prior.log_likelihood(model.mixture)
        transformations_ll = np.sum(np.array([
            tp.log_likelihood(t) for tp, t in
            zip(self.transformation_priors, model.transformations)]))

        return mixture_ll + transformations_ll


cdef class IsoTransformNoisePrior(Distribution):

    def __cinit__(self, mixture_prior, transformation_priors,
                  noise_prior):
        """Initialise with two separate priors.

        To fix the transformations, use a Dirac delta distribution
        as the transformation prior.

        The same prior is used for all transformations.

        Args:
            mixture_prior: GeneralIsoMixturePrior3D object.
            transformation_prior: list of TransformationPrior objects.
            noise_prior: Gamma prior on noise precision.
        """

        self._mixture_prior = mixture_prior
        self._transformation_priors = transformation_priors
        self._noise_prior = noise_prior

        self._num_projections = len(self._transformation_priors)

    @property
    def mixture_prior(self):
        return self._mixture_prior

    @property
    def transformation_priors(self):
        return self._transformation_priors

    @property
    def num_projections(self):
        return len(self._transformation_priors)

    def sample(self):
        """Draw a single sample from the prior."""

        mixture = self._mixture_prior.sample()
        mixture.set_use_fast_exp(False)
        transformations = [
            tp.sample() for tp in self._transformation_priors]
        precisions = self._noise_prior.sample(self._num_projections)

        return IsoTransformNoise(mixture, transformations, precisions)

    cpdef log_likelihood(self, IsoTransformNoise model):
        """Log-likelihood of model under prior.

        Args:
            model: IsoTransformNoise object.
        """

        return -self.energy(model)

    cpdef energy(self, IsoTransformNoise model):
        """Negative log-likelihood of model under this prior."""

        cdef:
            double E_mixture, E_precisions

        E_mixture = self.energy_mixture(model._mixture)
        E_precisions = self.energy_precisions(model._precisions)

        return E_mixture + E_precisions

    cpdef energy_mixture(self, IsoMixture3D mixture):
        """Negative log-likelihood of mixture under this prior.

        Everything except the transformations and Gaussian noise.
        """

        cdef:
            double E_means, E_weights, E_precision

        E_means = self.energy_means(mixture)
        E_weights = self.energy_weights(mixture)
        E_precision = self.energy_component_precision(mixture)

        return E_means + E_weights + E_precision

    cpdef energy_means(self, IsoMixture3D mixture):
        """Negative log-likelihood of means under this prior."""

        cdef:
            double energy
            int k, j
            double[:, :] means

        energy = 0.0
        means = mixture._means
        for k in range(mixture._num_components):
            for j in range(3):
                energy += (means[k, j] * means[k, j])

        c = 0.5 / self._mixture_prior._mean_prior._variance

        return c*energy

    cpdef energy_gradient_means(
        self, IsoMixture3D mixture, double[:] gradient):
        """Gradient of self.energy_means w.r.t. component means.

        Args:
            gradient: (K, 3) array for computed gradients.
        """

        cdef:
            int k, K, j
            double precision
            double[:,:] means

        K = mixture._num_components
        precision = 1.0 / self._mixture_prior._mean_prior._variance
        means = mixture._means
        for k in range(K):
            for j in range(3):
                gradient[3*k + j] = mixture._means[k, j] * precision


    cpdef energy_weights(self, IsoMixture3D mixture):
        """Return negative log-likelihood of weights under prior."""

        cdef:
            double energy, alpha, beta, part1, part2
            int K, k
            double[:] weights

        energy = 0.0
        alpha = self._mixture_prior._weight_prior._alpha
        beta = self._mixture_prior._weight_prior._beta
        weights = mixture._weights
        K = mixture._num_components

        part1 = 0.0
        part2 = 0.0
        for k in range(K):
            part1 += log(weights[k])
            part2 += weights[k]

        energy = (1.0 - alpha) * part1 + beta * part2

        constant = -(alpha * log(beta) - gammaln(alpha)) * K

        return energy + constant

    cpdef energy_gradient_weights(self, IsoMixture3D mixture,
                                  double[:] gradient):
        """Return gradient of energy w.r.t. weights."""

        cdef:
            int k, K
            double alpha, beta

        K = mixture._num_components
        alpha = self._mixture_prior._weight_prior._alpha
        beta = self._mixture_prior._weight_prior._beta

        for k in range(K):
            gradient[k] = (1.0 - alpha) / mixture._weights[k] + beta

        return gradient

    cpdef energy_component_precision(self, IsoMixture3D mixture):
        """Return negative log-likelihood of std under this prior."""

        cdef:
            double energy
            Gamma prior

        prior = self._mixture_prior._precision_prior
        energy = 0.0 - prior.log_likelihood(1.0 / mixture._variance)

        return energy

    cpdef double energy_precisions(self, double[:] precisions):
        """Return negative log-likelihood of precisions.

        Each precision is Gamma-distributed.
        """

        cdef:
            double energy
            Gamma prior

        prior = self._noise_prior
        for i in range(self._num_projections):
            energy -= prior.log_likelihood(precisions[i])

        return energy


cdef class TransformationPrior(Distribution):
    """Prior over transformations."""

    def __cinit__(self, rotation_prior, translation_prior):
        """Initialise with priors over the rotation and translation.

        Args:
            rotation_prior: A distribution over rotations.
                Either a DiracDistribution, or a UniformRotationPrior.
            translation_prior: A distribution over R^3.
                For instance a 3D Gaussian, or a DiracDistribution.
        """

        self._rotation_prior = rotation_prior
        self._translation_prior = translation_prior

    @property
    def rotation_prior(self):
        return self._rotation_prior

    @property
    def translation_prior(self):
        return self._translation_prior

    def sample(self):
        """Sample a Transformation object from the prior."""

        rotation = self.rotation_prior.sample()
        translation = np.array(self.translation_prior.sample()).flatten()

        return Transformation3D(rotation, translation)

    def log_likelihood(self, transformation):
        """Returns the log-likelihood of the given transformation.

        If any Dirac delta functions are involved, we just assume
        that the given parameters match those under the prior.
        Otherwise the log-likelihood would be -infinity.
        """

        rotation_ll = self.rotation_prior.log_likelihood(
            transformation.rotation)
        translation_ll = self.translation_prior.log_likelihood(
            transformation.translation)

        return rotation_ll + translation_ll


class TransformationPriors(object):
    """Factory class for creating TransformationPriors."""

    @staticmethod
    def create_fixed(transformation):

        return TransformationPrior(
            FixedRotationPrior(transformation.rotation),
            FixedTranslationPrior(transformation.translation))

    @staticmethod
    def create_uniform():
        """Returns a uniform rotation prior and a fixed translation prior.

        The translation prior is fixed to zero.
        """

        return TransformationPrior(
            UniformRotationPrior(), FixedTranslationPrior(np.zeros(3)))

    @staticmethod
    def create_uniform2D():
        """Returns a uniform rotation prior and a fixed translation prior.

        The translation prior is fixed to zero.
        """

        return TransformationPrior(
            UniformRotation2DPrior(), FixedTranslationPrior(np.zeros(2)))



cdef class RotationPrior(Distribution):
    """Distribution over Rotation3D object."""

    def log_likelihood(self, rotation):
        """Returns a constant, set to 0."""

        return 0.0


cdef class UniformRotationPrior(RotationPrior):
    """Uniform distribution over 3D rotations."""

    def sample(self):
        """Returns a uniform 3D rotation.

        Returns:
            Rotation3D object.
        """

        return Rotation3D.random()


cdef class UniformRotation2DPrior(RotationPrior):
    """Uniform distribution over 21D rotations."""

    def sample(self):
        """Returns a uniform 2D rotation.

        Returns:
            Rotation2D object.
        """

        return Rotation2D.random()


cdef class FixedRotationPrior(RotationPrior):
    """Dirac delta distribution over 3D rotations."""

    def __cinit__(self, rotation):
        """
        Args:
            rotation: Rotation3D object.
        """

        self._rotation = rotation

    def sample(self):
        """Returns the same rotation.

        Returns:
            Rotation3D object.
        """

        return self._rotation


cdef class TranslationPrior(Distribution):
    """Distribution over 3D translations or 2D translations."""

    def log_likelihood(self, translation):
        """Returns a constant, set to 0."""

        return 0.0


cdef class FixedTranslationPrior(TranslationPrior):
    """Dirac delta distribution over R^3 or R^2."""

    def __cinit__(self, translation):
        """
        Args:
            translation: double[:]
        """

        self._translation = translation

    def sample(self):
        """Returns the same translation."""

        return self._translation


class TranslationPriors(object):
    """Factory for generating TranslationPrior object."""

    @classmethod
    def default(cls):
        """Returns translation prior fixed to the origin."""

        translation = np.zeros(3, dtype='float')

        return FixedTranslationPrior(translation)

    @classmethod
    def normal(cls, variance):
        """Return Gaussian distribution as translation prior."""

        return IsoGaussianND(np.zeros(3, dtype='float'), variance)


cdef class IsoSubstructurePrior(Distribution):

    def __cinit__(self, IsoMixture3D mixture_fixed,
                  IsoMixturePrior3D mixture_prior, double ratio_prior,
                  TransformationPrior transformation_prior):
        """Initialise with prior on variable mixture, and substructure ratio.

        Args:
            mixture_fixed: The fixed substructure.
            mixture_prior: Prior on variable mixture.
            ratio_prior: Single double, alpha parameter of Beta prior.
                beta parameter is fixed by mixture weight prior.
        """

        self._mixture_fixed = mixture_fixed
        self._mixture_prior = mixture_prior
        self._ratio_prior = ratio_prior
        self._transformation_prior = transformation_prior

    @property
    def mixture_prior(self):
        return self._mixture_prior

    @property
    def ratio_prior(self):
        return self._ratio_prior

    def sample(self):
        """Draw a single sample from the prior."""

        mixture = self._mixture_prior.sample()
        ratio_alpha = self._ratio_prior
        alphas = np.array(self._mixture_prior._weights_prior.alpha)
        ratio_beta = (alphas - 1.0).sum() + 1
        ratio = np.random.beta(ratio_alpha, ratio_beta)
        substructure_transformation = self._transformation_prior.sample()

        return IsoSubstructure(self._mixture_fixed, mixture, ratio,
                               substructure_transformation)


