# Path: bionlp/utils/utils.py
# def theano_logsumexp(x, axis=None):
#     """
#     Compute log(sum(exp(x), axis=axis) in a numerically stable
#     fashion.
#     Parameters
#     ----------
#     x : tensor_like
#         A Theano tensor (any dimension will do).
#     axis : int or symbolic integer scalar, or None
#         Axis over which to perform the summation. `None`, the
#         default, performs over all axes.
#     Returns
#     -------
#     result : ndarray or scalar
#         The result of the log(sum(exp(...))) operation.
#     """
#     xmax = T.max(x,axis=axis, keepdims=True)
#     xmax_ = T.max(x,axis=axis)
#     return xmax_ + T.log(T.exp(x - xmax).sum(axis=axis))


from bionlp.utils.utils import theano_logsumexp
import lasagne,sys
import logging,theano
import theano.tensor as T
from __future__ import print_function

logging.basicConfig()
logger= logging.getLogger(__name__)

class CRFLayer(lasagne.layers.Layer):
    def __init__(self,incoming,Wp=lasagne.init.GlorotNormal(),**kwargs):
        super(CRFLayer,self).__init__(incoming, **kwargs)
        logger.info('Initializing CRF inference layer')
        self.num_labels=self.input_shape[-1] #getting the last dimension which is the number of labels.
        self.W= self.add_param(Wp, (self.num_labels,self.num_labels),name ='Wp')

    def get_CRF_params(self):
        return self.W

    def get_output_for(self,net_input,**kwargs):
        if 'unary' in kwargs and kwargs['unary']==True:
            return net_input

        logger.info('Initializing the messages')
        Wp=self.W
        unary_sequence = net_input.dimshuffle(1,0,2)    #Reshuffling the batched unary potential shape so that it can be used for word level iterations in theano.scan

        def forward_scan1(unary_sequence,forward_sm,Wp):
            forward_sm=forward_sm+unary_sequence
            forward_sm=theano_logsumexp(forward_sm.dimshuffle(0,1,'x')+Wp,1)
            return forward_sm

        def backward_scan1(unary_sequence,forward_sm,Wp):
            forward_sm=forward_sm+unary_sequence

# Path: bionlp/utils/utils.py
# def theano_logsumexp(x, axis=None):
#     """
#     Compute log(sum(exp(x), axis=axis) in a numerically stable
#     fashion.
#     Parameters
#     ----------
#     x : tensor_like
#         A Theano tensor (any dimension will do).
#     axis : int or symbolic integer scalar, or None
#         Axis over which to perform the summation. `None`, the
#         default, performs over all axes.
#     Returns
#     -------
#     result : ndarray or scalar
#         The result of the log(sum(exp(...))) operation.
#     """
#     xmax = T.max(x,axis=axis, keepdims=True)
#     xmax_ = T.max(x,axis=axis)
#     return xmax_ + T.log(T.exp(x - xmax).sum(axis=axis))


# Path: bionlp/taggers/rnn_feature/networks/crf_lstm_layer.py
from bionlp.utils.utils import theano_logsumexp
import lasagne,sys
import logging,theano
import theano.tensor as T

from __future__ import print_function

logging.basicConfig()
logger= logging.getLogger(__name__)

class CRFLayer(lasagne.layers.Layer):
    def __init__(self,incoming,Wp=lasagne.init.GlorotNormal(),**kwargs):
        super(CRFLayer,self).__init__(incoming, **kwargs)
        logger.info('Initializing CRF inference layer')
        self.num_labels=self.input_shape[-1] #getting the last dimension which is the number of labels.
        self.W= self.add_param(Wp, (self.num_labels,self.num_labels),name ='Wp')

    def get_CRF_params(self):
        return self.W

    def get_output_for(self,net_input,**kwargs):
        if 'unary' in kwargs and kwargs['unary']==True:
            return net_input

        logger.info('Initializing the messages')
        Wp=self.W
        unary_sequence = net_input.dimshuffle(1,0,2)    #Reshuffling the batched unary potential shape so that it can be used for word level iterations in theano.scan

        def forward_scan1(unary_sequence,forward_sm,Wp):
            forward_sm=forward_sm+unary_sequence
            forward_sm=theano_logsumexp(forward_sm.dimshuffle(0,1,'x')+Wp,1)
            return forward_sm

        def backward_scan1(unary_sequence,forward_sm,Wp):
            forward_sm=forward_sm+unary_sequence

            forward_sm=theano_logsumexp(forward_sm.dimshuffle(0,1,'x')+Wp.T,1)