"""
Sample or optimise rotations globally using either brute force or
perhaps later branch and bound. Mixture is assumed fixed.
"""

import numpy as np
cimport numpy as np
from numpy.random import multinomial

cimport cython

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

from cryoem.model.rotation cimport Rotation, Rotation3D, UniformRotations
from cryoem.model.math cimport log_sum_exp_c as log_sum_exp
from cryoem.data.griddata cimport GridDataNestedList, GridDataInt
from cryoem.data.grid cimport Grid2D
from cryoem.model.isotransform cimport IsoTransform, IsoTransformNoise
from cryoem.sampling.dos cimport States, StatesGridEnergies
from cryoem.sampling.dos cimport StatesGridEnergiesGaussian
from cryoem.model.transformation cimport Transformation, Transformation3D
from cryoem.model.isomixture cimport IsoMixture3D


cdef class BruteForceRotationSampler:
    """Sample or optimise rotations given mixture.

    Globally, using brute force.

    Later included translations. Only rotations are sampled, but
    every time the mixture as well as the translations might be different.
    """

    def __cinit__(self, GridDataNestedList data,
                  int sample_size):
        """
        Args:
            data: GridDataNestedList object. Note that a
                GridDataList object can easily be converted to this.
            sample_size: Number of random rotations to use.

        Note: std_range is just assumed to be 5. Could later
        include the option to vary this.
        """

        self._data = data

        self._num_projections = data.get_num_projections()

        self._counts_beta = 1.0 # For extracting the correct
            # GridDataInt object from 'data'.

        # Sample random rotations. This is only done once, right
        # at the beginning, and re-used throughout.
        self._rotations = UniformRotations(sample_size).generate_random()
        self._sample_size = sample_size

        self._use_symmetries = 0

    cpdef set_symmetries(self, list symmetries):
        """Use symmetry group.

        This only means that the mixture must first be expanded.
        """

        self._use_symmetries = 1
        self._symmetries = symmetries

    cpdef initialise_rotations(self, int sample_size):
        """Create uniformly random rotations.

        Returns:
            (sample_size, 9) double array, each row is a rotation.

        """

        cdef:
            int i, j
            double[:, ::1] rotations

        rotations = np.empty((sample_size, 9), dtype='double')
        for i in range(sample_size):
            r = np.array(Rotation3D.random().matrix).flatten()
            for j in range(9):
                rotations[i, j] = r[j]

        return rotations

    cpdef update_rotations(self, IsoTransform model, select_index):
        """Choose each rotation in the model given the current mixture.

        Args:
            select_index: Function from double[:] to int. Either sample
                index using weights, or select the index with the largest
                weight.
        """

        cdef:
            States states
            double[::1] energies = np.empty((self._sample_size, ),
                                            dtype='double')
            double[::1] weights = np.empty((self._sample_size, ),
                                         dtype='double')
            Transformation t
            double[:, :] matrix
            int i, j, k, index
            StatesGridEnergies energy_calculator
            Grid2D grid
            GridDataInt griddata
            double[:] translation
            IsoMixture3D mixture

        if self._use_symmetries == 1:
            print 'Using symmetries in global rotation sampler.'
            mixture = model._mixture.expand(self._symmetries)
        else:
            mixture = model._mixture
        
        states = States(self._rotations, mixture)
        grid = self._data.get(0, 1.0).grid

        # This is where the mixture is projected and evaluated.
        energy_calculator = states.get_grid_energy_calculator(grid)

        for i in range(self._num_projections):
            griddata = self._data.get(i, self._counts_beta)
            # Shift the griddata according to the current translation.
            t = model._transformations[i]
            translation = t._translation
            energy_calculator.energy_vector_in_place(griddata, energies)
            self.compute_weights(energies, weights)
            # Select index either by sampling or optimising.
            index = select_index(energies, weights)
            # Update rotation with this one.

            matrix = t._rotation._matrix
            for j in range(3):
                for k in range(3):
                    matrix[j, k] = self._rotations[index, j*3 + k]

    cpdef int sample_index(self, double[::1] energies,
                           double[::1] weights):
        """Sample index from multinomial with given weights as parameter."""

        cdef:
            int index

        index = np.where(multinomial(1, np.array(weights)) == 1)[0][0]

        return index

    cpdef int optimise_index(self, double[::1] energies,
                             double[::1] weights):
        """Return index minimising the energy."""

        cdef:
            int index

        index = np.argmin(np.array(energies))
        ## print('Weight of chosen energy / largest weight = ' +
        ##       '%1.3f' % weights[index])

        return index

    cpdef sample(self, IsoTransform model):
        """Sample each rotation in the model given the current mixture."""

        return self.update_rotations(model, self.sample_index)

    cpdef optimise(self, IsoTransform model):
        """Optimise each rotation in the model given the current mixture."""

        return self.update_rotations(model, self.optimise_index)

    cdef compute_weights(self, double[::1] energies, double[::1] weights):
        """Compute the weights of the energy distribution.

        These are just proportional to exp(-E), and normalised.
        """

        cdef:
            int i, n
            double c

        n = energies.shape[0]

        for i in range(n):
            weights[i] = -energies[i]

        # Compute normalisation constant.
        c = log_sum_exp(weights)

        # Normalise weights.
        for i in range(n):
            weights[i] = exp(weights[i] - c)

        # Sometimes they are not perfectly normalised for numerical reasons.
        c = 0.0
        for i in range(n):
            c += weights[i]
        for i in range(n):
            weights[i] /= c






cdef class BruteForceGaussianRotationSampler:
    """Sample or optimise rotations given mixture.

    Assumed Gaussian noise model. I.e. the model should be
    IsoTransformNoise instead of IsoTransform.

    Globally, using brute force.
    """

    def __cinit__(self, list data, int sample_size,
                  int unnormalised):
        """
        Args:
            data: List of GridDataReal objects.
            sample_size: Number of random rotations to use.
            unnormalised: Coefficient of 2D Gaussian.

        Note: std_range is just assumed to be 5. Could later
        include the option to vary this.
        """

        self._data = data
        self._unnormalised = unnormalised

        self._num_projections = len(data)

        # Sample random rotations. This is only done once, right
        # at the beginning, and re-used throughout.
        self._rotations = UniformRotations(
            sample_size).generate_random()
        self._sample_size = sample_size

        self._use_symmetries = 0

        # For scaling the energy function.
        self._annealing_beta = 1.0


    def set_annealing_beta(self, value):
        self._annealing_beta = value

    cpdef set_symmetries(self, list symmetries):
        """Use symmetry group.

        This only means that the mixture must first be expanded.
        """

        self._use_symmetries = 1
        self._symmetries = symmetries

    cpdef update_rotations(
        self, IsoTransformNoise model, select_index):
        """Choose each rotation in the model given the current mixture.

        Args:
            select_index: Function from double[:] to int. Either sample
                index using weights, or select the index with the 
                largest weight.
        """

        cdef:
            double[::1] energies = np.empty((self._sample_size, ),
                                            dtype='double')
            double[::1] weights = np.empty((self._sample_size, ),
                                         dtype='double')
            Transformation t
            double[:, :] matrix
            int i, j, k, index
            StatesGridEnergiesGaussian energy_calculator
            Grid2D grid
            double[:] translation
            IsoMixture3D mixture

        if self._use_symmetries == 1:
            print 'Using symmetries in global rotation sampler.'
            mixture = model._mixture.expand(self._symmetries)
        else:
            mixture = model._mixture
        
        grid = self._data[0].grid

        # This is where the mixture is projected and evaluated.
        energy_calculator = StatesGridEnergiesGaussian(
            self._rotations, mixture, grid, self._unnormalised)

        for i in range(self._num_projections):
            t = model._transformations[i]
            energy_calculator.energy_vector_in_place(
                self._data[i], energies, model._precisions[i])
            self.compute_weights(energies, weights)
            # Select index either by sampling or optimising.
            index = select_index(energies, weights)
            # Update rotation with this one.
            matrix = t._rotation._matrix
            for j in range(3):
                for k in range(3):
                    matrix[j, k] = self._rotations[index, j*3 + k]

    cpdef int sample_index(self, double[::1] energies,
                           double[::1] weights):
        """Sample index from multinomial with weights as parameter."""

        cdef:
            int index

        # print 'weights:'
        # print np.array(weights)
        index = np.where(multinomial(1, np.array(weights)) == 1)[0][0]

        return index

    cpdef int optimise_index(self, double[::1] energies,
                             double[::1] weights):
        """Return index minimising the energy."""

        cdef:
            int index

        index = np.argmin(np.array(energies))
        ## print('Weight of chosen energy / largest weight = ' +
        ##       '%1.3f' % weights[index])

        return index

    cpdef sample(self, IsoTransformNoise model):
        """Sample each rotation given the current mixture."""

        return self.update_rotations(model, self.sample_index)

    cpdef optimise(self, IsoTransformNoise model):
        """Optimise each rotation given the current mixture."""

        return self.update_rotations(model, self.optimise_index)

    cdef compute_weights(self, double[::1] energies,
                         double[::1] weights):
        """Compute the weights of the energy distribution.

        These are just proportional to exp(-E), and normalised.
        """

        cdef:
            int i, n
            double c


        n = energies.shape[0]

        for i in range(n):
            weights[i] = -energies[i] * self._annealing_beta

        # Compute normalisation constant.
        c = log_sum_exp(weights)

        # Normalise weights.
        for i in range(n):
            weights[i] = exp(weights[i] - c)

        # print 'weights:'
        # print np.sort(np.array(weights))[::-1]

        # Sometimes they are not perfectly normalised
        # for numerical reasons.
        c = 0.0
        for i in range(n):
            c += weights[i]
        for i in range(n):
            weights[i] /= c

