import nose.tools as nt
import numpy as np
import theano
import theano.tensor as T

import treeano
import treeano.nodes as tn
from treeano.sandbox.nodes import spatial_attention

fX = theano.config.floatX


def test_spatial_feature_point_node():
    network = tn.SequentialNode(
        "s",
        [tn.InputNode("i", shape=(2, 2, 2, 3)),
         spatial_attention.SpatialFeaturePointNode("fp")]
    ).network()

    fn = network.function(["i"], ["s"])

    x = np.zeros((2, 2, 2, 3), dtype=fX)
    idxs = np.array([[[0, 0],
                      [1, 0]],
                     [[0, 1],
                      [1, 2]]],
                    dtype=fX)
    ans = idxs / np.array([1, 2], dtype=fX)[None, None]
    for batch in range(2):
        for channel in range(2):
            i, j = idxs[batch, channel]
            x[batch, channel, i, j] = 1

    np.testing.assert_allclose(ans,
                               fn(x)[0],
                               rtol=1e-5,
                               atol=1e-8)

@@ -1001,8 +1001,456 @@
 l=1e-8)%0A
+%0A%0Adef test_pairwise_distance_node():%0A    # NOTE: only tests shape calculation%0A    network = tn.SequentialNode(%0A        %22s%22,%0A        %5Btn.InputNode(%22i%22, shape=(2, 2, 2, 3)),%0A         spatial_attention.SpatialFeaturePointNode(%22fp%22),%0A         spatial_attention.PairwiseDistanceNode(%22pd%22)%5D%0A    ).network()%0A%0A    fn = network.function(%5B%22i%22%5D, %5B%22s%22%5D)%0A    x = np.zeros((2, 2, 2, 3), dtype=fX)%0A%0A    nt.assert_equal((2, 4),%0A                    fn(x)%5B0%5D.shape)%0A

add pairwise distance test