################################################################################
# Copyright (C) 2013-2014 Jaakko Luttinen
#
# This file is licensed under the MIT License.
################################################################################


import warnings

import numpy as np

from bayespy.utils import misc

from .node import ensureparents
from .stochastic import Stochastic, Distribution

class ExponentialFamilyDistribution(Distribution):
    """
    Sub-classes implement distribution specific computations.
    """

    #
    # The following methods are for ExponentialFamily distributions
    #

    def compute_message_to_parent(self, parent, index, u_self, *u_parents):
        raise NotImplementedError()

    def compute_phi_from_parents(self, *u_parents, mask=True):
        raise NotImplementedError()

    def compute_moments_and_cgf(self, phi, mask=True):
        raise NotImplementedError()

    #
    # The following methods are for Mixture class
    #

    def compute_cgf_from_parents(self, *u_parents):
        raise NotImplementedError()

    def compute_fixed_moments_and_f(self, x, mask=True):
        raise NotImplementedError()

    def compute_logpdf(self, u, phi, g, f, ndims):
        """ Compute E[log p(X)] given E[u], E[phi], E[g] and
        E[f]. Does not sum over plates."""

        # TODO/FIXME: Should I take into account what is latent or
        # observed, or what is even totally ignored (by the mask).
        L = g + f
        for (phi_i, u_i, ndims_i) in zip(phi, u, ndims):
            # Axes to sum (dimensions of the variable, not the plates)
            axis_sum = tuple(range(-ndims_i,0))
            # Compute the term
            # TODO/FIXME: Use einsum!
            L = L + np.sum(
                np.where(u_i != 0, phi_i, 0) * u_i,
                axis=axis_sum
            )
        return L


    def compute_gradient(self, g, u, phi):
        r"""
        Compute the standard gradient with respect to the natural parameters.
        """

        raise NotImplementedError("Standard gradient not yet implemented for %s"
                                  % (self.__class__.__name__))



def useconstructor(__init__):
    def constructor_decorator(self, *args, **kwargs):
        if (self.dims is None or
            self._distribution is None or
            self._moments is None or
            self._parent_moments is None):

            (args, kwargs, dims, plates, dist, stats, pstats) = \
              self._constructor(*args, **kwargs)

            self.dims = dims
            self._distribution = dist
            self._moments = stats
            self._parent_moments = pstats
            self.plates = plates

        __init__(self, *args, **kwargs)

    return constructor_decorator

class ExponentialFamily(Stochastic):
    """
    A base class for nodes using natural parameterization `phi`.

    phi

    Sub-classes must implement the following static methods:
       _compute_message_to_parent(index, u_self, *u_parents)
       _compute_phi_from_parents(*u_parents, mask)
       _compute_moments_and_cgf(phi, mask)
       _compute_fixed_moments_and_f(x, mask=True)

    Sub-classes may need to re-implement:
    1. If they manipulate plates:
       _compute_weights_to_parent(index, weights)
       _compute_plates_to_parent(self, index, plates)
       _compute_plates_from_parent(self, index, plates)

    """

    # Sub-classes should overwrite this (possibly using _constructor)
    dims = None

    # Sub-classes should overwrite this
    _distribution = None

    @useconstructor
    def __init__(self, *parents, initialize=True, phi_bias=None, **kwargs):

        self.annealing = 1.0

        # Terms for the lower bound (G for latent and F for observed)
        self.g = np.array(np.nan)
        self.f = np.array(np.nan)

        self._phi_bias = phi_bias if phi_bias is not None else len(self.dims) * [0.0]

        super().__init__(*parents,
                         initialize=initialize,
                         dims=self.dims,
                         **kwargs)

        if not initialize:
            axes = len(self.plates)*(1,)
            self.phi = [misc.nans(axes+dim) for dim in self.dims]


    @classmethod
    @ensureparents
    def _constructor(cls, *parents, **kwargs):
        """
        Constructs distribution and moments objects.

        If __init__ uses useconstructor decorator, this method is called to
        construct distribution and moments objects.

        The method is given the same inputs as __init__. For some nodes, some of
        these can't be "static" class attributes, then the node class must
        overwrite this method to construct the objects manually.

        The point of distribution class is to move general distribution but
        not-node specific code. The point of moments class is to define the
        messaging protocols.
        """
        parent_plates = [cls._distribution.plates_from_parent(ind, parent.plates)
                         for (ind, parent) in enumerate(parents)]
        return (parents,
                kwargs,
                cls.dims,
                cls._total_plates(kwargs.get('plates'), *parent_plates),
                cls._distribution,
                cls._moments,
                cls._parent_moments)

    def _initialize_from_parent_moments(self, *u_parents):
        if not np.all(self.observed):
            # Update natural parameters using parents
            self._update_phi_from_parents(*u_parents)

            # Update moments
            mask = np.logical_not(self.observed)
            (u, g) = self._distribution.compute_moments_and_cgf(self.phi,
                                                                mask=mask)
            # TODO/FIXME/BUG: You should use observation mask in order to not
            # overwrite them!
            self._set_moments_and_cgf(u, g, mask=mask)


    def initialize_from_prior(self):
        u_parents = self._message_from_parents()
        self._initialize_from_parent_moments(*u_parents)


    def initialize_from_parameters(self, *args):
        u_parents = [p_mom.compute_fixed_moments(x)
                     for (p_mom, x) in zip(self._parent_moments, args)]
        self._initialize_from_parent_moments(*u_parents)


    def initialize_from_value(self, x, *args):
        # Update moments from value
        mask = np.logical_not(self.observed)
        u = self._moments.compute_fixed_moments(x, *args)
        # Check that the shape is correct
        for i in range(len(u)):
            ndim = len(self.dims[i])
            if ndim > 0:
                if np.shape(u[i])[-ndim:] != self.dims[i]:
                    raise ValueError("The initial value for node %s has invalid shape %s."
                                     % (np.shape(x)))
        self._set_moments_and_cgf(u, np.inf, mask=mask)

    def initialize_from_random(self):
        """
        Set the variable to a random sample from the current distribution.
        """
        #self.initialize_from_prior()
        X = self.random()
        self.initialize_from_value(X)


    def _update_phi_from_parents(self, *u_parents):

        # TODO/FIXME: Could this be combined to the function
        # _update_distribution_and_lowerbound ?
        # No, because some initialization methods may want to use this.

        # This makes correct broadcasting
        self.phi = [
            a + b
            for (a, b) in zip(
                self._distribution.compute_phi_from_parents(*u_parents),
                self._phi_bias
            )
        ]
        # Make sure phi has the correct number of axes. It makes life
        # a bit easier elsewhere.
        for i in range(len(self.phi)):
            axes = len(self.plates) + self.ndims[i] - np.ndim(self.phi[i])
            if axes > 0:
                # Add axes
                self.phi[i] = misc.add_leading_axes(self.phi[i], axes)
            elif axes < 0:
                # Remove extra leading axes
                first = -(len(self.plates)+self.ndims[i])
                sh = np.shape(self.phi[i])[first:]
                self.phi[i] = np.reshape(self.phi[i], sh)
            # Check that the shape is correct
            if not misc.is_shape_subset(np.shape(self.phi[i]),
                                         self.get_shape(i)):
                raise ValueError("Incorrect shape of phi[%d] in node class %s. "
                                 "Shape is %s but it should be broadcastable "
                                 "to shape %s."
                                 % (i,
                                    self.__class__.__name__,
                                    np.shape(self.phi[i]),
                                    self.get_shape(i)))

    def _set_moments_and_cgf(self, u, g, mask=True):
        self._set_moments(u, mask=mask)

        self.g = np.where(mask, g, self.g)

        return


    def get_riemannian_gradient(self):
        r"""
        Computes the Riemannian/natural gradient.
        """
        u_parents = self._message_from_parents()
        m_children = self._message_from_children()

        # TODO/FIXME: Put observed plates to zero?
        # Compute the gradient
        phi = [
            a + b
            for (a, b) in zip(
                self._distribution.compute_phi_from_parents(*u_parents),
                self._phi_bias
            )
        ]
        for i in range(len(self.phi)):
            phi[i] = self.annealing * (phi[i] + m_children[i]) - self.phi[i]
            phi[i] = phi[i] * np.ones(self.get_shape(i))

        return phi


    def get_gradient(self, rg):
        r""" Computes gradient with respect to the natural parameters.

        The function takes the Riemannian gradient as an input.  This is for
        three reasons: 1) You probably want to use the Riemannian gradient
        anyway so this helps avoiding accidental use of this function.  2) The
        gradient is computed by using the Riemannian gradient and chain rules.
        3) Probably you need both Riemannian and normal gradients anyway so you
        can provide it to this function to avoid re-computing it."""

        g = self._distribution.compute_gradient(rg, self.u, self.phi)
        for i in range(len(g)):
            g[i] /= self.annealing
        return g


    ## def update_parameters(self, d, scale=1.0):
    ##     r"""
    ##     Update the parameters of the VB distribution given a change.

    ##     The parameters should be such that they can be used for
    ##     optimization, that is, use log transformation for positive
    ##     parameters.
    ##     """
    ##     phi = self.get_parameters()
    ##     for i in range(len(phi)):
    ##         phi[i] = phi[i] + scale*d[i]
    ##     self.set_parameters(phi)
    ##     return


    def get_parameters(self):
        r"""
        Return parameters of the VB distribution.

        The parameters should be such that they can be used for
        optimization, that is, use log transformation for positive
        parameters.
        """
        return [np.copy(p) for p in self.phi]



    def _decode_parameters(self, x):
        return [np.copy(p) for p in x]


    def set_parameters(self, x):
        r"""
        Set the parameters of the VB distribution.

        The parameters should be such that they can be used for
        optimization, that is, use log transformation for positive
        parameters.
        """
        self.phi = self._decode_parameters(x)
        self._update_moments_and_cgf()
        return


    def _update_distribution_and_lowerbound(self, m_children, *u_parents):

        # Update phi first from parents..
        self._update_phi_from_parents(*u_parents)
        # .. then just add children's message
        self.phi = [self.annealing * (phi + m)
                    for (phi, m) in zip(self.phi, m_children)]

        # Update u and g
        self._update_moments_and_cgf()


    def _update_moments_and_cgf(self):
        """
        Update moments and cgf based on current phi.
        """
        # Mask for plates to update (i.e., unobserved plates)
        update_mask = np.logical_not(self.observed)

        # Compute the moments (u) and CGF (g)...
        (u, g) = self._distribution.compute_moments_and_cgf(self.phi,
                                                            mask=update_mask)
        # ... and store them
        self._set_moments_and_cgf(u, g, mask=update_mask)


    def observe(self, x, *args, mask=True):
        """
        Fix moments, compute f and propagate mask.
        """

        # Compute fixed moments
        (u, f) = self._distribution.compute_fixed_moments_and_f(x, *args,
                                                                mask=mask)

        # # Check the dimensionality of the observations
        # self._check_shape()
        # for (i,v) in enumerate(u):
        #     # This is what the dimensionality "should" be
        #     s = self.plates + self.dims[i]
        #     t = np.shape(v)
        #     if s != t:
        #         msg = "Dimensionality of the observations incorrect."
        #         msg += "\nShape of input: " + str(t)
        #         msg += "\nExpected shape: " + str(s)
        #         msg += "\nCheck plates."
        #         raise Exception(msg)

        # Set the moments. Shape checking is done there.
        self._set_moments(u, mask=mask, broadcast=False)

        self.f = np.where(mask, f, self.f)

        # Observed nodes should not be ignored
        self.observed = mask
        self._update_mask()

    def lower_bound_contribution(self, gradient=False, ignore_masked=True):
        r"""Compute E[ log p(X|parents) - log q(X) ]

        If deterministic annealing is used, the term E[ -log q(X) ] is
        divided by the anneling coefficient.  That is, phi and cgf of q
        are multiplied by the temperature (inverse annealing
        coefficient).

        """

        # Annealing temperature
        T = 1 / self.annealing

        # Messages from parents
        u_parents = self._message_from_parents()
        phi = [
            a # + b # TODO: Should the bias be added here or not?
            for (a, b) in zip(
                self._distribution.compute_phi_from_parents(*u_parents),
                self._phi_bias
            )
        ]
        # G from parents
        L = self._distribution.compute_cgf_from_parents(*u_parents)

        # G for unobserved variables (ignored variables are handled properly
        # automatically)
        latent_mask = np.logical_not(self.observed)

        # G and F
        if np.all(self.observed):
            z = np.nan
        elif T == 1:
            z = -self.g
        else:
            z = -T * self.g
            ## TRIED THIS BUT IT WAS WRONG:
            ## z = -T * self.g + (1-T) * self.f
            ## if np.any(np.isnan(self.f)):
            ##     warnings.warn("F(x) not implemented for node %s. This "
            ##                   "is required for annealed lower bound "
            ##                   "computation." % self.__class__.__name__)
            ##
            ## It was wrong because the optimal q distribution has f which is
            ## weighted by 1/T and here the f of q is weighted by T so the
            ## total weight is 1, thus it cancels out with f of p.

        L = L + np.where(self.observed, self.f, z)

        for (phi_p, phi_q, u_q, dims) in zip(phi, self.phi, self.u, self.dims):
            # Form a mask which puts observed variables to zero and
            # broadcasts properly
            latent_mask_i = misc.add_trailing_axes(
                                misc.add_leading_axes(
                                    latent_mask,
                                    len(self.plates) - np.ndim(latent_mask)),
                                len(dims))
            axis_sum = tuple(range(-len(dims),0))

            # Compute the term
            phi_q = np.where(latent_mask_i, phi_q, 0)
            # Apply annealing
            phi_diff = phi_p - T * phi_q
            # Handle 0 * -inf
            phi_diff = np.where(u_q != 0, phi_diff, 0)
            # TODO/FIXME: Use einsum here?
            Z = np.sum(phi_diff * u_q, axis=axis_sum)

            L = L + Z

        if ignore_masked:
            return (np.sum(np.where(self.mask, L, 0))
                    * self.broadcasting_multiplier(self.plates,
                                                   np.shape(L),
                                                   np.shape(self.mask))
                    * np.prod(self.plates_multiplier))
        else:
            return (np.sum(L)
                    * self.broadcasting_multiplier(self.plates,
                                                   np.shape(L))
                    * np.prod(self.plates_multiplier))


    def logpdf(self, X, mask=True):
        """
        Compute the log probability density function Q(X) of this node.
        """
        if mask is not True:
            raise NotImplementedError('Mask not yet implemented')
        (u, f) = self._distribution.compute_fixed_moments_and_f(X, mask=mask)
        Z = 0
        for (phi_d, u_d, dims) in zip(self.phi, u, self.dims):
            axis_sum = tuple(range(-len(dims),0))
            # TODO/FIXME: Use einsum here?
            Z = Z + np.sum(phi_d * u_d, axis=axis_sum)
            #Z = Z + misc.sum_multiply(phi_d, u_d, axis=axis_sum)

        return (self.g + f + Z)


    def pdf(self, X, mask=True):
        """
        Compute the probability density function of this node.
        """
        return np.exp(self.logpdf(X, mask=mask))


    def _save(self, group):
        """
        Save the state of the node into a HDF5 file.

        group can be the root
        """
        ## if name is None:
        ##     name = self.name
        ## subgroup = group.create_group(name)

        for i in range(len(self.phi)):
            misc.write_to_hdf5(group, self.phi[i], 'phi%d' % i)
        misc.write_to_hdf5(group, self.f, 'f')
        misc.write_to_hdf5(group, self.g, 'g')
        super()._save(group)


    def _load(self, group):
        """
        Load the state of the node from a HDF5 file.
        """
        # TODO/FIXME: Check that the shapes are correct!
        for i in range(len(self.phi)):
            phii = group['phi%d' % i][...]
            self.phi[i] = phii

        self.f = group['f'][...]
        self.g = group['g'][...]
        super()._load(group)


    def random(self):
        """
        Draw a random sample from the distribution.
        """
        return self._distribution.random(*(self.phi), plates=self.plates)