import warnings
import numpy as np
import theano
import theano.tensor as T
from theano.sandbox.rng_mrg import MRG_RandomStreams
import treeano
import treeano.nodes as tn


fX = theano.config.floatX


@treeano.register_node("normal_sample")
class NormalSampleNode(treeano.NodeImpl):

    input_keys = ("mu", "sigma")

    def compute_output(self, network, mu_vw, sigma_vw):
        # TODO look at shape of both mu and sigma
        shape = mu_vw.shape
        if any(s is None for s in shape):
            # NOTE: this uses symbolic shape - can be an issue with
            # theano.clone and random numbers
            # https://groups.google.com/forum/#!topic/theano-users/P7Mv7Fg0kUs
            warnings.warn("using symbolic shape for random number shape, "
                          "which can be an issue with theano.clone")
            shape = mu_vw.variable.shape
        # TODO save this state so that we can seed the rng
        srng = MRG_RandomStreams()
        sampled = srng.normal(shape,
                              avg=mu_vw.variable,
                              std=sigma_vw.variable,
                              dtype=fX)
        network.create_vw(
            "default",
            variable=theano.gradient.disconnected_grad(sampled),
            shape=mu_vw.shape,
            tags={"output"},
        )


@treeano.register_node("normal_REINFORCE")
class NormalREINFORCECostNode(treeano.NodeImpl):

    input_keys = ("state", "mu", "sigma", "reward", "sampled")

    def compute_output(self,
                       network,
                       state_vw,
                       mu_vw,
                       sigma_vw,
                       reward_vw,
                       sampled_vw):
        # want state to have dim (batch size x size of state)
        assert state_vw.ndim == 2
        # want mu to have dim (batch size x number of actions)
        assert mu_vw.ndim == 2

        state = state_vw.variable
        mu = mu_vw.variable
        sigma = sigma_vw.variable
        reward = reward_vw.variable
        sampled = sampled_vw.variable

        # create reward baseline
        bias = network.create_vw(
            name="bias",
            is_shared=True,
            shape=(),
            tags={"parameter", "bias"},
            default_inits=[],
        ).variable
        weight = network.create_vw(
            name="weight",
            is_shared=True,
            shape=(state_vw.shape[1],),
            tags={"parameter", "weight"},
            default_inits=[],
        ).variable
        constant_state = theano.gradient.disconnected_grad(state)
        baseline = ((weight.dimshuffle("x", 0) * constant_state).sum(axis=1)
                    + bias)
        # TODO monitor baseline
        constant_baseline = theano.gradient.disconnected_grad(baseline)

        # 1 / (sigma * sqrt(2 * pi)) * exp(-1/2 * ((t - mu) / sigma)^2)
        normal_pdf = (1 / (sigma * treeano.utils.as_fX(np.sqrt(2 * np.pi)))
                      * T.exp(-0.5 * T.sqr((sampled - mu) / sigma)))
        log_normal_pdf = T.log(normal_pdf)
        R = reward - constant_baseline
        # take sum of log pdf
        reinforce_cost = -(R * log_normal_pdf.sum(axis=1)).sum()
        # TODO add parameter as weight for baseline
        baseline_cost = T.sum((reward - baseline) ** 2)

        network.create_vw(
            name="default",
            # variable=reinforce_cost,
            variable=reinforce_cost + baseline_cost,
            shape=(),
            tags={"output", "monitor"},
        )

@@ -308,16 +308,62 @@
 %22sigma%22)
+%0A    hyperparameter_names = (%22deterministic%22,)
 %0A%0A    de
@@ -408,24 +408,179 @@
  sigma_vw):%0A
+        deterministic = network.find_hyperparameter(%5B%22deterministic%22%5D, False)%0A        if deterministic:%0A            res = mu_vw.variable%0A        else:%0A    
         # TO
@@ -617,32 +617,36 @@
 d sigma%0A        
+ 
+   
 shape = mu_vw.sh
@@ -645,24 +645,28 @@
 mu_vw.shape%0A
+    
         if a
@@ -707,16 +707,20 @@
         
+    
 # NOTE: 
@@ -779,16 +779,20 @@
         
+  
+  
 # theano
@@ -817,16 +817,20 @@
 numbers%0A
+    
         
@@ -912,16 +912,20 @@
         
+ 
+   
 warnings
@@ -1005,16 +1005,20 @@
         
+   
+ 
 %22which c
@@ -1056,32 +1056,36 @@
 e%22)%0A            
+    
 shape = mu_vw.va
@@ -1101,24 +1101,28 @@
 ape%0A        
+  
+  
 # TODO save 
@@ -1160,16 +1160,20 @@
 the rng%0A
+    
         
@@ -1203,31 +1203,31 @@
 s()%0A        
-sampled
+    res
  = srng.norm
@@ -1484,23 +1484,19 @@
 ed_grad(
-sampled
+res
 ),%0A     

add deterministic flag for REINFORCE sampling