from __future__ import division, absolute_import
from __future__ import print_function, unicode_literals

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

import treeano
import treeano.nodes as tn

fX = theano.config.floatX


def test_pool_output_shape_2d():
    def test_same(input_shape, local_sizes, strides, pads, ignore_border):
        res = tn.downsample.pool_output_shape(
            input_shape,
            (2, 3),
            local_sizes,
            strides,
            pads,
            ignore_border,
        )
        from theano.tensor.signal.downsample import max_pool_2d
        ans = max_pool_2d(
            T.constant(np.random.randn(*input_shape).astype(fX)),
            ds=local_sizes,
            st=strides,
            ignore_border=ignore_border,
            padding=pads,
        ).shape.eval()
        print(ans, res)
        np.testing.assert_equal(ans, res)

    # tests w/ ignore border
    test_same((1, 1, 5, 6), (2, 3), (1, 1), (0, 0), True)
    test_same((1, 1, 5, 6), (2, 3), (2, 2), (0, 0), True)
    test_same((1, 1, 1, 1), (2, 3), (2, 2), (0, 0), True)
    test_same((1, 1, 5, 6), (2, 3), (1, 1), (0, 1), True)
    test_same((1, 1, 5, 6), (2, 3), (2, 2), (1, 0), True)
    test_same((1, 1, 1, 1), (2, 3), (2, 2), (1, 1), True)

    # tests w/o ignore border, and stride <= pool_size
    test_same((1, 1, 5, 6), (2, 3), (1, 1), (0, 0), False)
    test_same((1, 1, 5, 6), (2, 3), (2, 2), (0, 0), False)
    test_same((1, 1, 1, 1), (2, 3), (2, 2), (0, 0), False)

    # tests w/o ignore border, and stride > pool_size
    test_same((1, 1, 5, 6), (2, 3), (3, 3), (0, 0), False)
    test_same((1, 1, 5, 6), (2, 3), (3, 3), (0, 0), False)
    test_same((1, 1, 1, 1), (2, 3), (3, 3), (0, 0), False)


def test_pool_output_shape_3d():
    def test_same(input_shape, local_sizes, strides, pads, ignore_border, ans):
        res = tn.downsample.pool_output_shape(
            input_shape,
            (2, 3, 4),
            local_sizes,
            strides,
            pads,
            ignore_border,