import unittest

import numpy
import six

import chainer
from chainer import cuda
from chainer import functions
from chainer import gradient_check
from chainer import testing
from chainer.testing import attr
from chainer.testing import condition


@testing.parameterize(*testing.product({
    'cover_all': [True, False],
}))
class TestUnpooling2D(unittest.TestCase):

    def setUp(self):
        self.N = 2
        self.n_channels = 3
        inh, inw = 2, 1
        self.x = numpy.arange(
            self.N * self.n_channels * inh * inw, dtype=numpy.float32)\
            .reshape(self.N, self.n_channels, inh, inw)
        numpy.random.shuffle(self.x)
        self.x = 2 * self.x / self.x.size - 1

        self.ksize = 2
        outh, outw = self.outsize = (4, 2)
        self.gy = numpy.random.uniform(
            -1, 1, (self.N, self.n_channels, outh, outw)).astype(numpy.float32)

    def check_forward(self, x_data):
        x = chainer.Variable(x_data)
        y = functions.unpooling_2d(x, self.ksize, outsize=self.outsize,
                                   cover_all=self.cover_all)
        self.assertEqual(y.data.dtype, numpy.float32)
        y_data = cuda.to_cpu(y.data)

        self.assertEqual(self.gy.shape, y_data.shape)
        for i in six.moves.range(self.N):
            for c in six.moves.range(self.n_channels):
                expect = numpy.array([
                    [self.x[i, c, 0, 0], self.x[i, c, 0, 0]],
                    [self.x[i, c, 0, 0], self.x[i, c, 0, 0]],
                    [self.x[i, c, 1, 0], self.x[i, c, 1, 0]],
                    [self.x[i, c, 1, 0], self.x[i, c, 1, 0]],
                ])
                gradient_check.assert_allclose(expect, y_data[i, c])

    @condition.retry(3)
    def test_forward_cpu(self):
        self.check_forward(self.x)

    @attr.gpu
    @condition.retry(3)
    def test_forward_gpu(self):
        self.check_forward(cuda.to_gpu(self.x))

    def check_backward(self, x_data, y_grad):
        gradient_check.check_backward(
            functions.Unpooling2D(self.ksize, outsize=self.outsize,
                                  cover_all=self.cover_all),
            x_data, y_grad)

    @condition.retry(3)
    def test_backward_cpu(self):
        self.check_backward(self.x, self.gy)

    @attr.gpu
    @condition.retry(3)
    def test_backward_gpu(self):
        self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.gy))


testing.run_module(__name__, __file__)

@@ -267,31 +267,286 @@
 ize(
-*testing.product(%7B
+%0A    # we assume insize as (2, 1)%0A    # standard output size which is estimated with get_deconv_outsize function%0A    %7B'cover_all': False, 'outsize': (4, 2)%7D,%0A    %7B'cover_all': True, 'outsize': (3, 1)%7D,%0A    %7B'cover_all': False, 'outsize': None, 'expected_outsize': (4, 2)%7D,
 %0A    
+%7B
 'cov
@@ -558,25 +558,207 @@
 l': 
-%5B
 True, 
-False%5D,%0A%7D)
+'outsize': None, 'expected_outsize': (3, 1)%7D,%0A    # another sizes which can be outsize of insize (2, 1)%0A    %7B'cover_all': False, 'outsize': (5, 2)%7D,%0A    %7B'cover_all': True, 'outsize': (4, 2)%7D,%0A
 )%0Acl
@@ -1194,16 +1194,32 @@
 ize 
-= (4, 2)
+or self.expected_outsize
 %0A   
@@ -1786,16 +1786,395 @@
 nnels):%0A
+                outsize = self.outsize or self.expected_outsize%0A                assert y_data.shape%5B2:%5D == outsize%0A                if outsize == (5, 2):%0A                    expect = numpy.zeros(outsize, dtype=numpy.float32)%0A                    expect%5B:2, :%5D = self.x%5Bi, c, 0, 0%5D%0A                    expect%5B2:4, :%5D = self.x%5Bi, c, 1, 0%5D%0A                elif outsize == (4, 2):%0A    
         
@@ -2204,16 +2204,20 @@
 array(%5B%0A
+    
         
@@ -2282,32 +2282,36 @@
                 
+  
+  
 %5Bself.x%5Bi, c, 0,
@@ -2328,32 +2328,36 @@
 x%5Bi, c, 0, 0%5D%5D,%0A
+    
                 
@@ -2414,32 +2414,36 @@
                 
+    
 %5Bself.x%5Bi, c, 1,
@@ -2484,17 +2484,367 @@
         
+  
-%5D
+  %5D)%0A                elif outsize == (3, 1):%0A                    expect = numpy.array(%5B%0A                        %5Bself.x%5Bi, c, 0, 0%5D%5D,%0A                        %5Bself.x%5Bi, c, 0, 0%5D%5D,%0A                        %5Bself.x%5Bi, c, 1, 0%5D%5D,%0A                    %5D)%0A                else:%0A                    raise ValueError('Unsupported outsize: %7B%7D'.format(outsize)
 )%0A      

Test of estimating outsize from insize for Unpooling2D