"""Density of states."""

import numpy as np
cimport numpy as np

from libc.math cimport exp, log, sqrt

cimport cython

import time

from numpy.random import multinomial

from cryoem.model.isotransform cimport IsoTransform
from cryoem.model.transformation cimport Transformation3D
from cryoem.model.isomixture cimport IsoMixture2D, IsoMixture3D
from cryoem.model.rotation cimport Rotation3D
from cryoem.data.griddata cimport GridDatasNested, GridDataNestedList
from cryoem.data.griddata cimport GridDataInt, GridDataReal
from cryoem.data.grid cimport Grid2D
from cryoem.model.math cimport log_sum_exp_c as log_sum_exp


@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class States:
    """Contains all states collected during sampling.

    Can evaluate the energy of each state for any given ensemble.

    To be used in calculating energies for DOS estimation,
    and also later for calculating extra energies needed for
    estimating the swap rate.
    """

    def __cinit__(self, double[:, ::1] rotations, IsoMixture3D mixture):
        """
        Args:
            rotations: (N, 9) array, each row a different rotation.
            mixture: Mixture for computing energies.
        """

        self._rotations = rotations
        self._mixture = mixture
        self._num_states = len(self._rotations)

    @property
    def num_states(self):
        return self._num_states

    @property
    def rotations(self):
        return self._rotations

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

    cpdef get_direct_energy_calculator(self):
        """Return object to calculate the energies of individual griddatas.

        Choose this if memory is an issue, counts are few, or the
        energies are needed for only a single griddata.

        The same amount of computation is done for every new image.
        """

        return StatesDirectEnergies(self)

    cpdef get_grid_energy_calculator(self, Grid2D grid):
        """Return object to calculate energies for many griddatas.

        Choose this if memory is not an issue (not too many states (< 1e6),
        and not too many pixels), or energies are needed for many images.

        The expensive computation is done right at the beginning, and
        obtaining the energies for new images is fast.
        """

        return StatesGridEnergies(self, grid)

    cpdef energy_at_beta_exact(self, double beta, GridDatasNested data,
                               double[::1] energy):
        """Computes energy vector without using interpolation.

        TODO: Do energy caching somewhere else. Get rid of this method.
        """

        cdef:
            int i, mi, mj
            double std_range = 5.0
            double ll
            Rotation3D rotation
            double[:, :] matrix
            GridDataInt griddata

        griddata = data.from_beta(beta)

        if self._state_energy is None:
            self._state_energy = StatesGridEnergies(self, griddata._grid)

        self._state_energy.energy_vector_in_place(griddata, energy)


    def save(self, parent_group, name='states'):
        """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 states group.
            name: Name of the new states group.
        """

        group = parent_group.create_group(name)
        dataset_rotations = group.create_dataset(
            'rotations', (<object>self._rotations).shape, np.double)
        dataset_rotations[...] = self._rotations
        self._mixture.save(group)

    @classmethod
    def load(cls, parent_group, name='states'):
        """Load States object from group in HFD5 database.
        """

        group = parent_group[name]
        rotations = group['rotations'][:]
        mixture = IsoMixture3D.load(group)

        states = cls(rotations, mixture)

        return states


@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class StatesDirectEnergies:
    """Calculation of state energy vector for images / replicas.

    Only evaluates projected mixture at pixels required by current image.

    Memory efficient for large numbers of states.
    """

    def __cinit__(self, States states):

        self._states = states
        self._num_states = states._num_states

    cpdef energy_vector_in_place(self, GridDataInt griddata,
                                 double[::1] energies):
        """Return vector with energy for each state at this griddata."""

        cdef:
            int i, mi, mj, K
            double std_range = 5.0
            double ll
            Rotation3D rotation
            double[:, :] matrix
            double[:, ::1] rotations
            IsoMixture3D transformed
            IsoMixture2D projected
            Transformation3D transformation

        rotations = self._states._rotations
        mixture = self._states._mixture

        # Init memory.
        K = mixture._num_components
        transformed = IsoMixture3D(np.zeros([K, 3], dtype='float'),
                                   1.0, np.zeros(K, dtype='float'))
        projected = IsoMixture2D(np.zeros([K, 2], dtype='float'),
                                 1.0, np.zeros(K, dtype='float'))
        transformation = Transformation3D.random()

        for i in range(self._num_states):
            # Update 'transformation' with current rotation.
            rotation = transformation._rotation
            matrix = rotation._matrix
            for mi in range(3):
                for mj in range(3):
                    matrix[mi, mj] = rotations[i, mi * 3 + mj]
            # Compute energy.
            mixture._transform(transformation, transformed)
            transformed._project(projected)

            ll = projected.log_likelihood_grid_sparse(griddata, std_range)
            energies[i] = -ll


    cpdef energy_matrix_in_place(self, list griddatas,
                                 double[:, :] energies):
        """Return matrix with energy for each state at each griddata."""

        """Compute energy of each rotation at each griddata.

        Args:
            griddatas: List of M GridDataInt objects.
            energies: (len(griddatas), num_states), initialised.

        Return:
            (N, M) array, where N is the number of states.

        TODO: compute griddata_diffs, for faster energy evaluation.
        But first make sure that the simpler approach works.
        """

        cdef:
            int i, k, mi, mj
            double std_range = 5.0
            double ll
            Rotation3D rotation
            double[:, :] matrix, rotations
            IsoMixture3D transformed
            IsoMixture2D projected
            Transformation3D transformation

        rotations = self._states._rotations
        mixture = self._states._mixture

        # Init memory.
        K = mixture._num_components
        transformed = IsoMixture3D(np.zeros([K, 3], dtype='float'),
                                   1.0, np.zeros(K, dtype='float'))
        projected = IsoMixture2D(np.zeros([K, 2], dtype='float'),
                                 1.0, np.zeros(K, dtype='float'))
        transformation = Transformation3D.random()

        ## self._num_replicas = len(griddatas)

        for k, griddata in enumerate(griddatas):

            # Compute k'th row.
            for i in range(self._states._num_states):
                # Update 'transformation' with current rotation.
                rotation = transformation._rotation
                matrix = rotation._matrix
                for mi in range(3):
                    for mj in range(3):
                        matrix[mi, mj] = rotations[i, mi * 3 + mj]
                # Compute energy.
                mixture._transform(transformation, transformed)
                transformed._project(projected)

                ll = projected.log_likelihood_grid_sparse(griddata,
                                                          std_range)
                energies[k, i] = -ll


    cpdef energy_vector(self, GridDataInt griddata):
        """Return vector with energy for each state at this griddata."""

        cdef:
            double[::1] energies

        energies = np.empty((self._num_states, ), dtype='double')

        self.energy_vector_in_place(griddata, energies)

        return energies

    cpdef energy_matrix(self, list griddatas):
        """Return matrix with energy for each state at each griddata."""

        cdef:
            double[:, :] energies

        energies = np.empty((len(griddatas), self._num_states),
                            dtype='double')

        self.energy_matrix_in_place(griddatas, energies)

        return energies



@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class StatesGridEnergies:
    """Fast calculation of state energy vector for new replicas."""

    def __cinit__(self, States states,
                  Grid2D grid):
        """
        Args:
            rotations: State for which energy will be calculated.
            mixture: 3D mixture to rotate and project using state.
            grid: Grid onto which mixture should be projected.

        """

        cdef:
            Transformation3D transformation
            IsoMixture3D transformed
            IsoMixture2D projected
            Rotation3D rotation
            double[:, :] matrix
            double[:] v
            int mi, mj, K, i, j
            double threshold, vj
            double[:, ::1] rotations = states._rotations
            IsoMixture3D mixture = states._mixture

        self._num_states = rotations.shape[0]
        self._grid_size = grid.size

        # Initialise memory.
        K = mixture._num_components
        transformation = Transformation3D.random()
        transformed = IsoMixture3D(np.zeros([K, 3], dtype='float'),
                                   1.0, np.zeros(K, dtype='float'))
        projected = IsoMixture2D(np.zeros([K, 2], dtype='float'),
                                 1.0, np.zeros(K, dtype='float'))
        self._energy = np.empty((self._num_states, self._grid_size),
                                dtype='float')

        rotation = transformation._rotation
        matrix = rotation._matrix
        std_range = 10
        threshold = 1e-320

        # Compute the energy at every pixel for every state.
        for i in range(self._num_states):

            for mi in range(3):
                for mj in range(3):
                    matrix[mi, mj] = rotations[i, mi * 3 + mj]

            mixture._transform(transformation, transformed)
            transformed._project(projected)
            griddata = projected.evaluate_grid_sparse(grid, std_range)

            v = griddata._values

            for j in range(self._grid_size):
                vj = v[j]
                if vj < threshold:
                    self._energy[i, j] = -log(threshold)
                else:
                    self._energy[i, j] = -log(vj)

    cpdef energy_vector_in_place(self, GridDataInt griddata,
                                 double[::1] energies):
        """Compute energies for all states at this replica.

        Args:
            griddata: Image at which to evaluate all states.
                griddata.grid should be equal to the original grid
                used in the constructor.
            energies: Double array of length 'num_states'.

        This method essentially just computes the product of a matrix
        with a vector.
        """

        cdef:
            double[:, :] energy = self._energy
            long[:] values = griddata._values
            int i, j

        for i in range(self._num_states):
            energies[i] = 0.0
            for j in range(self._grid_size):
                energies[i] += values[j] * energy[i, j]

    cpdef energy_matrix_in_place(self, list griddatas,
                                 double[:, :] energies):
        """Essentially just a number of matrix-vector multiplications."""

        cdef:
            double[:, :] energy = self._energy
            long[:] values
            int i, j, k
            GridDataInt griddata

        for k, griddata in enumerate(griddatas):
            values = griddata._values
            for i in range(self._num_states):
                for j in range(self._grid_size):
                    energies[k, i] += values[j] * energy[i, j]

    cpdef energy_vector(self, GridDataInt griddata):
        """Return vector with energy for each state at this griddata."""

        cdef:
            double[::1] energies

        energies = np.empty((self._num_states, ), dtype='double')

        self.energy_vector_in_place(griddata, energies)

        return energies

    cpdef energy_matrix(self, list griddatas):
        """Return matrix with energy for each state at each griddata."""

        cdef:
            double[:, :] energies

        energies = np.empty((len(griddatas), self._num_states),
                            dtype='double')

        self.energy_matrix_in_place(griddatas, energies)

        return energies


@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class StatesGridEnergiesGaussian:
    """Fast calculation of state energy vector for new replicas.

    Uses Gaussian noise model.
    """

    def __cinit__(self, double[:, ::1] rotations,
                  IsoMixture3D mixture,
                  Grid2D grid,
                  unnormalised):
        """
        Args:
            rotations: State for which energy will be calculated.
            mixture: 3D mixture to rotate and project using state.
            grid: Grid onto which mixture should be projected.
        """

        cdef:
            Transformation3D transformation
            IsoMixture3D transformed
            IsoMixture2D projected
            Rotation3D rotation
            double[:, :] matrix
            double[:] v
            int mi, mj, K, i, j
            double threshold, vj

        self._num_states = rotations.shape[0]
        self._grid_size = grid.size

        print 'num states = %d' % self._num_states


        # Initialise memory.
        K = mixture._num_components
        transformation = Transformation3D.random()
        transformed = IsoMixture3D(np.zeros([K, 3], dtype='float'),
                                   1.0, np.zeros(K, dtype='float'))
        projected = IsoMixture2D(np.zeros([K, 2], dtype='float'),
                                 1.0, np.zeros(K, dtype='float'))
        self._predicted = np.empty((self._num_states, self._grid_size),
                                   dtype='float')
        evaluated = GridDataReal(
            grid, np.zeros(self._grid_size, dtype='float'))

        rotation = transformation._rotation
        matrix = rotation._matrix
        std_range = 5

        # Evaluate projection at every pixel - predicted values.
        # Then just need observed values to compute energy.
        for i in range(self._num_states):
            for mi in range(3):
                for mj in range(3):
                    matrix[mi, mj] = rotations[i, mi * 3 + mj]

            mixture._transform(transformation, transformed)
            transformed._project(projected)
            projected.evaluate_grid_in_place(
                grid, std_range, evaluated, unnormalised)

            for j in range(self._grid_size):
                self._predicted[i, j] = evaluated._values[j]

    cpdef energy_vector_in_place(self, GridDataReal griddata,
                                 double[::1] energies,
                                 double beta_i):
        """Compute energies for all states at this replica.

        Args:
            griddata: Image at which to evaluate all states.
                griddata.grid should be equal to the original grid
                used in the constructor.
            energies: Double array of length 'num_states'.
            beta_i: Noise precision for this image.

        This method just computes the residual between two vectors.
        """

        cdef:
            int i, j
            double[:, :] predicted = self._predicted
            double[:] observed = griddata._values
            
        # print 'energies before:'
        # print np.array(energies).shape
        # print 'num states = %d' % self._num_states
        # print np.array(energies)


        for i in range(self._num_states):
            energies[i] = 0.0
            for j in range(self._grid_size):
                energies[i] += (predicted[i, j] - observed[j])**2
            energies[i] *= 0.5 * beta_i

        # print 'energies:'
        # print np.array(energies)


    cpdef energy_vector(self, GridDataReal griddata, double beta_i):
        """Return vector with energy for each state at this griddata.
        """

        cdef:
            double[::1] energies

        energies = np.empty((self._num_states, ), dtype='double')

        self.energy_vector_in_place(griddata, energies, beta_i)

        return energies






@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class DOSEstimator:
    """Estimate density of states from energies."""

    def __cinit__(self, double[:, :] energy, double tol=1e-10):

        self._energy = energy
        self._tol = tol

        self._num_replicas = energy.shape[0]
        self._num_states = energy.shape[1]

        self._verbose = 0

    cpdef set_verbose(self, verbose):
        self._verbose = verbose

    cdef double rel_error(self, double a, double b):

        cdef:
            double c

        c = (a - b) / (a + b)

        if c > 0:
            return c
        else:
            return -c

    cpdef run(self, int num_steps):
        """Find ML estimate of DOS."""

        cdef:
            double[::1] log_dos, free_energy, state_array, replica_array
            int i, replica_index, state_index
            double N, log_N, ll, ll_previous, lse

        # Initialise.
        log_dos = np.zeros((self._num_states,), dtype='double')
        N = self._num_states / self._num_replicas
            # N = number of states per replica, assumed same for each.
        log_N = log(N)
        free_energy = np.empty((self._num_replicas, ), dtype='double')
        # Just two arrays for storing temp results.
        state_array = np.empty((self._num_states,), dtype='double')
        replica_array = np.empty((self._num_replicas,), dtype='double')

        ll_previous = 0.0 # Just a dummy value which won't be used.
            # For testing convergence.

        for i in range(num_steps):

            # Update free energies.
            for replica_index in range(self._num_replicas):
                for state_index in range(self._num_states):
                    state_array[state_index] = (
                        -self._energy[replica_index, state_index] +
                        log_dos[state_index])
                free_energy[replica_index] = -log_sum_exp(state_array)
            ## print np.array(free_energy)

            # Update density of states.
            for state_index in range(self._num_states):
                for replica_index in range(self._num_replicas):
                    replica_array[replica_index] = (
                        log_N - self._energy[replica_index, state_index] +
                        free_energy[replica_index])
                log_dos[state_index] = -log_sum_exp(replica_array)
            ## print np.array(log_dos)

            # Normalise density of states.
            lse = log_sum_exp(log_dos)
            for state_index in range(self._num_states):
                log_dos[state_index] -= lse

            # Compute log-likelihood.
            ll = -N * sum_array(free_energy)

            if i > 0 and self.rel_error(ll, ll_previous) < self._tol:
                break
            ll_previous = ll

            if i > 0 and self._verbose > 0 and i % self._verbose == 0:
                print '%4d: %1.6f' % (i, ll)

        self._free_energy = free_energy

        return log_dos

    cpdef double[:] get_free_energy(self):
        return self._free_energy


@cython.boundscheck(False)
@cython.cdivision(True)
@cython.wraparound(False)
cdef class DOS:

    def __cinit__(self, States states, double[::1] log_dos):
        self._states = states
        self._log_dos = log_dos

        self._num_states = states._num_states
        self._state_array = np.empty((self._num_states, ), dtype='double')
        self._log_p1 = np.empty((self._num_states, ), dtype='double')
        self._log_p2 = np.empty((self._num_states, ), dtype='double')
        self._energy_1 = np.empty((self._num_states, ), dtype='double')
        self._energy_2 = np.empty((self._num_states, ), dtype='double')

    @property
    def states(self):
        return self._states

    @property
    def log_dos(self):
        return self._log_dos

    cpdef double swap_rate(self, double beta_1, double beta_2,
                           GridDatasNested data):
        """Estimate the swap rate between the given betas."""

        cdef:
            double[::1] energy_1, energy_2, log_p1, log_p2
            double c1, c2, a, b, c, sr, accuracy
            int i, j

        ## times = []
        ## times.append(time.time())

        # Get energy vectors.
        self._states.energy_at_beta_exact(beta_1, data, self._energy_1)
        self._states.energy_at_beta_exact(beta_2, data, self._energy_2)

        ## times.append(time.time())

        # Compute log-normalisation constants.
        c1 = self.log_normalisation_constant(self._energy_1)
        c2 = self.log_normalisation_constant(self._energy_2)

        ## times.append(time.time())

        # Compute probability vectors.
        # (Could also just compute these, and then normalise them.)
        self.compute_log_probability_vector(
            self._energy_1, c1, self._log_p1)
        self.compute_log_probability_vector(
            self._energy_2, c2, self._log_p2)

        ## times.append(time.time())

        # Estimate swap rate.
        # Note that this could be twice as fast - only need to look at
        # half of the matrix.

        # Determine threshold below which we don't need to compute
        # the exponent, because it's too close to 0.
        accuracy = -log(self._num_states * self._num_states * 100.0)

        log_p1 = self._log_p1
        log_p2 = self._log_p2
        sr = 0.0
        for i in range(self._num_states):
            for j in range(self._num_states):
                a = log_p1[i] + log_p2[j]
                b = log_p1[j] + log_p2[i]
                if a < b:
                    c = a
                else:
                    c = b
                if c > accuracy:
                    sr += exp(c)

        ## times.append(time.time())

        ## times = np.array(times)
        ## l = len(times)
        ## diffs = times[1:] - times[:(l-1)]
        ## p = diffs / np.sum(diffs)
        ## print((100*p).astype('i'))


        return sr


    cdef double log_normalisation_constant(self, double[::1] energy):
        """Estimate the normalisation constant."""

        cdef:
            int i
            double[::1] log_dos = self._log_dos
            double[::1] state_array = self._state_array

        for i in range(self._num_states):
            state_array[i] = log_dos[i] - energy[i]

        return log_sum_exp(state_array)

    cdef compute_log_probability_vector(self,
                                        double[::1] energy,
                                        double log_constant,
                                        double[::1] p_vector):
        """Compute the vector needed for estimating the swap rate."""

        cdef:
            int i
            double[::1] log_dos = self._log_dos

        for i in range(self._num_states):
            p_vector[i] = log_dos[i] - energy[i] - log_constant

    def save(self, parent_group, name='dos'):
        """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 dos group.
            name: Name of the new dos group.
        """

        group = parent_group.create_group(name)
        self._states.save(group)
        dataset_log_dos = group.create_dataset(
            'log_dos', (<object>self._log_dos).shape, np.double)
        dataset_log_dos[...] = self._log_dos

    @classmethod
    def load(cls, parent_group, name='dos'):
        """Load DOS object from group in HDF5 database."""

        group = parent_group[name]
        states = States.load(group)
        log_dos = group['log_dos'][:]

        return cls(states, log_dos)




class ScheduleOptimiser(object):
    """Optimise beta schedule to achieve a target exchange rate."""

    def __init__(self, dos, data):
        """
        Args:
            dos: DOS object, density of states.
            data: GridDataNested object.
        """

        self.dos = dos
        self.data = data

    def optimise(self, beta_low, target_rate=0.23):
        """
        Optimise beta schedule.

        Args:
            beta_low: double, typically corresponding to no counts.
                So could be 1/C, where C is the total number of counts.
            target_rate: double, desired swap rate between neighbouring
                replicas. Typically 0.23.

        Return:
            Array of beta values from 'beta_low' to 1 such that the
            estimated swap rate between neighbouring replicas is
            exactly 'target' (except for the last pair, which will be
            at least 'target').
        """

        from scipy.optimize import bisect

        log_beta_low = np.log10(beta_low)
        log_beta_current = log_beta_low
        log_betas_optimized = [log_beta_current]
        swap_rate = self.dos.swap_rate

        while swap_rate(10**log_beta_current, 1.0, self.data) < target_rate:

            def f(log_beta2):
                return (swap_rate(10**log_beta_current,
                                  10**log_beta2, self.data) -
                        target_rate)

            log_beta_next = bisect(f, log_beta_current, 0.0, xtol=1e-3)
            log_betas_optimized.append(log_beta_next)
            log_beta_current = log_beta_next

        if len(log_betas_optimized) == 1:
            # Add an extra beta in the middle to have at least three.
            log_betas_optimized.append(10 ** (log_beta_low / 2.0))

        log_betas_optimized.append(0.0)

        return 10**np.array(log_betas_optimized)




cdef double sum_array(double[::1] x):

    cdef:
        double acc
        int i

    acc = 0.0
    for i in range(x.shape[0]):
        acc += x[i]

    return acc


