import numpy as np
cimport numpy as np

from libc.math cimport exp, log, sqrt


from cryoem.model.isomixture cimport IsoMixture1D
from cryoem.data.griddata cimport GridDataReal
from cryoem.data.grid cimport Grid1D



cdef class ScaledIsoMixture1D:
    """Scaled 1D isotropic mixture model.

    Essentially just a 1D mixture model with a scaling factor and
    Gaussian noise added at the end.
    """

    def __cinit__(self, IsoMixture1D mixture, double scale,
                  double noise_variance):
        """
        Args:
            mixture: IsoMixture1D object, integrates to 1.
            scale: Double, scaling factor.
            noise_variance: Noise level.

        The forward model is (scale*mixture) + noise, where the
        noise is Gaussian with given variance and zero mean.
        """

        self._mixture = mixture
        self._scale = scale
        self._noise_variance = noise_variance

    @property
    def mixture(self):
        return self._mixture

    @property
    def scale(self):
        return self._scale

    @property
    def noise_variance(self):
        return self._noise_variance

    cpdef GridDataReal evaluate_grid(self, Grid1D grid):
        """Evaluate scaled mixture on grid without adding noise."""

        cdef:
            GridDataReal evaluated
            double[:] noise, y

        evaluated = self._mixture.evaluate_grid(grid)
        values = evaluated._values

        y = np.zeros(grid.shape, dtype='double')

        for i in range(len(y)):
            y[i] = values[i] * self._scale

        return GridDataReal(grid, y)


    cpdef GridDataReal sample_grid(self, Grid1D grid):
        """Evaluate scaled mixture on grid and add noise."""

        cdef:
            GridDataReal evaluated
            double[:] noise, y

        evaluated = self._mixture.evaluate_grid(grid)
        values = evaluated._values

        y = np.zeros(grid.shape, dtype='double')

        noise = np.random.randn(grid.shape) * sqrt(self._noise_variance)

        for i in range(len(y)):

            y[i] = values[i] * self._scale + noise[i]

        return GridDataReal(grid, y)

    cpdef copy(self):
        """Returns a deep copy of the scaled mixture."""

        return ScaledIsoMixture1D(self._mixture.copy(),
                                  self._scale,
                                  self._noise_variance)
