# Path: nn_layers.py
# def per_word_averaging_layer(dist, dist_mask):
#     """
#     :param proj: Output of the LSTM layer
#     :param dist_mask: 4D index tensor
#     :return: The per-word averages.
#     """
# 
#     dist = dist * dist_mask
# 
#     dist = dist.dimshuffle(1, 2, 0, 3)
#     dist_mask = tensor.cast(dist_mask, theano.config.floatX)
#     dist_mask = dist_mask.dimshuffle(1, 2, 0, 3)
#     divider = dist_mask.sum(axis=0)
#     divider += tensor.eq(divider, numpy_floatX(0.0))
# 
#     dist = dist * dist_mask
#     tmp = tensor.cast(dist.sum(axis=0), theano.config.floatX)
#     tmp /= divider
#     return tmp.dimshuffle(1, 0, 2)


import unittest
import theano
import theano.tensor
import numpy
import logging
from nn_layers import per_word_averaging_layer
        # n_chars -> maximum character sequence length (e.g. 140)
        # n_samples = minibatch size
        # n_proj = number of combined word/character embeddings
        n_chars, n_samples, n_proj = 4, 2, 4
        # Allocate a matrix which represents the LSTM layer output
        L = numpy.zeros((maxw, maxwlen, n_samples, n_proj))
        L[0, 0, 0, :] = [0.2, 0.3, 0.4, 0.5]
        L[0, 1, 0, :] = [-0.2, -0.4, 0.6, 0.7]
        L[1, 0, 0, :] = [0.8, -0.2, 0.3, -0.1]
        L[0, 0, 1, :] = [-0.4, 0.2, 0.8, 0.9]

        W = numpy.zeros((maxw, maxwlen, n_samples, n_proj))
        W[0, 0, 0, :] = numpy.ones(n_proj) # 1st word, 1st character, 1st tweet
        W[0, 1, 0, :] = numpy.ones(n_proj) # 1st word, 2nd character, 1st tweet
        W[1, 0, 0, :] = numpy.ones(n_proj) # 2nd word, 3rd character, 1st tweet
        W[0, 0, 1, :] = numpy.ones(n_proj) # 1st word, 1st character, 2nd tweet

        # The expected output
        O = numpy.zeros((n_samples, 2, n_proj))
        O[0, 0, :] = [0.0, -0.05, 0.50, 0.60]  # First word, first tweet
        O[1, 0, :] = [0.8, -0.2, 0.3, -0.1]  # Second word, first tweet
        O[0, 1, :] = [-0.4, 0.2, 0.8, 0.9]  # First word, second tweet

        return n_chars, n_samples, n_proj, L, W, O

    def test_forward(self):
        n_chars, n_samples, n_proj, L, W, O = WordAveragingOpTests.get_std()
        LArg = theano.tensor.dtensor4()
        WArg = theano.tensor.dtensor4()


# Path: nn_layers.py
# def per_word_averaging_layer(dist, dist_mask):
#     """
#     :param proj: Output of the LSTM layer
#     :param dist_mask: 4D index tensor
#     :return: The per-word averages.
#     """
# 
#     dist = dist * dist_mask
# 
#     dist = dist.dimshuffle(1, 2, 0, 3)
#     dist_mask = tensor.cast(dist_mask, theano.config.floatX)
#     dist_mask = dist_mask.dimshuffle(1, 2, 0, 3)
#     divider = dist_mask.sum(axis=0)
#     divider += tensor.eq(divider, numpy_floatX(0.0))
# 
#     dist = dist * dist_mask
#     tmp = tensor.cast(dist.sum(axis=0), theano.config.floatX)
#     tmp /= divider
#     return tmp.dimshuffle(1, 0, 2)


# Path: test_nn_layers.py
import unittest
import theano
import theano.tensor
import numpy
import logging
from nn_layers import per_word_averaging_layer

        # n_chars -> maximum character sequence length (e.g. 140)
        # n_samples = minibatch size
        # n_proj = number of combined word/character embeddings
        n_chars, n_samples, n_proj = 4, 2, 4
        # Allocate a matrix which represents the LSTM layer output
        L = numpy.zeros((maxw, maxwlen, n_samples, n_proj))
        L[0, 0, 0, :] = [0.2, 0.3, 0.4, 0.5]
        L[0, 1, 0, :] = [-0.2, -0.4, 0.6, 0.7]
        L[1, 0, 0, :] = [0.8, -0.2, 0.3, -0.1]
        L[0, 0, 1, :] = [-0.4, 0.2, 0.8, 0.9]

        W = numpy.zeros((maxw, maxwlen, n_samples, n_proj))
        W[0, 0, 0, :] = numpy.ones(n_proj) # 1st word, 1st character, 1st tweet
        W[0, 1, 0, :] = numpy.ones(n_proj) # 1st word, 2nd character, 1st tweet
        W[1, 0, 0, :] = numpy.ones(n_proj) # 2nd word, 3rd character, 1st tweet
        W[0, 0, 1, :] = numpy.ones(n_proj) # 1st word, 1st character, 2nd tweet

        # The expected output
        O = numpy.zeros((n_samples, 2, n_proj))
        O[0, 0, :] = [0.0, -0.05, 0.50, 0.60]  # First word, first tweet
        O[1, 0, :] = [0.8, -0.2, 0.3, -0.1]  # Second word, first tweet
        O[0, 1, :] = [-0.4, 0.2, 0.8, 0.9]  # First word, second tweet

        return n_chars, n_samples, n_proj, L, W, O

    def test_forward(self):
        n_chars, n_samples, n_proj, L, W, O = WordAveragingOpTests.get_std()
        LArg = theano.tensor.dtensor4()
        WArg = theano.tensor.dtensor4()


        out = per_word_averaging_layer(LArg, WArg)