# Path: amdtk/densities/efd.py
# class EFDPrior(PersistentModel, metaclass=abc.ABCMeta):
#     """Abstract base class for a prior from the EFD."""
# 
#     @abc.abstractstaticmethod
#     def _log_partition_func(natural_params):
#         pass
# 
#     @abc.abstractstaticmethod
#     def _grad_log_partition_func(natural_params):
#         pass
# 
#     def __init__(self):
#         self._natural_params = None
#         self._log_partition =  None
#         self._grad_log_partition = None
# 
#     def correct_np_value(self, value):
#         """Correct the value of the natural parameters.
# 
#         When doing a gradient ascent, it may happen that the update
#         leads to values outside the domain of the natural parameters.
#         This function simply project back the gradient to the closest
#         point inside the domain of the natural parameters.
# 
#         The default behavior of this function is to let the natural
#         parameters untouch. It is the responsibilities of the
#         subclass to implement the correct behavior.
# 
#         Parameters
#         ----------
#         value : numpy.ndarray
#             New value of the natural parameters
# 
#         Returns
#         -------
#         corrected_value : numpy.ndarray
#             Corrected value of the natural parameters.
# 
#         """
#         return value
# 
#     @property
#     def natural_params(self):
#         """Vector of natural parameters."""
#         return self._natural_params
# 
#     @natural_params.setter
#     def natural_params(self, value):
#         #corrected_value = self.correct_np_value(value)
#         corrected_value = value
#         self._natural_params = corrected_value
#         self._log_partition = self._log_partition_func(corrected_value)
#         self._grad_log_partition = \
#             self._grad_log_partition_func(corrected_value)
# 
#     @property
#     def log_partition(self):
#         """Log-partition of the distribution."""
#         return self._log_partition
# 
#     @property
#     def grad_log_partition(self):
#         """Gradient of the log-partition.
# 
#         Gradient of the log-partition with respect to the natural
#         parameters. This corresponds to the expected value of the
#         sufficient statistics of the current density/distribution.
# 
#         """
#         return self._grad_log_partition
# 
#     def kl_div(self, dist):
#         """Kullback-Leibler divergence.
# 
#         Compute the Kullback-Leibler divergence between the current
#         distribution and the given distribution.
# 
#         Parameters
#         ----------
#         dist : EFDPrior
#             Distribution of the same type as the current distribution
#             (i.e. "self").
# 
# 
#         -------
#         div : float
#             Results of the KL divergence in nats.
# 
#         """
#         # Expected value of the sufficient statistics with respect to
#         # the current distribution
#         expected_value = self.grad_log_partition
# 
#         # Natural parameters of the current and given distributions.
#         nparams1 = self.natural_params
#         nparams2 = dist.natural_params
# 
#         # Log-partition of the current and given distributions.
#         log_partition1 = self.log_partition
#         log_partition2 = dist.log_partition
# 
#         # Compute the KL divergence.
#         retval = (nparams1 - nparams2).dot(expected_value)
#         retval += log_partition2 - log_partition1
# 
#         return retval
# 
#     # PersistentModel interface implementation.
#     # -----------------------------------------------------------------
# 
#     def to_dict(self):
#         return {
#             'class': self.__class__,
#             'natural_params': self._natural_params
#         }
# 
#     @classmethod
#     def load_from_dict(cls, model_data):
#         model = cls.__new__(model_data['class'])
#         model.natural_params = model_data['natural_params']
#         return model
# 
#     # -----------------------------------------------------------------


import theano
import theano.tensor as T
import numpy as np
from .efd import EFDPrior
# ]
#
# log_partition(kappa, mean, rate, scale) =
#   gammaln(rate) - rate * log(scale) - .5 * log(kappa)
#
# log_partition(np1, np2, np3, np4) =
#   gammaln(.5 * (np4 + 1)) - .5 * (np4 + 1) log(.5 *
#       (np1 - (np2 ** 2) / np3)) - .5 * log(np3)
#




def _log_partition_symfunc():
    natural_params = T.vector()
    size = natural_params.shape[0] // 4
    np1, np2, np3, np4 = T.split(natural_params, 4 * [size], 4)

    log_Z = T.sum(T.gammaln(.5 * (np4 + 1)))
    log_Z += T.sum(- .5 * (np4 + 1) * T.log(.5 * (np1 - (np2 ** 2) / np3)))
    log_Z += T.sum(-.5 * T.log(np3))

    func = theano.function([natural_params], log_Z)
    grad_func = theano.function([natural_params],
                                T.grad(T.sum(log_Z), natural_params))
    return func, grad_func

_lp_func, _grad_lp_func = _log_partition_symfunc()



# Path: amdtk/densities/efd.py
# class EFDPrior(PersistentModel, metaclass=abc.ABCMeta):
#     """Abstract base class for a prior from the EFD."""
# 
#     @abc.abstractstaticmethod
#     def _log_partition_func(natural_params):
#         pass
# 
#     @abc.abstractstaticmethod
#     def _grad_log_partition_func(natural_params):
#         pass
# 
#     def __init__(self):
#         self._natural_params = None
#         self._log_partition =  None
#         self._grad_log_partition = None
# 
#     def correct_np_value(self, value):
#         """Correct the value of the natural parameters.
# 
#         When doing a gradient ascent, it may happen that the update
#         leads to values outside the domain of the natural parameters.
#         This function simply project back the gradient to the closest
#         point inside the domain of the natural parameters.
# 
#         The default behavior of this function is to let the natural
#         parameters untouch. It is the responsibilities of the
#         subclass to implement the correct behavior.
# 
#         Parameters
#         ----------
#         value : numpy.ndarray
#             New value of the natural parameters
# 
#         Returns
#         -------
#         corrected_value : numpy.ndarray
#             Corrected value of the natural parameters.
# 
#         """
#         return value
# 
#     @property
#     def natural_params(self):
#         """Vector of natural parameters."""
#         return self._natural_params
# 
#     @natural_params.setter
#     def natural_params(self, value):
#         #corrected_value = self.correct_np_value(value)
#         corrected_value = value
#         self._natural_params = corrected_value
#         self._log_partition = self._log_partition_func(corrected_value)
#         self._grad_log_partition = \
#             self._grad_log_partition_func(corrected_value)
# 
#     @property
#     def log_partition(self):
#         """Log-partition of the distribution."""
#         return self._log_partition
# 
#     @property
#     def grad_log_partition(self):
#         """Gradient of the log-partition.
# 
#         Gradient of the log-partition with respect to the natural
#         parameters. This corresponds to the expected value of the
#         sufficient statistics of the current density/distribution.
# 
#         """
#         return self._grad_log_partition
# 
#     def kl_div(self, dist):
#         """Kullback-Leibler divergence.
# 
#         Compute the Kullback-Leibler divergence between the current
#         distribution and the given distribution.
# 
#         Parameters
#         ----------
#         dist : EFDPrior
#             Distribution of the same type as the current distribution
#             (i.e. "self").
# 
# 
#         -------
#         div : float
#             Results of the KL divergence in nats.
# 
#         """
#         # Expected value of the sufficient statistics with respect to
#         # the current distribution
#         expected_value = self.grad_log_partition
# 
#         # Natural parameters of the current and given distributions.
#         nparams1 = self.natural_params
#         nparams2 = dist.natural_params
# 
#         # Log-partition of the current and given distributions.
#         log_partition1 = self.log_partition
#         log_partition2 = dist.log_partition
# 
#         # Compute the KL divergence.
#         retval = (nparams1 - nparams2).dot(expected_value)
#         retval += log_partition2 - log_partition1
# 
#         return retval
# 
#     # PersistentModel interface implementation.
#     # -----------------------------------------------------------------
# 
#     def to_dict(self):
#         return {
#             'class': self.__class__,
#             'natural_params': self._natural_params
#         }
# 
#     @classmethod
#     def load_from_dict(cls, model_data):
#         model = cls.__new__(model_data['class'])
#         model.natural_params = model_data['natural_params']
#         return model
# 
#     # -----------------------------------------------------------------


# Path: amdtk/densities/normal_gamma.py
import theano
import theano.tensor as T
import numpy as np
from .efd import EFDPrior

# ]
#
# log_partition(kappa, mean, rate, scale) =
#   gammaln(rate) - rate * log(scale) - .5 * log(kappa)
#
# log_partition(np1, np2, np3, np4) =
#   gammaln(.5 * (np4 + 1)) - .5 * (np4 + 1) log(.5 *
#       (np1 - (np2 ** 2) / np3)) - .5 * log(np3)
#




def _log_partition_symfunc():
    natural_params = T.vector()
    size = natural_params.shape[0] // 4
    np1, np2, np3, np4 = T.split(natural_params, 4 * [size], 4)

    log_Z = T.sum(T.gammaln(.5 * (np4 + 1)))
    log_Z += T.sum(- .5 * (np4 + 1) * T.log(.5 * (np1 - (np2 ** 2) / np3)))
    log_Z += T.sum(-.5 * T.log(np3))

    func = theano.function([natural_params], log_Z)
    grad_func = theano.function([natural_params],
                                T.grad(T.sum(log_Z), natural_params))
    return func, grad_func

_lp_func, _grad_lp_func = _log_partition_symfunc()



class NormalGamma(EFDPrior):