from __future__ import absolute_import, print_function, division
import unittest

from nose.plugins.skip import SkipTest
import numpy as np

import theano
import theano.tensor as T
from theano.tests import unittest_tools as utt

from theano.tensor.signal import conv

from theano.tensor.basic import _allclose


class TestSignalConv2D(unittest.TestCase):

    def setUp(self):
        utt.seed_rng()

    def validate(self, image_shape, filter_shape, out_dim, verify_grad=True):

        image_dim = len(image_shape)
        filter_dim = len(filter_shape)
        input = T.TensorType('float64', [False] * image_dim)()
        filters = T.TensorType('float64', [False] * filter_dim)()

        bsize = image_shape[0]
        if image_dim != 3:
            bsize = 1
        nkern = filter_shape[0]
        if filter_dim != 3:
            nkern = 1

        # THEANO IMPLEMENTATION ############
        # we create a symbolic function so that verify_grad can work
        def sym_conv2d(input, filters):
            return conv.conv2d(input, filters)
        output = sym_conv2d(input, filters)
        assert output.ndim == out_dim
        theano_conv = theano.function([input, filters], output)

        # initialize input and compute result
        image_data = np.random.random(image_shape)
        filter_data = np.random.random(filter_shape)
        theano_output = theano_conv(image_data, filter_data)

        # REFERENCE IMPLEMENTATION ############
        out_shape2d = np.array(image_shape[-2:]) - np.array(filter_shape[-2:]) + 1
        ref_output = np.zeros(tuple(out_shape2d))

        # reshape as 3D input tensors to make life easier
        image_data3d = image_data.reshape((bsize,) + image_shape[-2:])
        filter_data3d = filter_data.reshape((nkern,) + filter_shape[-2:])
        # reshape theano output as 4D to make life easier
        theano_output4d = theano_output.reshape((bsize, nkern,) +
                                                theano_output.shape[-2:])

        # loop over mini-batches (if required)
        for b in range(bsize):

            # loop over filters (if required)
            for k in range(nkern):

                image2d = image_data3d[b, :, :]
                filter2d = filter_data3d[k, :, :]
                output2d = np.zeros(ref_output.shape)
                for row in range(ref_output.shape[0]):
                    for col in range(ref_output.shape[1]):
                        output2d[row, col] += (
                            image2d[row:row + filter2d.shape[0],
                                    col:col + filter2d.shape[1]] *
                            filter2d[::-1, ::-1]
                            ).sum()

                self.assertTrue(_allclose(theano_output4d[b, k, :, :],
                                          output2d))

        # TEST GRADIENT ############
        if verify_grad:
            utt.verify_grad(sym_conv2d, [image_data, filter_data])

    def test_basic(self):
        """
        Basic functionality of nnet.conv.ConvOp is already tested by
        its own test suite.  We just have to test whether or not
        signal.conv.conv2d can support inputs and filters of type
        matrix or tensor3.
        """
        if(not theano.tensor.nnet.conv.imported_scipy_signal and
           theano.config.cxx == ""):
            raise SkipTest("conv2d tests need SciPy or a c++ compiler")

        self.validate((1, 4, 5), (2, 2, 3), out_dim=4, verify_grad=True)
        self.validate((7, 5), (5, 2, 3), out_dim=3, verify_grad=False)
        self.validate((3, 7, 5), (2, 3), out_dim=3, verify_grad=False)
        self.validate((7, 5), (2, 3), out_dim=2, verify_grad=False)

    def test_fail(self):
        """
        Test that conv2d fails for dimensions other than 2 or 3.
        """
        self.assertRaises(Exception, conv.conv2d, T.dtensor4(), T.dtensor3())
        self.assertRaises(Exception, conv.conv2d, T.dtensor3(), T.dvector())

    def test_bug_josh_reported(self):
        """
        Test refers to a bug reported by Josh, when due to a bad merge these
        few lines of code failed. See
        http://groups.google.com/group/theano-dev/browse_thread/thread/8856e7ca5035eecb
        """
        m1 = theano.tensor.matrix()
        m2 = theano.tensor.matrix()
        conv.conv2d(m1, m2)
