"""Isotropic mixture model together with transformations.

This describes the model parameters needed to obtain the 2D marginals.
"""

import numpy as np
cimport numpy as np
cimport cython

from libc.math cimport exp, log, sqrt
from libc.math cimport M_PI

cdef inline int int_max(int a, int b): return a if a >= b else b
cdef inline int int_min(int a, int b): return a if a <= b else b


from cryoem.model.isomixture cimport IsoMixture3D
from cryoem.model.transformation cimport Transformation3D
from cryoem.model.math cimport exp_fast, exp_fast2


cdef class IsoTransform:
    """Isotropic mixture model together with list of tranformations."""

    def __cinit__(self, mixture, transformations):
        """Initialise with iso mixture and list of transformations.

        Args:
            mixture: IsoMixture object, 3D mixture.
            tranformations: List of Transformation objects.
        """

        self._mixture = mixture
        self._transformations = transformations

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

    def set_mixture(self, value):
        self._mixture = value

    @property
    def transformations(self):
        return self._transformations

    def set_transformations(self, value):
        self._transformations = value

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

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

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

    def project(self):
        """Applies the various transformations to the mixture and project.

        Returns:
            List of IsoMixture2D objects.
        """

        return [self._mixture.transform(t).project()
                for t in self._transformations]

    def save(self, parent_group, name='isotransform'):
        """Save to group in HDF5 database.

        Assume that the group parent_group[name] doesn't exist.

        Args:
            parent_group: Group within which to create the isotransform
                group.
            name: Name of the new isotransform group.
        """

        group = parent_group.create_group(name)
        self.mixture.save(group)
        group.attrs['num_transformations'] = len(self.transformations)
        for i, t in enumerate(self.transformations):
            t.save(group, 'transformation_%d' % i)

    @classmethod
    def load(cls, parent_group, name='isotransform'):
        """Loads IsoTransform from group in HDF5 database."""

        group = parent_group[name]
        mixture = IsoMixture3D.load(group)
        P = group.attrs['num_transformations']
        transformations = [
            Transformation3D.load(group, 'transformation_%d' % i)
            for i in range(P)]

        return cls(mixture, transformations)

    def log_likelihoods(self, data, std_range=5.0):
        """Returns log-likelihood of projections.

        Args:
            data: List of 2D Griddata objects.
        """

        projections = self.project()
        lls = [p.log_likelihood_grid_sparse(d, std_range=std_range)
               for p, d in zip(projections, data)]

        return np.array(lls)

    def log_likelihood(self, data):
        """Returns log-likelihood of projections.

        Args:
            data: List of 2D Griddata objects.
        """

        return np.sum(self.log_likelihoods(data))

    def copy(self):
        """Returns a deep copy of this model."""

        return IsoTransform(self._mixture.copy(),
                            [t.copy() for t in self._transformations])

    def expand(self, symmetries):
        """Return model with more mixture components."""

        model_expanded = self.copy()
        model_expanded.set_mixture(model_expanded.mixture.expand(symmetries))

        return model_expanded

    def save_rotation_cylinders(self, filename, radius, length,
                                colours, model_name='rotations'):
        """Save rotations as cylinders in cmm format for chimera."""

        from cryoem.data.chimera_utils import MarkerSet
        from numpy.linalg.linalg import norm

        mset = MarkerSet(model_name)

        a = np.array([1.0, 0.0, 0.0]) # x-axis

        for i, t in enumerate(self._transformations):
            r = t.rotation
            p = r.inverse_rotate(a).flatten()
            p = length * p / norm(p)
            mset.add_cylinder([0, 0, 0], p, radius, colours[i])

        mset.write(filename)


cdef class IsoTransformNoise:
    """General isotropic mixture model, transformations and Gaussian noise.

    'General' refers to the weights being unnormalised.
    """

    def __cinit__(self, mixture, transformations, precisions):
        """Initialise with iso mixture and list of transformations.

        Args:
            mixture: IsoMixture3D object, 3D mixture.
            transformations: List of Transformation objects.
            precisions: double[] of same length as transformations.
                Precision of Gaussian noise for each projection direction.
        """

        self._mixture = mixture
        self._transformations = transformations
        self._precisions = precisions

        self._num_projections = len(transformations)
        self._num_components = self._mixture._num_components

        # For caching during energy computations:
        self._grid_size = 0

        # If true, remove the normalisation constant from
        # the Gaussian components.
        self._unnormalised = 0

    def set_unnormalised(self, value=True):
        """ Remove normalisation constant from Gaussian components."""
        if value:
            self._unnormalised = 1
        else:
            self._unnormalised = 0

    @property
    def unnormalised(self):
        return self._unnormalised

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

    def set_mixture(self, value):
        self._mixture = value

    @property
    def transformations(self):
        return self._transformations

    def set_transformations(self, value):
        self._transformations = value

    @property
    def precisions(self):
        return np.array(self._precisions)

    cpdef set_all_precisions(self, float value):
        """Set all precisions to the same value."""
        cdef int i
        for i in range(self._num_projections):
            self._precisions[i] = value

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

    @property
    def num_projections(self):
        return self._num_projections

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

    cdef init_memory(self, data):
        """Assign memory needed to compute energy and energy gradients.

        Args:
            data: List of GridDataReal objects.
        """

        print 'Initialise memory for IsoTransformNoise energy calculations.'

        P = len(data)

        self._grid = data[0].grid
        self._grid_size = self._grid.size

        K = self._mixture._num_components
        self._transformed = IsoMixture3D(
            np.zeros([K, 3], dtype='float'),
            1.0, np.zeros(K, dtype='float'))
        self._projected = IsoMixture2D(
            np.zeros([K, 2], dtype='float'),
            1.0, np.zeros(K, dtype='float'))
        self._projected.set_use_fast_exp(False)

        N = self._grid_size

        # Add GridDataReal of right size, filled with zeros.
        values = np.zeros(N, dtype='float')
        self._griddata = GridDataReal(self._grid, values)

        # # Compute inverse transformations.
        # self._inverse_transformations = [
        #     t.invert() for t in self.transformations]

        # 3D points for applying inverse transformations.
        self._point = np.zeros([1, 3], dtype='float')
        self._point2 = np.zeros([1, 3], dtype='float')
        self._point3 = np.zeros(3, dtype='float')

        # Add memory for computing second order partial derivatives.
        self._weight_accum = np.empty((P, N), dtype='float')
        self._precision_accum = np.empty((P, N), dtype='float')
        self._residuals2 = np.empty((N, ), dtype='float')
        self._residuals = np.empty((N, ), dtype='float')
        self._predicted = np.empty((N, ), dtype='float')        

    def project(self):
        """Applies the various transformations to the mixture and project.

        Returns:
            List of IsoMixture2D objects.
        """

        return [self._mixture.transform(t).project()
                for t in self._transformations]

    def project_grid(self, grid, std_range=5.0):
        """Project and evaluate on 2D grid."""

        projections = self.project()
        data = [m.evaluate_grid_sparse(grid, std_range)
                for m in projections]

        return data

    def sample(self, grid, std_range=5.0):
        """Samples list of GridDataReal objects from this model.

        Args:
            grid: Grid2D object on which to evaluate projections.

        Projects model in various directions and adds Gaussian noise.

        TODO: perhaps add shifting and scaling.
        """

        projections = self.project()

        data = [m.evaluate_grid_sparse(grid, std_range)
                for m in projections]

        n = grid.size
        noisy_data = [
            GridDataReal(grid, gd.values + np.random.randn(n) /
                         np.sqrt(precision))
            for gd, precision in zip(data, np.array(self._precisions))]

        return noisy_data

    cdef residual(self, double[:] a, double[:] b):
        """Return sum of squared errors."""

        cdef:
            int i
            double s

        s = 0.0
        for i in range(a.shape[0]):
            s += (a[i] - b[i])**2
        return s

    cpdef double energy_single(self, data,
                               int projection_index, double std_range):
        """Return negative log-likelihood of single image given model.

        The noise precision is ignored (it's included later).
        """

        cdef:
            int i
            GridDataReal gd
            Transformation transformation
            IsoMixture3D mixture
            double[:] observed, predicted
            double energy, residual

        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)


        i = projection_index

        transformation = self._transformations[i]
        self._mixture._transform(transformation, self._transformed)
        self._transformed._project(self._projected)

        self._projected.evaluate_grid_in_place(
            self._grid, std_range, self._griddata, self._unnormalised)

        # Compute residual (sum of squared errors).
        gd = data[i]
        observed = gd._values
        predicted = self._griddata._values
        residual = self.residual(observed, predicted)

        return 0.5 * residual

    def log_likelihood(self, data, std_range=5.0):
        """Return log-likelihood of data.

        Args:
            data: List of GridDataReal objects.

        """

        return -self.energy(data, std_range)

    cpdef energy(self, data, double std_range):
        """Return negative log-likelihood of given data.

        Energy is the negative log-likelihood of self._data given this
        model. Project mixture in all projection directions, evaluate on
        2D grids, compute residuals, and use noise precisions.

        Memory needed to compute result is assigned the first
        time method is called, and every time the grid changes.

        25.9.2014
        Adding term which depends only on noise precision.
        """

        cdef:
            int i
            double energy, residual, sum_log_betas, t2
            double[:] observed, predicted
            GridDataReal gd
            Transformation transformation
            IsoMixture3D mixture


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        energy = 0.0
        sum_log_betas = 0.0

        for i in range(self._num_projections):
            energy += self._precisions[i] * self.energy_single(
                data, i, std_range)
            sum_log_betas += log(self._precisions[i])

        t2 = -self._grid_size * sum_log_betas

        return energy + t2


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

        Note: Translations are assumed to be zero here.

        Args:
            data: List of GridDataReal objects.
            gradients: (3*K,) array, will contain calculated gradients.
        """

        cdef:
            int i, j, k, l, a0, b0, a1, b1, mean_index_0, mean_index_1
            int index_0, index_1, local_steps_0, local_steps_1
            double c1, c2, mean_0, mean_1, lambda_k, beta_i,
            double component_precision, local_range, p0, p1
            double[:] observed, residuals
            double[::1] p3
            double[:,:] p, p2, m
            GridDataReal gd
            Grid2D grid
            Transformation transformation, inverse_transformation
            IsoMixture3D mixture


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        grid = self._grid
        p = self._point
        p2 = self._point2
        p3 = self._point3

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        # Initialise gradients.
        for k in range(self._num_components):
            for j in range(3):
                gradients[3*k + j] = 0.0

        for i in range(self._num_projections):
            transformation = self._transformations[i]
            self._mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            self._projected.evaluate_grid_in_place(
                self._grid, std_range, self._griddata,
                self._unnormalised)

            # Compute residuals, replacing predicted grid values.
            gd = data[i]
            observed = gd._values
            residuals = self._griddata._values # predicted
            for j in range(self._grid_size):
                residuals[j] = observed[j] - residuals[j]

            # Evaluate part of gradients for this projection direction.
            inverse_transformation = transformation.invert()
            m = inverse_transformation._rotation._matrix

            c1 = 1.0 / (2.0 * M_PI * self._projected._variance)
            if self._unnormalised == 1:
                c1 = 1.0
            c2 = -0.5 / self._projected._variance
            beta_i = self._precisions[i]
            component_precision = 1.0 / self._mixture._variance
            for k in range(self._num_components):

                mean_0 = self._projected._means[k, 0]
                mean_1 = self._projected._means[k, 1]
                lambda_k = self._mixture._weights[k]

                # Compute indices of grid point near component mean.
                mean_index_0 = (int)((mean_0 - grid._origin_0)
                                      / grid._spacing_0)
                mean_index_1 = (int)((mean_1 - grid._origin_1)
                                      / grid._spacing_1)
                # Compute the indices of the grid points to consider.
                a0 = int_max(mean_index_0 - local_steps_0, 0)
                b0 = int_min(mean_index_0 + local_steps_0 + 2,
                             grid._shape_0)
                for index_0 in range(a0, b0):
                    a1 = int_max(mean_index_1 - local_steps_1, 0)
                    b1 = int_min(mean_index_1 + local_steps_1 + 2,
                                 grid._shape_1)
                    for index_1 in range(a1, b1):
                        # Evaluate component at grid point.
                        j = index_0 * grid._shape_1 + index_1
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index_0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index_1)
                        p0 = mean_0 - grid_point_0
                        p1 = mean_1 - grid_point_1
                        dist2 = p0*p0 + p1*p1
                        ## value = c1 * exp_fast2(c2 * dist2)
                        value = c1 * exp(c2 * dist2)

                        p3[0] = m[0, 0] * p0 + m[0, 1] * p1
                        p3[1] = m[1, 0] * p0 + m[1, 1] * p1
                        p3[2] = m[2, 0] * p0 + m[2, 1] * p1

                        # Increment gradient.
                        for l in range(3):
                            gradients[3*k + l] += (
                                residuals[j] * value * p3[l] *
                                beta_i * lambda_k *
                                component_precision)


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

        Args:
            data: List of GridDataReal objects.
            gradients: (K,) array, will contain calculated gradients.

        """

        cdef:
            int i, j, k, l, a0, b0, a1, b1, mean_index_0, mean_index_1
            int index_0, index_1, local_steps_0, local_steps_1
            double c1, c2, mean_0, mean_1, beta_i
            double component_precision, local_range, p0, p1
            double[:] observed, residuals
            GridDataReal gd
            Grid2D grid
            Transformation transformation
            IsoMixture3D mixture


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        grid = self._grid

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        # Initialise gradients.
        for k in range(self._num_components):
            gradients[k] = 0.0

        for i in range(self._num_projections):
            transformation = self._transformations[i]
            self._mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            self._projected.evaluate_grid_in_place(
                self._grid, std_range, self._griddata, 0)

            # Compute residuals, replacing predicted grid values.
            gd = data[i]
            observed = gd._values
            residuals = self._griddata._values # predicted
            for j in range(self._grid_size):
                residuals[j] = observed[j] - residuals[j]

            # Evaluate part of gradients for this projection direction.
            c1 = 1.0 / (2.0 * M_PI * self._projected._variance)
            c2 = -0.5 / self._projected._variance
            beta_i = self._precisions[i]
            component_precision = 1.0 / self._mixture._variance
            for k in range(self._num_components):

                mean_0 = self._projected._means[k, 0]
                mean_1 = self._projected._means[k, 1]

                # Compute indices of grid point near component mean.
                mean_index_0 = (int)((mean_0 - grid._origin_0)
                                      / grid._spacing_0)
                mean_index_1 = (int)((mean_1 - grid._origin_1)
                                      / grid._spacing_1)
                # Compute the indices of the grid points to consider.
                a0 = int_max(mean_index_0 - local_steps_0, 0)
                b0 = int_min(mean_index_0 + local_steps_0 + 2,
                             grid._shape_0)
                for index_0 in range(a0, b0):
                    a1 = int_max(mean_index_1 - local_steps_1, 0)
                    b1 = int_min(mean_index_1 + local_steps_1 + 2,
                                 grid._shape_1)
                    for index_1 in range(a1, b1):
                        # Evaluate component at grid point.
                        j = index_0 * grid._shape_1 + index_1
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index_0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index_1)
                        p0 = mean_0 - grid_point_0
                        p1 = mean_1 - grid_point_1
                        dist2 = p0*p0 + p1*p1
                        value = c1 * exp(c2 * dist2)

                        # Increment gradient.
                        gradients[k] -= residuals[j] * value * beta_i

    cpdef energy_gradient_weight(self, data, double std_range):
        """Return gradient of energy w.r.t. single common weight.

        Args:
            data: List of GridDataReal objects.
        """

        cdef:
            int i, j, k, l, a0, b0, a1, b1, mean_index_0, mean_index_1
            int index_0, index_1, local_steps_0, local_steps_1
            double c1, c2, mean_0, mean_1, beta_i
            double component_precision, local_range, p0, p1, gradient
            double[:] observed, residuals
            GridDataReal gd
            Grid2D grid
            Transformation transformation
            IsoMixture3D mixture


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        grid = self._grid

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        # Initialise gradient.
        gradient = 0.0

        for i in range(self._num_projections):
            transformation = self._transformations[i]
            self._mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            self._projected.evaluate_grid_in_place(
                self._grid, std_range, self._griddata,
                self._unnormalised)

            # Compute residuals, replacing predicted grid values.
            gd = data[i]
            observed = gd._values
            residuals = self._griddata._values # predicted
            for j in range(self._grid_size):
                residuals[j] = observed[j] - residuals[j]

            # Evaluate part of gradients for this projection direction.
            c1 = 1.0 / (2.0 * M_PI * self._projected._variance)
            if self._unnormalised == 1:
                c1 = 1.0
            c2 = -0.5 / self._projected._variance
            beta_i = self._precisions[i]
            ## component_precision = 1.0 / self._mixture._variance
            for k in range(self._num_components):

                mean_0 = self._projected._means[k, 0]
                mean_1 = self._projected._means[k, 1]

                # Compute indices of grid point near component mean.
                mean_index_0 = (int)((mean_0 - grid._origin_0)
                                      / grid._spacing_0)
                mean_index_1 = (int)((mean_1 - grid._origin_1)
                                      / grid._spacing_1)
                # Compute the indices of the grid points to consider.
                a0 = int_max(mean_index_0 - local_steps_0, 0)
                b0 = int_min(mean_index_0 + local_steps_0 + 2,
                             grid._shape_0)
                for index_0 in range(a0, b0):
                    a1 = int_max(mean_index_1 - local_steps_1, 0)
                    b1 = int_min(mean_index_1 + local_steps_1 + 2,
                                 grid._shape_1)
                    for index_1 in range(a1, b1):
                        # Evaluate component at grid point.
                        j = index_0 * grid._shape_1 + index_1
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index_0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index_1)
                        p0 = mean_0 - grid_point_0
                        p1 = mean_1 - grid_point_1
                        dist2 = p0*p0 + p1*p1
                        value = c1 * exp(c2 * dist2)

                        # Increment gradient.
                        gradient -= residuals[j] * value * beta_i

        return gradient


    cpdef energy_gradient_component_precision(self, data, double std_range):
        """Return gradient of energy w.r.t. the component precision.

        Args:
            data: List of GridDataReal objects.
        """

        cdef:
            int i, j, k, l, a0, b0, a1, b1, mean_index_0, mean_index_1
            int index_0, index_1, local_steps_0, local_steps_1
            double c1, c2, mean_0, mean_1, beta_i, lambda_k
            double component_precision, local_range, p0, p1, gradient
            double[:] observed, residuals
            GridDataReal gd
            Grid2D grid
            Transformation transformation
            IsoMixture3D mixture


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        grid = self._grid

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        # Initialise gradient.
        gradient = 0.0

        for i in range(self._num_projections):
            transformation = self._transformations[i]
            self._mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            self._projected.evaluate_grid_in_place(
                self._grid, std_range, self._griddata,
                self._unnormalised)

            # Compute residuals, replacing predicted grid values.
            gd = data[i]
            observed = gd._values
            residuals = self._griddata._values # predicted
            for j in range(self._grid_size):
                residuals[j] = observed[j] - residuals[j]

            # Evaluate part of gradients for this projection direction.
            c1 = 1.0 / (2.0 * M_PI * self._projected._variance)
            if self._unnormalised == 1:
                c1 = 1.0
            c2 = -0.5 / self._projected._variance
            beta_i = self._precisions[i]
            ## component_precision = 1.0 / self._mixture._variance
            for k in range(self._num_components):

                mean_0 = self._projected._means[k, 0]
                mean_1 = self._projected._means[k, 1]
                lambda_k = self._mixture._weights[k]

                # Compute indices of grid point near component mean.
                mean_index_0 = (int)((mean_0 - grid._origin_0)
                                      / grid._spacing_0)
                mean_index_1 = (int)((mean_1 - grid._origin_1)
                                      / grid._spacing_1)
                # Compute the indices of the grid points to consider.
                a0 = int_max(mean_index_0 - local_steps_0, 0)
                b0 = int_min(mean_index_0 + local_steps_0 + 2,
                             grid._shape_0)
                for index_0 in range(a0, b0):
                    a1 = int_max(mean_index_1 - local_steps_1, 0)
                    b1 = int_min(mean_index_1 + local_steps_1 + 2,
                                 grid._shape_1)
                    for index_1 in range(a1, b1):
                        # Evaluate component at grid point.
                        j = index_0 * grid._shape_1 + index_1
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index_0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index_1)
                        p0 = mean_0 - grid_point_0
                        p1 = mean_1 - grid_point_1
                        dist2 = p0*p0 + p1*p1
                        value = c1 * exp(c2 * dist2) * dist2

                        # Increment gradient.
                        gradient += (
                            residuals[j] * value * beta_i * 0.5*lambda_k)

        return gradient

    cpdef energy_gradient_noise_precision(self, data, double std_range):
        """Return gradient of energy w.r.t. the noise precision.

        Assume that all noise precisions are the same.

        Args:
            data: List of GridDataReal objects.
        """

        cdef:
            int i
            double gamma_a, gamma_b, gradient, beta
            GridDataReal gd


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        gamma_b = 0.0

        for i in range(self._num_projections):
            gamma_b += self.energy_single(data, i, std_range)

        gamma_a = self._num_projections * self._grid_size

        beta = self._precisions[0]
        gradient = -(gamma_a - 1.0) / beta + gamma_b

        return gradient

    cpdef energy_gradient_weight_precision(
        self, data,
        double[:] gradients,
        double[:,:] gradients2,
        double std_range):
        """Return gradients of energy w.r.t. the weight and precision.

        Both first and second gradients are component.
        W.r.t. a single weight and the component precision.

        Individual Gaussians are assumed to be unnormalised.

        Args:
            data: List of GridDataReal objects.
            gradients: (2, ) array with weight and precision gradient.
            gradients2: Symmetric (2, 2) matrix with Hessian.
        """

        cdef:
            int i, j, k, l, a0, b0, a1, b1, mean_index_0, mean_index_1
            int index_0, index_1, local_steps_0, local_steps_1
            int P, K, N
            double c1, c2, mean_0, mean_1, beta_i, lambda_k
            double beta, lambda_0
            double component_precision, local_range, p0, p1, gradient
            double weight_sum_squares, precision_sum_squares
            double gradient_weight, gradient_precision
            double gradient2_precision, gradient2_mixed
            double[:] observed, residuals
            GridDataReal gd
            Grid2D grid
            Transformation transformation
            IsoMixture3D mixture


        P = self._num_projections
        K = self._num_components
        N = self._grid_size

        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        grid = self._grid

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        # Initialise accumulators.
        gradient_weight = 0.0
        gradient_precision = 0.0
        gradient2_precision = 0.0 # 2nd term of 2nd derivative
        gradient2_mixed = 0.0
        for i in range(P):
            for j in range(N):
                self._weight_accum[i, j] = 0.0
                self._precision_accum[i, j] = 0.0 # 1st term of 2nd derivative

        # Evaluate each component at nearby pixels.
        for i in range(self._num_projections):
            transformation = self._transformations[i]
            self._mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            self._projected.evaluate_grid_in_place(
                self._grid, std_range, self._griddata,
                self._unnormalised)

            # Compute residuals, replacing predicted grid values.
            gd = data[i]
            observed = gd._values
            residuals = self._griddata._values # predicted
            # Note that we reuse the memory where the predicted
            # values are stored for storing the residuals.
            for j in range(self._grid_size):
                self._residuals2[j] = observed[j] - 2.0 * residuals[j]
                residuals[j] = observed[j] - residuals[j]

            # Evaluate part of gradients for this projection direction.
            c2 = -0.5 / self._projected._variance
            for k in range(self._num_components):

                mean_0 = self._projected._means[k, 0]
                mean_1 = self._projected._means[k, 1]
                lambda_k = self._mixture._weights[k]

                # Compute indices of grid point near component mean.
                mean_index_0 = (int)((mean_0 - grid._origin_0)
                                      / grid._spacing_0)
                mean_index_1 = (int)((mean_1 - grid._origin_1)
                                      / grid._spacing_1)
                # Compute the indices of the grid points to consider.
                a0 = int_max(mean_index_0 - local_steps_0, 0)
                b0 = int_min(mean_index_0 + local_steps_0 + 2,
                             grid._shape_0)
                for index_0 in range(a0, b0):
                    a1 = int_max(mean_index_1 - local_steps_1, 0)
                    b1 = int_min(mean_index_1 + local_steps_1 + 2,
                                 grid._shape_1)
                    for index_1 in range(a1, b1):
                        # Evaluate component at grid point.
                        j = index_0 * grid._shape_1 + index_1
                        grid_point_0 = (grid._origin_0 +
                                        grid._spacing_0 * index_0)
                        grid_point_1 = (grid._origin_1 +
                                        grid._spacing_1 * index_1)
                        p0 = mean_0 - grid_point_0
                        p1 = mean_1 - grid_point_1
                        dist2 = p0*p0 + p1*p1
                        v = exp(c2 * dist2)

                        # Adjust accummulators.
                        gradient_weight += residuals[j] * v
                        gradient_precision += residuals[j] * v * dist2
                        gradient2_precision -= (
                            residuals[j] * v * dist2 * dist2)
                        gradient2_mixed += self._residuals2[j] * v * dist2
                        self._weight_accum[i, j] += v
                        self._precision_accum[i, j] += v * dist2

        # Compute gradients.
        beta = self._precisions[0] # Assume they're all the same.
        lambda_0 = self._mixture._weights[0] # Again all the same.
        # Weight, first derivative.
        gradients[0] = -beta * gradient_weight
        # Component precision, first derivative.
        gradients[1] = 0.5 * beta * lambda_0 * gradient_precision
        # Weight and precision, second derivative.
        weight_sum_squares = 0.0
        precision_sum_squares = 0.0
        for i in range(P):
            for j in range(N):
                weight_sum_squares += self._weight_accum[i, j]**2
                precision_sum_squares += self._precision_accum[i, j]**2
        gradients2[0, 0] = beta * weight_sum_squares
        gradients2[1, 1] = 0.25 * beta * lambda_0 * (
            lambda_0 * precision_sum_squares + gradient2_precision)
        gradients2[0, 1] = 0.5 * beta * gradient2_mixed
        gradients2[1, 0] = gradients2[0, 1]


    cpdef optimal_weight(self, data, double std_range):
        """Return optimal weight.

        Assumes single weight for all components,
        and single noise precision for all images.
        """

        cdef:
            int i, N
            double energy, residual, num, den, weight
            GridDataReal gd
            Transformation transformation
            IsoMixture3D mixture
            double[:] observed, predicted


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)
        N = self._grid_size

        num = 0.0
        den = 0.0

        weight = self._mixture._weights[0]

        for i in range(self._num_projections):

            transformation = self._transformations[i]
            self._mixture._transform(transformation, self._transformed)
            self._transformed._project(self._projected)

            self._projected.evaluate_grid_in_place(
                self._grid, std_range, self._griddata,
                self._unnormalised)
            predicted = self._griddata._values

            gd = data[i]
            observed = gd._values

            for l in range(N):
                p = (predicted[l] / weight)
                num += observed[l] * p
                den += p*p

        return num / den

    cpdef energy_gradient_rotation(
        self, data, int projection_index, double[:, :] gradient,
        double std_range):
        """Compute 3x3 gradients of energy w.r.t. rotation matrix.

        Args:
            data: List of GridDataReal objects.
            gradient: (3, 3) array, will contain calculated gradients.
        """

        cdef:
            int i, j, k, l, a0, b0, a1, b1, mean_index_0, mean_index_1
            int index_0, index_1, local_steps_0, local_steps_1
            int ri, rj
            double c1, c2, mean_0, mean_1, lambda_k, beta_i, c3
            double component_precision, local_range, p0, p1
            double[:] observed, residuals, p, u
            double[::1] p3
            # double[:,:] p, p2, m
            GridDataReal gd
            Grid2D grid
            Transformation transformation, inverse_transformation
            IsoMixture3D mixture


        gd = data[0]
        if self._grid_size != gd._grid._size:
            self.init_memory(data)

        p = np.zeros(2, dtype='float')
        u = np.zeros(3, dtype='float')

        grid = self._grid

        # Calculate the number of steps to move in each direction.
        local_range = sqrt(self._mixture._variance) * std_range
        local_steps_0 = (int)(local_range / grid._spacing_0)
        local_steps_1 = (int)(local_range / grid._spacing_1)

        # Initialise gradient.
        for ri in range(3):
            for rj in range(3):
                gradient[ri, rj] = 0.0

        i = projection_index

        transformation = self._transformations[i]
        self._mixture._transform(transformation, self._transformed)
        self._transformed._project(self._projected)

        self._projected.evaluate_grid_in_place(
            self._grid, std_range, self._griddata,
            self._unnormalised)

        # Compute residuals, replacing predicted grid values.
        gd = data[i]
        observed = gd._values
        self._predicted = self._griddata._values
        for j in range(self._grid_size):
            self._residuals[j] = observed[j] - self._predicted[j]

        c1 = 1.0 / (2.0 * M_PI * self._projected._variance)
        if self._unnormalised == 1:
            c1 = 1.0
        c2 = -0.5 / self._projected._variance
        beta_i = self._precisions[i]
        component_precision = 1.0 / self._mixture._variance
        for k in range(self._num_components):

            mean_0 = self._projected._means[k, 0]
            mean_1 = self._projected._means[k, 1]
            lambda_k = self._mixture._weights[k]

            # Compute indices of grid point near component mean.
            mean_index_0 = (int)((mean_0 - grid._origin_0)
                                  / grid._spacing_0)
            mean_index_1 = (int)((mean_1 - grid._origin_1)
                                  / grid._spacing_1)
            # Compute the indices of the grid points to consider.
            a0 = int_max(mean_index_0 - local_steps_0, 0)
            b0 = int_min(mean_index_0 + local_steps_0 + 2,
                         grid._shape_0)
            for index_0 in range(a0, b0):
                a1 = int_max(mean_index_1 - local_steps_1, 0)
                b1 = int_min(mean_index_1 + local_steps_1 + 2,
                             grid._shape_1)
                for index_1 in range(a1, b1):
                    # Evaluate component at grid point.
                    j = index_0 * grid._shape_1 + index_1
                    grid_point_0 = (grid._origin_0 +
                                    grid._spacing_0 * index_0)
                    grid_point_1 = (grid._origin_1 +
                                    grid._spacing_1 * index_1)

                    p[0] = mean_0 - grid_point_0
                    p[1] = mean_1 - grid_point_1
                    dist2 = p[0]**2 + p[1]**2
                    v = exp(c2 * dist2)

                    for rj in range(3):
                        u[rj] = self._mixture._means[k, rj]

                    # Increment gradient.
                    # Should have a minus sign here.
                    c3 = (beta_i * component_precision * lambda_k *
                          self._residuals[j] * v)
                    for ri in range(2):
                        for rj in range(3):
                            gradient[ri, rj] += c3 * p[ri] * u[rj]

    def save(self, parent_group, name='isotransform'):
        """Save to group in HDF5 database.

        Assume that the group parent_group[name] doesn't exist.

        Args:
            parent_group: Group within which to create the isotransform
                group.
            name: Name of the new isotransform group.
        """

        group = parent_group.create_group(name)
        self.mixture.save(group)
        group.attrs['num_transformations'] = len(self.transformations)
        for i, t in enumerate(self.transformations):
            t.save(group, 'transformation_%d' % i)
        precisions = np.array(self.precisions)
        dataset = group.create_dataset('precisions', precisions.shape,
                                       np.double)
        dataset[...] = precisions


    @classmethod
    def load(cls, parent_group, name='isotransform'):
        """Loads IsoTransform from group in HDF5 database."""

        group = parent_group[name]
        mixture = IsoMixture3D.load(group)
        P = group.attrs['num_transformations']
        transformations = [
            Transformation3D.load(group, 'transformation_%d' % i)
            for i in range(P)]
        precisions = group['precisions'][:]

        return cls(mixture, transformations, precisions)

    def copy(self):
        """Returns a deep copy of this model."""

        return IsoTransformNoise(self._mixture.copy(),
                                 [t.copy() for t in self._transformations],
                                 np.array(self._precisions).copy())



