import numpy as np
import tensorflow as tf

import bayesflow as bf
import bayesflow.util as util


# TODO: do we really need separate classes to represent Q
# distributions? the "variational models" perspective is that they
# should just be (un)ConditionalDistribution objects.

class QDistribution(object):
    def __init__(self, shape):
        self.output_shape=shape
    
    def sample_stochastic_inputs(self):
        raise Exception("not implemented")

    def entropy(self):
        raise Exception("not implemented")

    def density(self):
        raise Exception("not implemented")
        
class ObservedQDistribution(QDistribution):
    def __init__(self, observed_val):
        shape = observed_val.shape
        super(ObservedQDistribution, self).__init__(shape=shape)
        
        self.mean = tf.constant(observed_val, name="mean")
        self.stddev = tf.zeros_like(self.mean, name="stddev")
        self.variance = tf.identity(self.stddev)
        self.sample = self.mean 
        
    def sample_stochastic_inputs(self):
        return {}
    
    def entropy(self):
        return 0.0
    
class GaussianQDistribution(QDistribution):
    
    def __init__(self, shape):
        super(GaussianQDistribution, self).__init__(shape=shape)
        
        init_mean = np.float32(np.random.randn(*shape))
        init_log_stddev = np.float32(np.ones(shape) * -10)
        self.mean = tf.Variable(init_mean, name="mean")
        self.log_stddev = tf.Variable(init_log_stddev, name="log_stddev")
        self.stddev = tf.exp(self.log_stddev)
        self.variance = tf.square(self.stddev)        
        self.stochastic_eps = tf.placeholder(dtype=self.mean.dtype, shape=shape, name="eps")
        self.sample = self.stochastic_eps * self.stddev + self.mean        
        
        self._entropy = tf.reduce_sum(bf.dists.gaussian_entropy(variance=self.variance))
        
    def sample_stochastic_inputs(self):
        sampled_eps = np.random.randn(*self.output_shape)
        return {self.stochastic_eps : sampled_eps}
    
    def entropy(self):
        return self._entropy
    
class BernoulliQDistribution(QDistribution):
    
    def __init__(self, shape):
        super(BernoulliQDistribution, self).__init__(shape=shape)
        
        init_log_odds = np.float32(np.random.randn(*shape))
        self.logodds = tf.Variable(init_log_odds, name="logodds")
        self.probs, _ = bf.transforms.logit(self.logodds)

        self.stochastic_eps = tf.placeholder(dtype=self.probs.dtype, shape=shape, name="eps")
        self.sample = self.stochastic_eps < self.probs     
        
        self._entropy = tf.reduce_sum(bf.dists.bernoulli_entropy(p=self.probs))
        
    def sample_stochastic_inputs(self):
        sampled_eps = np.random.rand(*self.output_shape)
        return {self.stochastic_eps : sampled_eps}
    
    def entropy(self):
        return self._entropy

@@ -1103,16 +1103,1720 @@
 urn 0.0%0A
+%0Aclass DeltaQDistribution(QDistribution):%0A    def __init__(self, tf_value):%0A%0A        shape = util.extract_shape(tf_value)%0A        super(DeltaQDistribution, self).__init__(shape=shape)%0A        %0A        self.mean = tf_value%0A        self.stddev = tf.zeros_like(self.mean, name=%22stddev%22)%0A        self.variance = tf.identity(self.stddev)%0A        self.sample = self.mean %0A        %0A    def sample_stochastic_inputs(self):%0A        return %7B%7D%0A    %0A    def entropy(self):%0A        return 0.0%0A%0Aclass PointwiseTransformedQDistribution(QDistribution):%0A    def __init__(self, parent_q, transform, implicit=False):%0A%0A        %0A        super(PointwiseTransformedQDistribution, self).__init__(shape=parent_q.output_shape)%0A        self.sample, self.log_jacobian = transform(parent_q.sample)%0A%0A        # an %22implicit%22 transformation is a formal object associated with a deterministic%0A        # transformation in the model, where the parent q distribution is associated with%0A        # the untransformed variable. In this case the ELBO expectations are with respect%0A        # to the parent Q distribution, so the jacobian of the transformation does not appear.%0A        # By contrast, a non-implicit use of a transformed Q distribution would be to create%0A        # a new type of distribution (eg, lognormal by exponentiating a Gaussian parent)%0A        # that could then itself be associated with stochastic variables in the graph.%0A        self.implicit = implicit%0A        self.parent_q = parent_q%0A        %0A    def sample_stochastic_inputs(self):%0A        return %7B%7D%0A    %0A    def entropy(self):%0A        if self.implicit:%0A            return 0.0%0A        else:%0A            return self.parent_q.entropy() + self.log_jacobian%0A        %0A    
     %0Acla
@@ -2881,32 +2881,59 @@
 it__(self, shape
+, monte_carlo_entropy=False
 ):%0A        super
@@ -3504,32 +3504,33 @@
 ean        %0A
+%0A
         
 %0A        sel
@@ -3509,33 +3509,218 @@
       %0A%0A        
+# HACK
 %0A
+        if monte_carlo_entropy:%0A            self._entropy = -(tf.reduce_sum(bf.dists.gaussian_log_density(self.sample, mean=self.mean, variance=self.variance)))%0A        else:%0A    
         self._en
@@ -3788,24 +3788,26 @@
 .variance))%0A
+%0A%0A
         %0A   

add delta and pointwise-transformed distributions