def Softmax(x, temp = 1):
    """ 
    Softmax Units. 

    Applies row-wise softmax  to the input supplied.
            
    Args:
        x: could be a ``theano.tensor`` or a ``theano.shared`` or ``numpy`` arrays or 
            ``python lists``.
        temp: temperature of type ``float``. Mainly used during distillation, normal 
                softmax prefer ``T=1``. 
    Notes:
        Refer [3] for details.
    
        .. [#]  Hinton, Geoffrey, Oriol Vinyals, and Jeff Dean. "Distilling the knowledge in
                a neural network." arXiv preprint arXiv:1503.02531 (2015).       
                        
    Returns: 
        same as input: returns a row-wise softmax output of the same shape as the input.
    """
    if temp != 1:
        expo = T.exp(x / float(temp)) # at this moment this is mini_batch_size X num_classes.
        normalizer = T.sum(expo,axis=1,keepdims=True)  # at this moment this is mini_batch_size X 1.
        T.addbroadcast(normalizer,1)    
        return expo / normalizer
    else:
        return T.nnet.softmax(x)               

    def fprop(self, x, mode='train'):
        if mode == 'test':
            # this is for use during test/validation time
            x_avg = self.params.getParameter('x_avg')
        elif mode == 'calculate':
            x_avg = x.mean(self.norm_axis, keepdims=True)
        elif mode == 'train':
            # otherwise calculate the batch mean and std
            x_avg = x.mean(self.norm_axis, keepdims=True)
            
            # the following trick is learend from lasagne implementation
            running_mean = theano.clone(self.params.getParameter('x_avg'), share_inputs=False)
            
            running_mean_udpate = ((1 - self.alpha) * running_mean
                                    +self.alpha * x_avg)

 
            # set a default update for them
            running_mean.default_update = running_mean_udpate
 
            x_avg += 0 * running_mean
        else:
            raise "mode can only take ['train', 'test', 'calculate']"
        
        self.x_avg = x_avg
        x_avg = T.addbroadcast(x_avg, *self.norm_axis)
        beta = T.addbroadcast(self.params.getParameter('beta'), *self.norm_axis)
        
        bn_x = x / (x_avg + 1e-18) * beta
        return bn_x if self.actFunc is None else self.actFunc(bn_x)


# End BatchExpNormLayer
#-------------------------------------------------------------------------------

    def output(self, input):

        W_shuffled = self.W.dimshuffle(3, 0, 1, 2)  # c01b to bc01

        print "input ndim", input.ndim

        conv_out = dnn.dnn_conv(img=input,
                                        kerns=W_shuffled,
                                        subsample=(self.stride, self.stride),
                                        border_mode=self.padsize)

        conv_out = conv_out + self.b.dimshuffle('x', 0, 'x', 'x')

        if self.batch_norm:
            conv_out = (conv_out - T.mean(conv_out, axis = (0,2,3), keepdims = True)) / (1.0 + T.std(conv_out, axis=(0,2,3), keepdims = True))
            conv_out = conv_out * T.addbroadcast(self.bn_std,0,2,3) + T.addbroadcast(self.bn_mean, 0,2,3)

        self.out_store = conv_out

        if self.activation == "relu":
            self.out = T.maximum(0.0, conv_out)
        elif self.activation == "tanh":
            self.out = T.tanh(conv_out)
        elif self.activation == None:
            self.out = conv_out


        return T.specify_shape(self.out, (self.batch_size, self.out_channels, self.in_length / self.stride, self.in_length / self.stride))

    def get_state(self):
        st = super(LatentTypeWithTuningCurve, self).get_state()

        # The filters are non-identifiable as we can negate both the
        # temporal and the spatial filters and get the same net effect.
        # By convention, choose the sign that results in the most
        # positive temporal filter.
        sign = T.sgn(T.sum(self.stim_resp_t, axis=0))
        T.addbroadcast(sign, 0)

        # Similarly, we can trade a constant between the spatial and temporal
        # pieces. By convention, set the temporal filter to norm 1.
        Z = T.sqrt(T.sum(self.stim_resp_t**2, axis=0))
        T.addbroadcast(Z, 0)

        # Compute the normalized temporal response
        stim_resp_t = sign*(1.0/Z)*self.stim_resp_t

        # Finally, reshape the spatial component as necessary
        if self.spatial_ndim == 2:
            stim_resp_x = sign*Z*self.stim_resp_x
            stim_resp_x = T.reshape(stim_resp_x,
                                    self.spatial_shape + (self.R,))
        else:
            stim_resp_x = sign*Z*self.stim_resp_x

        st.update({'stim_response_x' : stim_resp_x,
                   'stim_response_t' : stim_resp_t})

        return st

	def resample_step(self):
		
		idx=self.theano_rng.multinomial(pvals=T.reshape(self.weights_now,(1,self.npcl))).T
		s_samp=T.sum(self.s_now*T.addbroadcast(idx,1),axis=0)
		h_samp=T.sum(self.h_now*T.addbroadcast(idx,1),axis=0)
		
		return T.cast(s_samp,'float32'), T.cast(h_samp,'float32')

    def output(self, input_raw):

        if self.flatten_input:
            input = input_raw.flatten(2)
        else:
            input = input_raw

        lin_output = T.dot(input, self.W) + self.b

        if self.batch_norm:
            lin_output = (lin_output - T.mean(lin_output, axis = 0, keepdims = True)) / (1.0 + T.std(lin_output, axis = 0, keepdims = True))
            lin_output = (lin_output * T.addbroadcast(self.bn_std,0) + T.addbroadcast(self.bn_mean,0))

        self.out_store = lin_output

        if self.activation == None: 
            activation = lambda x: x
        elif self.activation == "relu": 
            activation = lambda x: T.maximum(0.0, x)
        elif self.activation == "exp": 
            activation = lambda x: T.exp(x)
        elif self.activation == "tanh":
            activation = lambda x: T.tanh(x)
        elif self.activation == 'softplus':
            activation = lambda x: T.nnet.softplus(x)
        else: 
            raise Exception("Activation not found")

        out = activation(lin_output)

        #if self.residual:
        #    return out + input_raw
        #else:
        #    return out
        return out

    def output(self, input):

        if self.unflatten_input != None:
            input = T.reshape(input, self.unflatten_input)

        conv_out = deconv(input, self.W, subsample=(2, 2), border_mode=(2,2))

        conv_out = conv_out + self.b.dimshuffle('x', 0, 'x', 'x')

        if self.batch_norm:
            conv_out = (conv_out - conv_out.mean(axis = (0,2,3), keepdims = True)) / (1.0 + conv_out.std(axis = (0,2,3), keepdims = True))
            conv_out = conv_out * T.addbroadcast(self.bn_std,0,2,3) + T.addbroadcast(self.bn_mean,0,2,3)

        if self.activation == "relu":
            out = T.maximum(0.0, conv_out)
        elif self.activation == "tanh":
            out = T.tanh(conv_out)
        elif self.activation == None:
            out = conv_out
        else:
            raise Exception()


        self.params = {'W' : self.W, 'b' : self.b}
        if self.batch_norm:
            self.params["mu"] = self.bn_mean
            self.params["sigma"] = self.bn_std

        return out

def keep_max(input, theta, k, sent_mask):
    sig_input = T.nnet.sigmoid(T.dot(input, theta))
    sent_mask = sent_mask.dimshuffle(0, 'x', 1, 'x')
    sig_input = sig_input * sent_mask
    #sig_input = T.dot(input, theta)
    if k == 0:
        result = input * T.addbroadcast(sig_input, 3)
        return result, sig_input

    # get the sorted idx
    sort_idx = T.argsort(sig_input, axis=2)
    k_max_ids = sort_idx[:,:,-k:,:]
    dim0, dim1, dim2, dim3 = k_max_ids.shape
    batchids = T.repeat(T.arange(dim0), dim1*dim2*dim3)
    mapids = T.repeat(T.arange(dim1), dim2*dim3).reshape((1, dim2*dim3))
    mapids = T.repeat(mapids, dim0, axis=0).flatten()
    rowids = k_max_ids.flatten()
    colids = T.arange(dim3).reshape((1, dim3))
    colids = T.repeat(colids, dim0*dim1*dim2, axis=0).flatten()
    sig_mask = T.zeros_like(sig_input)
    choosed = sig_input[batchids, mapids, rowids, colids]
    sig_mask = T.set_subtensor(sig_mask[batchids, mapids, rowids, colids], 1)
    input_mask = sig_mask * sig_input
    result = input * T.addbroadcast(input_mask, 3)
    return result, sig_input

 def get_output_for(self, input, deterministic=False, **kwargs):
     if deterministic:
         # use stored mean and std
         mean = self.mean
         std = self.std
     else:
         # use this batch's mean and std
         mean = input.mean(self.axes, keepdims=True)
         std = input.std(self.axes, keepdims=True)
         # and update the stored mean and std:
         # we create (memory-aliased) clones of the stored mean and std
         running_mean = theano.clone(self.mean, share_inputs=False)
         running_std = theano.clone(self.std, share_inputs=False)
         # set a default update for them
         running_mean.default_update = (1 - self.alpha) * running_mean + self.alpha * mean
         running_std.default_update = (1 - self.alpha) * running_std + self.alpha * std
         # and include them in the graph so their default updates will be
         # applied (although the expressions will be optimized away later)
         mean += 0 * running_mean
         std += 0 * running_std
     std += self.epsilon
     mean = T.addbroadcast(mean, *self.axes)
     std = T.addbroadcast(std, *self.axes)
     beta = T.addbroadcast(self.beta, *self.axes)
     gamma = T.addbroadcast(self.gamma, *self.axes)
     normalized = (input - mean) * (gamma / std) + beta
     return self.nonlinearity(normalized)

def keep_max(input, theta, k):
    """
    :type input: theano.tensor.tensor4
    :param input: the input data
                
    :type theta: theano.tensor.matrix
    :param theta: the parameter for sigmoid function
                            
    :type k: int 
    :param k: the number k used to define top k sentence to remain
    """
    sig_input = T.nnet.sigmoid(T.dot(input, theta))
    if k == 0: # using all the sentences
        result = input * T.addbroadcast(sig_input, 3)
        return result, sig_input

    # get the sorted idx
    sort_idx = T.argsort(sig_input, axis=2)
    k_max_ids = sort_idx[:,:,-k:,:]
    dim0, dim1, dim2, dim3 = k_max_ids.shape
    batchids = T.repeat(T.arange(dim0), dim1*dim2*dim3)
    mapids = T.repeat(T.arange(dim1), dim2*dim3).reshape((1, dim2*dim3))
    mapids = T.repeat(mapids, dim0, axis=0).flatten()
    rowids = k_max_ids.flatten()
    colids = T.arange(dim3).reshape((1, dim3))
    colids = T.repeat(colids, dim0*dim1*dim2, axis=0).flatten()
    # construct masked data
    sig_mask = T.zeros_like(sig_input)
    choosed = sig_input[batchids, mapids, rowids, colids]
    sig_mask = T.set_subtensor(sig_mask[batchids, mapids, rowids, colids], 1)

    input_mask = sig_mask * sig_input
    result = input * T.addbroadcast(input_mask, 3)
    return result, sig_input

    def log_p(self, L):
        """ Compute log prob of the given value under this prior
            Input: L ~ NxD
        """
        assert L.ndim == 2, "L must be 2d!"
        # Compute pairwise L2 norm
        L1 = L.dimshuffle(0,'x',1)     # Nx1xD
        L2 = L.dimshuffle('x',0,1)     # 1xNxD
        T.addbroadcast(L1,1)
        T.addbroadcast(L2,0)

        # Compute pairwise distances
        D = ((L1-L2)**2).sum(axis=2)

        # Compute the kernel
        K = T.exp(-D / self.sigma**2)

        # Log prob is the log determinant of the pairwise distances
        lp_det = T.log(self.d(K))

        # Also multiply by a spherical Gaussian with standard deviation of 'bound'
        # to prevent points from diverging to infinity
        lp_gauss = self.gaussian.log_p(L)

        return lp_det + lp_gauss

    def output(self, input):

        if self.unflatten_input != None:
            input = T.reshape(input, self.unflatten_input)

        W_shuffled = self.W.val.dimshuffle(3, 0, 1, 2)  # c01b to bc01

        conv_out = dnn.dnn_conv(img=input,
                                        kerns=W_shuffled,
                                        subsample=(self.convstride, self.convstride),
                                        border_mode=self.padsize)

        conv_out = conv_out + self.b.val.dimshuffle('x', 0, 'x', 'x')

        if self.batch_norm:
            conv_out = (conv_out - T.mean(conv_out, axis = (0,2,3), keepdims = True)) / (1.0 + T.std(conv_out, axis=(0,2,3), keepdims = True))
            conv_out = conv_out * T.addbroadcast(self.bn_std,0,2,3) + T.addbroadcast(self.bn_mean, 0,2,3)

        self.out_store = conv_out

        if self.activation == "relu":
            self.out = T.maximum(0.0, conv_out)
        elif self.activation == "tanh":
            self.out = T.tanh(conv_out)
        elif self.activation == None:
            self.out = conv_out

        #if self.residual:
        #    print "USING RESIDUAL"
        #    self.out += input


        return self.out

def local_contrast_normalize(X, window, img_shape):
    """Return normalized X and the convolution transform
    """
    batchsize, channels, R, C = img_shape
    assert window.shape[0] == 1
    assert window.shape[1] == channels
    N = window.shape[2]
    assert window.shape[3] == N
    blur = tlinear.Conv2d(
            filters=sharedX(window, 'LCN_window'),
            img_shape=img_shape,
            border_mode='full')
    N2 = N//2
    # remove global mean
    X = X - X.mean(axis=[1, 2, 3]).dimshuffle(0, 'x', 'x', 'x')

    #remove local mean
    blurred_x = tensor.addbroadcast(blur.lmul(X), 1)
    x2c = X - blurred_x[:, :, N2:R + N2, N2:C + N2]

    # standardize contrast
    blurred_x2c_sqr = tensor.addbroadcast(blur.lmul(x2c ** 2), 1)
    x2c_lcn =  x2c / tensor.sqrt((10 + blurred_x2c_sqr[:, :, N2:R + N2, N2:C + N2]))

    return x2c_lcn, blur

    def _ct(self, other):
        ''' Helper function to make tensors dimensions compatible'''
        if (other.var_set == self.var_set):
            return (self.pt_tensor, other.pt_tensor)
        union_var_set = other.scope.union(self.scope)
        vidx1 = frozenset(self.var_indices)
        vidx2 = frozenset(other.var_indices)
        union_indices = vidx1.union(vidx2)

        shape1 = []
        shape2 = []
        b1 = []
        b2 = []
        u1 = []
        u2 = []

        for i,vidx in enumerate(sorted(union_indices)):
            if (vidx in vidx1):
                shape1.append(self.discrete_pgm.cardinalities[vidx])
                u1.append(i)
            else:
                shape1.append(1)
                b1.append(i)
            if (vidx in vidx2):
                shape2.append(self.discrete_pgm.cardinalities[vidx])
                u2.append(i)
            else:
                shape2.append(1)
                b2.append(i)
        t1 = T.addbroadcast(T.unbroadcast(self.pt_tensor.reshape(shape1, len(shape1)), *u1), *b1)
        t2 = T.addbroadcast(T.unbroadcast(other.pt_tensor.reshape(shape2, len(shape2)), *u2), *b2)
        return (t1, t2)

    def embed(self,x, y, kth):

        hidden = self.hidden_k(x,self.superw,self.dicw, kth)
        size = y.ndim

        y = T.addbroadcast(y,size - 1)
        embedding = T.sum(hidden*y,0)/T.addbroadcast(T.cast(T.sum(y,0), 'int16'), size - 2)
        return embedding

    def __Theano_build__(self):
        Td = T.tensor3('Td')
        Ty = T.ivector('Ty')
        Tlr = T.scalar('Tlr')
        #Talpha = T.TensorType(dtype='float32', broadcastable=(0, 1, 1))('alpha')
        A = theano.shared(np.ones((self.D.shape[0]))\
                .astype('float32').reshape(-1, 1, 1), 'A')
        Ttriple   = [T.ivector('triple'+x) for x in ['i', 'j', 'l']]
        Tneighbor = [T.ivector('neighbor'+x) for x in ['i', 'j']]

        d = (Td * T.addbroadcast(A, 1, 2)).sum(0)
        pull_error, _ = theano.scan(
            fn = lambda i, j, d: d[i, j],
            sequences=[Tneighbor[0], Tneighbor[1]],
            outputs_info=None,
            non_sequences=[d])
        pull_error = pull_error.sum()

        push_error, _ = theano.scan(
            fn = lambda i, j, l, d: T.neq(Ty[i], Ty[l]) * T.maximum((d[i]-d[j]) - (d[i]-d[l]) +1, 0),
            sequences=[Ttriple[0], Ttriple[1], Ttriple[2]],
            outputs_info=None,
            non_sequences=[d])
#       zerocount = T.eq(linalg.diag(mask*T.maximum(lossij - lossil + 1, 0)), 0).sum()

        error = pull_error.sum() + push_error.sum()
        grad = T.grad(error, A)
        newA = A - Tlr*grad #T.maximum(A - Tlr*grad, 0)
        updates = [(A, newA/newA.sum())]

        self.Ttrain = theano.function(Ttriple+Tneighbor+[Tlr], 
                Tlr*grad,
                givens={Td: self.D, 
                        Ty: self.y},
                updates=updates,
                allow_input_downcast=True,
                on_unused_input='warn')

        self.Tloss = theano.function(Ttriple+Tneighbor, 
                error,
                givens={Td: self.D, Ty: self.y},
                allow_input_downcast=True)


#       eig, eigv = linalg.eig((d+d.T)/2.0)
#       self.Tmineig = theano.function([], 
#               T.min(eig),
#               givens={Td: self.D},
#               allow_input_downcast=True)

        self.Tmindist = theano.function([], 
                T.min(d),
                givens={Td: self.D},
                allow_input_downcast=True)


        self.Ttransform = theano.function([Td], (Td*T.addbroadcast(A, 1, 2)).sum(0), allow_input_downcast=True)
        self.TA = A

    def __init__(self, model, latent):
        """ Initialize the stochastic block model for the adjacency matrix
        """
        self.model = model
        self.prms = model['network']['graph']
        self.N = model['N']
        self.N_dims = self.prms['N_dims']

        # Get the latent location
        self.location = latent[self.prms['locations']]
        self.Lm = self.location.Lm
        # self.location_prior = create_prior(self.prms['location_prior'])
        #
        # # Latent distance model has NxR matrix of locations L
        # self.L = T.dvector('L')
        # self.Lm = T.reshape(self.L, (self.N, self.N_dims))

        # Compute the distance between each pair of locations
        # Reshape L into a Nx1xD matrix and a 1xNxD matrix, then add the requisite
        # broadcasting in order to subtract the two matrices
        L1 = self.Lm.dimshuffle(0,'x',1)     # Nx1xD
        L2 = self.Lm.dimshuffle('x',0,1)     # 1xNxD
        T.addbroadcast(L1,1)
        T.addbroadcast(L2,0)
        #self.D = T.sqrt(T.sum((L1-L2)**2, axis=2))
        #self.D = T.sum((L1-L2)**2, axis=2)

        # It seems we need to use L1 norm for now because
        # Theano doesn't properly compute the gradients of the L2
        # norm. (It gives NaNs because it doesn't realize that some
        # terms will cancel out)
        # self.D = (L1-L2).norm(1, axis=2)
        self.D = T.pow(L1-L2,2).sum(axis=2)

        # There is a distance scale, \delta
        self.delta = T.dscalar(name='delta')

        # Define complete adjacency matrix
        self.A = T.bmatrix('A')

        # The probability of A is exponentially decreasing in delta
        # self.pA = T.exp(-1.0*self.D/self.delta)
        self.pA = T.exp(-0.5*self.D/self.delta**2)

        if 'rho_refractory' in self.prms:
            self.pA += T.eye(self.N) * (self.prms['rho_refractory']-self.pA)
            # self.pA[np.diag_indices(self.N)] = self.prms['rho_refractory']

        # Allow for scaling the log likelihood of the graph so that we can do
        # Annealed importance sampling
        self.lkhd_scale = theano.shared(value=1.0, name='lkhd_scale')

        # Define log probability
        self.lkhd = T.sum(self.A * T.log(self.pA) + (1 - self.A) * T.log(1 - self.pA))
        # self.log_p = self.lkhd_scale * self.lkhd + self.location_prior.log_p(self.Lm)
        self.log_p = self.lkhd_scale * self.lkhd

    def __init__(self, model):
        self.model = model
        self.imp_model = model['impulse']

        # Number of presynaptic neurons
        self.N = model['N']

        # Get parameters of the prior
        self.alpha = self.imp_model['alpha']

        # Create a basis for the impulse responses response
        self.basis = create_basis(self.imp_model['basis'])
        (_,self.B) = self.basis.shape
        # The basis is interpolated once the data is specified
        self.initialize_basis()

        # Initialize memory for the filtered spike train
        self.ir = theano.shared(name='ir',
                                value=np.zeros((1,self.N,self.B)))

        # Define Dirichlet distributed weights by normalizing gammas
        # The variables are log-gamma distributed
        self.lng = T.dvector('w_lng')
        self.g = T.exp(self.lng)
        self.g2 = T.reshape(self.g, [self.N,self.B])
        self.g_sum = T.reshape(T.sum(self.g2, axis=1), [self.N,1])
        
        # Normalize the gammas to get a Dirichlet draw
        T.addbroadcast(self.g_sum, 1)
        self.w_ir2 = self.g2 / self.g_sum
        self.w_ir2.name = 'w_ir'

        # Repeat them (in a differentiable manner) to create a 3-tensor
        self.w_ir3 = T.reshape(self.w_ir2, [1,self.N,self.B])

        # Make w_ir3 broadcastable in the 1st dim
        T.addbroadcast(self.w_ir3,0)

        # Take the elementwise product of the filtered stimulus and
        # the repeated weights to get the weighted impulse current along each
        # impulse basis dimension. Then sum over bases to get the
        # total coupling current from each presynaptic neurons at
        # all time points
        self.I_imp = T.sum(self.ir*self.w_ir3, axis=2)

        # Log probability of a set of independent log-gamma r.v.'s
        # This is log p(log(g)) under the prior. Since we are taking the
        # log, we multiply by a factor of g to ensure normalization and
        # thus the \alpha-1 in the exponent becomes \alpha
        self.log_p = -self.B*self.N*scipy.special.gammaln(self.alpha) \
                     + T.sum(self.alpha*self.lng) \
                     - T.sum(self.g)

        # Define a helper variable for the impulse response
        # after projecting onto the basis
        self.impulse = T.dot(self.w_ir2, T.transpose(self.ibasis))

 def get_t_weights(self, t):
     """
     Generate vector of weights allowing selection of current timestep.
     (if t is not an integer, the weights will linearly interpolate)
     """
     n_seg = self.trajectory_length
     t_compare = T.arange(n_seg, dtype=theano.config.floatX).reshape((1,n_seg))
     diff = abs(T.addbroadcast(t,1) - T.addbroadcast(t_compare,0))
     t_weights = T.max(T.join(1, (-diff+1).reshape((n_seg,1)), T.zeros((n_seg,1))), axis=1)
     return t_weights.reshape((-1,1))

 def node_update(self, s_, h_, m_) :
     """
     Update params in Attention.
     """
     preact = tensor.dot(h_, self.params[self._p(self.prefix, 'U')])
     preact += tensor.addbroadcast(tensor.dot(s_, self.params[self._p(self.prefix, 'W')]).dimshuffle('x', 0, 1), 0)
     preact = tensor.dot(tensor.tanh(preact), self.params[self._p(self.prefix, 'va')]) * m_
     alpha = tensor.nnet.softmax(preact.dimshuffle(1, 0)).dimshuffle(1, 0, 'x')
     c = (h_ * tensor.addbroadcast(alpha, 2)).sum(axis=0) # c is (samples,2*hidden)
 
     return c, alpha