import itertools
from unittest import TestCase

from theano.compile.pfunc import pfunc
from theano import tensor
from theano.tests import unittest_tools

import numpy

# Skip test if cuda_ndarray is not available.
from nose.plugins.skip import SkipTest
import theano.sandbox.cuda as cuda_ndarray
if cuda_ndarray.cuda_available == False:
    raise SkipTest('Optional package cuda disabled')

import theano.sandbox.cuda as tcn

from theano.tensor.signal.downsample import DownsampleFactorMax, DownsampleFactorMaxGrad

import theano.compile.mode
from theano.tensor.tests.test_blas import BaseGemv, TestBlasStrides, TestGer
from theano.sandbox.cuda.blas import gpu_gemv_no_inplace, gpu_gemv_inplace
from theano.sandbox.cuda.blas import gpu_ger_inplace, gpu_ger_no_inplace


if theano.config.mode=='FAST_COMPILE':
    mode_with_gpu = theano.compile.mode.get_mode('FAST_RUN').including('gpu')
    mode_without_gpu = theano.compile.mode.get_mode('FAST_RUN').excluding('gpu')
else:
    mode_with_gpu = theano.compile.mode.get_default_mode().including('gpu')
    mode_without_gpu = theano.compile.mode.get_default_mode().excluding('gpu')

def my_rand(*shape):
    return theano._asarray(numpy.random.rand(*shape),dtype='float32')

def transpose(cuda_mat):
    # The easiest way to transpose a cuda matrix for now
    return tcn.dimshuffle(cuda_mat, [1, 0])

def test_dot22():
    def cmp(a_shp, b_shp):
        a0 = my_rand(*a_shp)
        a = tcn.shared_constructor(a0, 'a')

        b = tensor.fmatrix()

        f = pfunc([b], [], updates=[(a, tensor.dot(a,b))], mode=mode_with_gpu)

        bval = my_rand(*b_shp)
        f(bval)

        assert numpy.allclose(numpy.dot(a0, bval), a.get_value())

        # Try with a matrix equal to a0, but with strides in both dims
        a.set_value(a0)
        a.set_value(
                a.get_value(borrow=True, return_internal_type=True)[::-1, ::-1],
                borrow=True)
        f(bval)

    cmp((3,4),(4,5))
    cmp((0,4),(4,5))
    cmp((3,4),(4,0))
    cmp((3,0),(0,5))
    cmp((0,4),(4,0))
    cmp((0,0),(0,0))

def test_dot22scalar():
    def cmp(a_shp, b_shp):
        a = tensor.fmatrix()
        b = tensor.fmatrix()
        scalar = tensor.fscalar()
        av = my_rand(*a_shp)
        bv = my_rand(*b_shp)

        f = theano.function([a,b], tensor.dot(a,b)*numpy.asarray(4, 'float32'), mode=mode_with_gpu)
        f2 = theano.function([a,b], tensor.dot(a,b)*numpy.asarray(4, 'float32'))
        t=f.maker.env.toposort()
        assert len(t)==4
        assert isinstance(t[0].op,tcn.GpuFromHost)
        assert isinstance(t[1].op,tcn.GpuFromHost)
        assert isinstance(t[2].op,tcn.blas.GpuDot22Scalar)
        assert isinstance(t[3].op,tcn.HostFromGpu)
        assert numpy.allclose(f(av,bv),f2(av,bv))

        f = theano.function([a,b,scalar], tensor.dot(a,b)*scalar, mode=mode_with_gpu)
        f2 = theano.function([a,b,scalar], tensor.dot(a,b)*scalar)
        t=f.maker.env.toposort()
        assert len(t)==4
        assert isinstance(t[0].op,tcn.GpuFromHost)
        assert isinstance(t[1].op,tcn.GpuFromHost)
        assert isinstance(t[2].op,tcn.blas.GpuDot22Scalar)
        assert isinstance(t[3].op,tcn.HostFromGpu)
        assert numpy.allclose(f(av,bv,0.5),f2(av,bv,0.5))

    cmp((3,4),(4,5))
    cmp((0,4),(4,5))
    cmp((3,4),(4,0))
    cmp((3,0),(0,5))
    cmp((0,4),(4,0))
    cmp((0,0),(0,0))

def test_gemm():
    def cmp(a_shp, b_shp):
        a0 = my_rand(*a_shp)
        a = tcn.shared_constructor(a0, 'a')

        b = tensor.fmatrix('b')
        c = tensor.fmatrix('c')

        f = pfunc([b,c], [], updates=[(a, tensor.dot(a,b) + tensor.exp(c))], mode=mode_with_gpu)
        assert any([node.op == tcn.blas.gpu_gemm_inplace for node in f.maker.env.toposort()])

        bval = my_rand(*b_shp)
        cval = my_rand(a_shp[0],b_shp[1])
        f(bval,cval)

        assert numpy.allclose(numpy.dot(a0, bval)+numpy.exp(cval), a.get_value())

        # Try with a matrix equal to a0, but with strides in both dims
        a.set_value(a0)
        a.set_value(
                a.get_value(borrow=True, return_internal_type=True)[::-1, ::-1],
                borrow=True)
        f(bval, cval)

    cmp((3,4),(4,5))
    cmp((0,4),(4,5))
    cmp((3,4),(4,0))
    cmp((3,0),(0,5))
    cmp((0,4),(4,0))
    cmp((0,0),(0,0))

def test_gemm_no_inplace():

    def cmp(a_shp, b_shp):
        a0 = my_rand(*a_shp)
        a = tcn.shared_constructor(a0, 'a')
        cval = my_rand(a_shp[0], b_shp[1])
        c = tcn.shared_constructor(cval.copy(), 'c')

        b = tcn.fmatrix('b')
        b2 = tcn.fmatrix('b2')

        f = pfunc([b,b2], [tensor.dot(a,b2) + c], updates=[(a, tensor.dot(a,b) + c)], mode=mode_with_gpu)

        assert any([node.op == tcn.blas.gpu_gemm_no_inplace for node in f.maker.env.toposort()])
        bval = my_rand(*b_shp)
        bval2 = my_rand(*b_shp)
        rval = f(bval,bval2)

        assert numpy.allclose(numpy.dot(a0, bval)+cval, a.get_value())
        assert numpy.allclose(numpy.dot(a0, bval2)+cval, rval)

        # Try with a matrix equal to a0, but with strides in both dims
        a.set_value(a0)
        a.set_value(
                a.get_value(borrow=True, return_internal_type=True)[::-1, ::-1],
                borrow=True)
        f(bval, bval2)

    cmp((3,4),(4,5))
    cmp((0,4),(4,5))
    cmp((3,4),(4,0))
    cmp((3,0),(0,5))
    cmp((0,4),(4,0))
    cmp((0,0),(0,0))


class TestBlasStridesGpu(TestBlasStrides):
    dtype = 'float32'
    shared = staticmethod(tcn.shared_constructor)
    mode = mode_with_gpu

def test_dot22_strides():
    def cmp(b_shp, c_shp):
        av = numpy.zeros((0, 0), dtype='float32')
        bv = my_rand(*b_shp)
        cv = my_rand(*c_shp)

        a = tcn.shared_constructor(av, 'a')
        b = tcn.shared_constructor(bv, 'b')
        c = tcn.shared_constructor(cv, 'c')

        b_t = tcn.shared_constructor(bv.T, 'b.T')
        c_t = tcn.shared_constructor(cv.T, 'c.T')

        b_gpu = b.get_value(borrow=False, return_internal_type=True)
        c_gpu = c.get_value(borrow=False, return_internal_type=True)
        bt_gpu = b_t.get_value(borrow=False, return_internal_type=True)
        ct_gpu = c_t.get_value(borrow=False, return_internal_type=True)

        f_nn = pfunc([], [], updates={a: tensor.dot(b, c)},
                mode=mode_with_gpu)
        f_nt = pfunc([], [], updates={a: tensor.dot(b, c_t.T)},
                mode=mode_with_gpu)
        f_tn = pfunc([], [], updates={a: tensor.dot(b_t.T, c)},
                mode=mode_with_gpu)
        f_tt = pfunc([], [], updates={a: tensor.dot(b_t.T, c_t.T)},
                mode=mode_with_gpu)

        # Try with all stride patterns, and all transposed pattern
        for step_signs in itertools.product((-1, 1), repeat=4):
            for step in (1, 2):
                b_step1, b_step2, c_step1, c_step2 = (s * step
                        for s in step_signs)

                b.set_value(b_gpu.copy()[::b_step1, ::b_step2], borrow=True)
                c.set_value(c_gpu.copy()[::c_step1, ::c_step2], borrow=True)
                b_t.set_value(bt_gpu.copy()[::b_step2, ::b_step1], borrow=True)
                c_t.set_value(ct_gpu.copy()[::c_step2, ::c_step1], borrow=True)

                # Numpy result
                a_n = numpy.dot(bv[::b_step1, ::b_step2],
                                cv[::c_step1, ::c_step2])

                f_nn()
                assert numpy.allclose(a.get_value(), a_n)

                f_nt()
                assert numpy.allclose(a.get_value(), a_n)

                f_tn()
                assert numpy.allclose(a.get_value(), a_n)

                f_tt()
                assert numpy.allclose(a.get_value(), a_n)

    cmp((3, 4), (4, 5))
    cmp((1, 4), (4, 5))
    cmp((3, 4), (4, 1))
    cmp((3, 1), (1, 1))
    cmp((1, 4), (4, 1))
    cmp((3, 1), (1, 5))
    cmp((0, 4), (4, 5))
    cmp((0, 4), (4, 1))
    cmp((0, 1), (1, 5))
    cmp((3, 4), (4, 0))
    cmp((3, 0), (0, 5))
    cmp((0, 4), (4, 0))
    cmp((0, 0), (0, 0))


def test_dot22scalar_strides():
    def cmp(b_shp, c_shp):
        av = numpy.zeros((0, 0), dtype='float32')
        bv = my_rand(*b_shp)
        cv = my_rand(*c_shp)
        l = numpy.float32(0.2)

        a = tcn.shared_constructor(av, 'a')
        b = tcn.shared_constructor(bv, 'b')
        c = tcn.shared_constructor(cv, 'c')

        b_t = tcn.shared_constructor(bv.T, 'b.T')
        c_t = tcn.shared_constructor(cv.T, 'c.T')

        b_gpu = b.get_value(borrow=False, return_internal_type=True)
        c_gpu = c.get_value(borrow=False, return_internal_type=True)
        bt_gpu = b_t.get_value(borrow=False, return_internal_type=True)
        ct_gpu = c_t.get_value(borrow=False, return_internal_type=True)

        f_nn = pfunc([], [], updates={a: l * tensor.dot(b, c)},
                mode=mode_with_gpu)
        f_nt = pfunc([], [], updates={a: l * tensor.dot(b, c_t.T)},
                mode=mode_with_gpu)
        f_tn = pfunc([], [], updates={a: l * tensor.dot(b_t.T, c)},
                mode=mode_with_gpu)
        f_tt = pfunc([], [], updates={a: l * tensor.dot(b_t.T, c_t.T)},
                mode=mode_with_gpu)

        # Try with all stride patterns, and all transposed pattern
        for step_signs in itertools.product((-1, 1), repeat=4):
            for step in (1, 2):
                b_step1, b_step2, c_step1, c_step2 = (s * step
                        for s in step_signs)

                b.set_value(b_gpu.copy()[::b_step1, ::b_step2], borrow=True)
                c.set_value(c_gpu.copy()[::c_step1, ::c_step2], borrow=True)
                b_t.set_value(bt_gpu.copy()[::b_step2, ::b_step1], borrow=True)
                c_t.set_value(ct_gpu.copy()[::c_step2, ::c_step1], borrow=True)

                # Numpy result
                a_n = l * numpy.dot(bv[::b_step1, ::b_step2],
                                    cv[::c_step1, ::c_step2])

                f_nn()
                assert numpy.allclose(a.get_value(), a_n)

                f_nt()
                assert numpy.allclose(a.get_value(), a_n)

                f_tn()
                assert numpy.allclose(a.get_value(), a_n)

                f_tt()
                assert numpy.allclose(a.get_value(), a_n)

    cmp((3, 4), (4, 5))
    cmp((1, 4), (4, 5))
    cmp((3, 4), (4, 1))
    cmp((3, 1), (1, 1))
    cmp((1, 4), (4, 1))
    cmp((3, 1), (1, 5))
    cmp((0, 4), (4, 5))
    cmp((0, 4), (4, 1))
    cmp((0, 1), (1, 5))
    cmp((3, 4), (4, 0))
    cmp((3, 0), (0, 5))
    cmp((0, 4), (4, 0))
    cmp((0, 0), (0, 0))


def test_gemm_strides():
    def cmp(a_shp, b_shp, c_shp):
        av = my_rand(*a_shp)
        bv = my_rand(*b_shp)
        cv = my_rand(*c_shp)
        l = numpy.float32(0.2)

        a = tcn.shared_constructor(av, 'a')
        b = tcn.shared_constructor(bv, 'b')
        c = tcn.shared_constructor(cv, 'c')

        a_t = tcn.shared_constructor(av.T, 'a.T')
        b_t = tcn.shared_constructor(bv.T, 'b.T')
        c_t = tcn.shared_constructor(cv.T, 'c.T')

        a_gpu = a.get_value(borrow=False, return_internal_type=True)
        b_gpu = b.get_value(borrow=False, return_internal_type=True)
        c_gpu = c.get_value(borrow=False, return_internal_type=True)
        bt_gpu = b_t.get_value(borrow=False, return_internal_type=True)
        ct_gpu = c_t.get_value(borrow=False, return_internal_type=True)

        f_nnn = pfunc([], [], updates={a: (l * a + tensor.dot(b, c))},
                mode=mode_with_gpu)
        f_nnt = pfunc([], [], updates={a: (l * a + tensor.dot(b, c_t.T))},
                mode=mode_with_gpu)
        f_ntn = pfunc([], [], updates={a: (l * a + tensor.dot(b_t.T, c))},
                mode=mode_with_gpu)
        f_ntt = pfunc([], [], updates={a: (l * a + tensor.dot(b_t.T, c_t.T))},
                mode=mode_with_gpu)
        f_tnn = pfunc([], [], updates={a_t: (l * a_t + tensor.dot(b, c).T)},
                mode=mode_with_gpu)
        f_tnt = pfunc([], [], updates={a_t: (l * a_t + tensor.dot(b, c_t.T).T)},
                mode=mode_with_gpu)
        f_ttn = pfunc([], [], updates={a_t: (l * a_t + tensor.dot(b_t.T, c).T)},
                mode=mode_with_gpu)
        f_ttt = pfunc([], [], updates={a_t: (l * a_t + tensor.dot(b_t.T, c_t.T).T)},
                mode=mode_with_gpu)

        # Try with all stride patterns, and all transposed pattern
        for step_signs in itertools.product((-1, 1), repeat=6):
            for step in (1, 2):
                a_step1, a_step2, b_step1, b_step2, c_step1, c_step2 = \
                        (s * step for s in step_signs)

                b.set_value(b_gpu.copy()[::b_step1, ::b_step2], borrow=True)
                c.set_value(c_gpu.copy()[::c_step1, ::c_step2], borrow=True)
                b_t.set_value(bt_gpu.copy()[::b_step2, ::b_step1], borrow=True)
                c_t.set_value(ct_gpu.copy()[::c_step2, ::c_step1], borrow=True)

                # Numpy results
                a_n = (l * av[::a_step1, ::a_step2]
                       + numpy.dot(bv[::b_step1, ::b_step2],
                                   cv[::c_step1, ::c_step2]))
                at_n = (l * av[::a_step1, ::a_step2].T
                        + numpy.dot(bv[::b_step1, ::b_step2],
                                    cv[::c_step1, ::c_step2]).T)

                # a's value is updated, so we need to reinitialize it each time
                a.set_value(a_gpu.copy()[::a_step1, ::a_step2], borrow=True)
                f_nnn()
                assert numpy.allclose(a.get_value(), a_n)

                a.set_value(a_gpu.copy()[::a_step1, ::a_step2], borrow=True)
                f_nnt()
                assert numpy.allclose(a.get_value(), a_n)

                a.set_value(a_gpu.copy()[::a_step1, ::a_step2], borrow=True)
                f_ntn()
                assert numpy.allclose(a.get_value(), a_n)

                a.set_value(a_gpu.copy()[::a_step1, ::a_step2], borrow=True)
                f_ntt()
                assert numpy.allclose(a.get_value(), a_n)

                a_t.set_value(transpose(a_gpu.copy())[::a_step2, ::a_step1],
                        borrow=True)
                f_tnn()
                assert numpy.allclose(a_t.get_value(), at_n)

                a_t.set_value(transpose(a_gpu.copy())[::a_step2, ::a_step1],
                        borrow=True)
                f_tnt()
                assert numpy.allclose(a_t.get_value(), at_n)

                a_t.set_value(transpose(a_gpu.copy())[::a_step2, ::a_step1],
                        borrow=True)
                f_ttn()
                assert numpy.allclose(a_t.get_value(), at_n)

                a_t.set_value(transpose(a_gpu.copy())[::a_step2, ::a_step1],
                        borrow=True)
                f_ttt()
                assert numpy.allclose(a_t.get_value(), at_n)


    cmp((3, 5), (3, 4), (4, 5))
    cmp((1, 5), (1, 4), (4, 5))
    cmp((3, 1), (3, 4), (4, 1))
    cmp((3, 1), (3, 1), (1, 1))
    cmp((1, 1), (1, 4), (4, 1))
    cmp((3, 5), (3, 1), (1, 5))
    cmp((0, 5), (0, 4), (4, 5))
    cmp((0, 1), (0, 4), (4, 1))
    cmp((0, 5), (0, 1), (1, 5))
    cmp((3, 0), (3, 4), (4, 0))
    cmp((3, 5), (3, 0), (0, 5))
    cmp((0, 0), (0, 4), (4, 0))
    cmp((0, 0), (0, 0), (0, 0))


def test_gemv_strides():
    def cmp(a_shp, b_shp, c_shp):
        av = my_rand(a_shp)
        bv = my_rand(*b_shp)
        cv = my_rand(c_shp)
        l = numpy.float32(0.2)

        a = tcn.shared_constructor(av, 'a')
        b = tcn.shared_constructor(bv, 'b')
        c = tcn.shared_constructor(cv, 'c')
        b_t = tcn.shared_constructor(bv.T, 'b.T')

        a_gpu = a.get_value(borrow=False, return_internal_type=True)
        b_gpu = b.get_value(borrow=False, return_internal_type=True)
        c_gpu = c.get_value(borrow=False, return_internal_type=True)

        f_n = pfunc([], [], updates={a: (a + l * tensor.dot(b, c))},
                mode=mode_with_gpu)

        f_t = pfunc([], [], updates={a: (a + l * tensor.dot(b_t.T, c))},
                mode=mode_with_gpu)

        # Try with all stride patterns, and all transposed pattern
        for step_signs in itertools.product((1, -1), repeat=4):
            for step in (1, 2):
                a_step, b_step1, b_step2, c_step = (s * step
                        for s in step_signs)

                a.set_value(a_gpu.copy()[::a_step], borrow=True)
                b.set_value(b_gpu.copy()[::b_step1, ::b_step2],
                        borrow=True)
                b_t.set_value(transpose(b_gpu.copy())[::b_step2, ::b_step1],
                        borrow=True)
                c.set_value(c_gpu.copy()[::c_step], borrow=True)

                a_n = (av[::a_step]
                        + l * numpy.dot(bv[::b_step1, ::b_step2], cv[::c_step]))
                f_n()
                assert numpy.allclose(a.get_value(), a_n), (a.get_value(), a_n)

                a.set_value(a_gpu.copy()[::a_step], borrow=True)
                f_t()
                assert numpy.allclose(a.get_value(), a_n), (a.get_value(), a_n)

    cmp(3, (3, 5), 5)
    cmp(1, (1, 5), 5)
    cmp(3, (3, 1), 1)
    cmp(0, (0, 5), 5)
    cmp(3, (3, 0), 0)
    cmp(0, (0, 1), 1)
    cmp(1, (1, 0), 0)
    cmp(0, (0, 0), 0)


def test_ger_strides():
    def cmp(a_shp, b_shp, c_shp):
        av = my_rand(*a_shp)
        bv = my_rand(b_shp)
        cv = my_rand(c_shp)
        l = numpy.float32(0.2)

        a = tcn.shared_constructor(av, 'a')
        b = tcn.shared_constructor(bv, 'b')
        c = tcn.shared_constructor(cv, 'c')
        a_t = tcn.shared_constructor(av.T, 'a.T')

        a_gpu = a.get_value(borrow=False, return_internal_type=True)
        b_gpu = b.get_value(borrow=False, return_internal_type=True)
        c_gpu = c.get_value(borrow=False, return_internal_type=True)

        f_n = pfunc([], [], updates={a: (a + l * tensor.outer(b, c))},
                mode=mode_with_gpu)

        f_t = pfunc([], [], updates={a_t: (a_t + l * tensor.outer(b, c).T)},
                mode=mode_with_gpu)

        # Try with all stride patterns, and all transposed patterns
        for step_signs in itertools.product((1, -1), repeat=4):
            for step in (1, 2):
                a_step1, a_step2, b_step, c_step = (s * step
                        for s in step_signs)

                a.set_value(a_gpu.copy()[::a_step1, ::a_step2], borrow=True)
                a_t.set_value(transpose(a_gpu.copy())[::a_step1, ::a_step2],
                        borrow=True)
                b.set_value(b_gpu.copy()[::b_step], borrow=True)
                c.set_value(c_gpu.copy()[::c_step], borrow=True)

                f_n()
                n_n = (av[::a_step1, ::a_step2]
                        + l * numpy.outer(bv[::b_step], cv[::c_step]))
                assert numpy.allclose(a.get_value(), n_n), (a.get_value(), n_n)

                f_t()
                n_t = (av.T[::a_step1, ::a_step2]
                        + l * numpy.outer(bv[::b_step], cv[::c_step]).T)
                assert numpy.allclose(a_t.get_value(), n_t),\
                        (a_t.get_value(), n_t)

    cmp((3, 5), 3, 5)
    cmp((1, 5), 1, 5)
    cmp((3, 1), 3, 1)
    cmp((0, 5), 0, 5)
    cmp((3, 0), 3, 0)
    cmp((0, 1), 0, 1)
    cmp((1, 0), 1, 0)
    cmp((0, 0), 0, 0)


def test_outer():
    x = tcn.shared_constructor(my_rand(8,), 'x')
    y = tcn.shared_constructor(my_rand(6,), 'y')

    x_val = x.get_value().copy()
    y_val = y.get_value().copy()

    f = pfunc([], tensor.outer(x, y), mode=mode_with_gpu)
    assert numpy.allclose(numpy.outer(x_val, y_val), f())

    f = pfunc([], tensor.outer(x[::2], y), mode=mode_with_gpu)
    assert numpy.allclose(numpy.outer(x_val[::2], y_val), f())

    f = pfunc([], tensor.outer(x, y[::3]), mode=mode_with_gpu)
    assert numpy.allclose(numpy.outer(x_val, y_val[::3]), f())

    f = pfunc([], tensor.outer(x[::2], y[::3]), mode=mode_with_gpu)
    assert numpy.allclose(numpy.outer(x_val[::2], y_val[::3]), f())

    f = pfunc([], tensor.outer(x[::-1], y), mode=mode_with_gpu)
    assert numpy.allclose(numpy.outer(x_val[::-1], y_val), f())

    f = pfunc([], tensor.outer(x, y[::-1]), mode=mode_with_gpu)
    assert numpy.allclose(numpy.outer(x_val, y_val[::-1]), f())

if 0:
    # This is commented out because it doesn't make sense...
    # tcn.blas has no op called DownsampleFactorMax
    # tcn.blas has an op called GpuDownsampleFactorMax, but that op requires arguments that are
    # CudaNdarrayType variables... so rethink this test?
    def test_maxpool():
        """TODO: test the gpu version!!! """
        for d0, d1, r_true, r_false in [(4,4,[[[[5,7],[13,15]]]],[[[[5,7],[13,15]]]]),
                                        (5,5,[[[[6, 8],[ 16, 18], [ 21, 23]]]],
                                         [[[[6, 8, 9],[ 16, 18, 19], [ 21, 23, 24]]]])]:
            for border,ret in [(True,r_true),(False, r_false)]:
                ret=numpy.array(ret)
                a = tcn.blas.DownsampleFactorMax((2,2),border)
                dmatrix4 = tensor.TensorType("float32", (False, False, False, False))
                b = dmatrix4()
                f = pfunc([b], [a(b)], mode=mode_with_gpu)

                bval = numpy.arange(0,d0*d1).reshape(1,1,d0,d1)
                r = f(bval)[0]
    #            print bval, bval.shape, border
                print r, r.shape
                assert (ret==r).all()

def test_downsample():
    import random
    shps = [ (1, 1, 1, 12),
            (1, 1, 2, 2),
            (1, 1, 1, 1),
            (1,1,4,4),
            (1, 1, 10, 11),
            (1, 2, 2, 2),
            (3,5,4,4),
            (25, 1, 7, 7),
            (1, 1, 12, 12),
            (1, 1, 2, 14),
            (1, 1, 12, 14),
            (1, 1, 14, 14),
            (1, 1, 16, 16),
            (1, 1, 18, 18),
            (1, 1, 24, 24),
            (1, 6, 24, 24),
            (10, 1, 24, 24),
            (10, 6, 24, 24),
            (30, 6, 12, 12),
            (30, 2, 24, 24),
            (30, 6, 24, 24),
            (10, 10, 10, 11),
            (1,1,10,1025),
            (1,1,10,1023),
            (1,1,1025,10),
            (1,1,1023,10),
             ]

    numpy.random.RandomState(unittest_tools.fetch_seed()).shuffle(shps)

    for shp in shps:
        for ds in (2, 2), (3,2), (1,1):
            if ds[0] > shp[2]: continue
            if ds[1] > shp[3]: continue
            #GpuDownsampleFactorMax don't having more then 512 columns in the output tensor
            if float(shp[3])/ds[1]>512: continue
            for ignore_border in (True, False):
                print 'test_downsample', shp, ds, ignore_border
                ds_op = DownsampleFactorMax(ds, ignore_border=ignore_border)

                a = tcn.shared_constructor(my_rand(*shp), 'a')
                f = pfunc([], ds_op(tensor.as_tensor_variable(a)), mode=mode_with_gpu)
                f2 = pfunc([], ds_op(tensor.as_tensor_variable(a)), mode=mode_without_gpu)
                assert any([isinstance(node.op, tcn.blas.GpuDownsampleFactorMax) for node in
                            f.maker.env.toposort()])
                assert any([isinstance(node.op, DownsampleFactorMax) for node in
                            f2.maker.env.toposort()])
                assert numpy.allclose(f(),f2())

                g = pfunc([], tensor.grad(ds_op(tensor.as_tensor_variable(a)).sum(),a), mode=mode_with_gpu)
                g2 = pfunc([], tensor.grad(ds_op(tensor.as_tensor_variable(a)).sum(),a), mode=mode_without_gpu)
                assert any([isinstance(node.op, tcn.blas.GpuDownsampleFactorMaxGrad)
                            for node in g.maker.env.toposort()])
                assert any([isinstance(node.op, DownsampleFactorMaxGrad)
                            for node in g2.maker.env.toposort()])
                assert numpy.allclose(g(),g2())

                #We already check that the gpu version return the same value as the gpu version
                #for GpuDownsampleFactorMaxGrad. So no need to call verify_grad here.


class TestGpuGemv(TestCase, BaseGemv,
                  unittest_tools.TestOptimizationMixin):
    mode = mode_with_gpu
    dtype = 'float32'

    # As all input are transfered to the gpu, this allow to make all
    # the gemv inplace.
    gemv = gpu_gemv_inplace
    gemv_inplace = gpu_gemv_inplace

class TestGpuGemvNoTransfer(TestCase, BaseGemv,
                  unittest_tools.TestOptimizationMixin):
    mode = mode_with_gpu
    dtype = 'float32'

    # Mimic shared constructors registry
    @staticmethod
    def shared(val):
        try:
            return tcn.shared_constructor(val)
        except TypeError:
            return theano.shared(val)

    # In this test, inputs are not always transfered to GPU
    gemv = gpu_gemv_no_inplace
    gemv_inplace = gpu_gemv_inplace


class TestVectorMatrixDot(TestCase):
    ### Tolerance factor used in this tests
    atol = 1e-6
    ##########################

    def test_dot_vm(self):
        ''' Test vector dot matrix '''
        v = theano.shared( numpy.array(numpy.random.rand(2), dtype='float32'))
        m = theano.shared( numpy.array(numpy.random.rand(2,5),
                                       dtype='float32'))
        no_gpu_f = theano.function([], theano.dot(v,m), mode = mode_without_gpu)
        gpu_f    = theano.function([], theano.dot(v,m), mode = mode_with_gpu)
        #gpu_f2 is needed to test the case when the input is not on the gpu
        #but the output is moved to the gpu.
        gpu_f2   = theano.function([], tcn.gpu_from_host(theano.dot(v,m)), mode = mode_with_gpu)

        # Assert they produce the same output
        assert numpy.allclose(no_gpu_f(), gpu_f(), atol=self.atol)
        assert numpy.allclose(no_gpu_f(), gpu_f2(), atol=self.atol)
        # Assert that the gpu version actually uses gpu
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f.maker.env.toposort() ]) == 1
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f2.maker.env.toposort() ]) == 1

        # Check double-strided m
        m.set_value(
                m.get_value(borrow=True, return_internal_type=True)[::-1, ::-1],
                borrow=True)
        assert numpy.allclose(no_gpu_f(), gpu_f(), atol=self.atol)
        assert numpy.allclose(no_gpu_f(), gpu_f2(), atol=self.atol)


    def test_dot_mv(self):
        ''' Test matrix dot vector '''
        v = theano.shared( numpy.array(numpy.random.rand(2), dtype='float32'))
        m = theano.shared( numpy.array(numpy.random.rand(5,2),
                                       dtype='float32'))
        no_gpu_f = theano.function([], theano.dot(m,v), mode = mode_without_gpu)
        gpu_f    = theano.function([], theano.dot(m,v), mode = mode_with_gpu)
        #gpu_f2 is needed to test the case when the input is not on the gpu
        #but the output is moved to the gpu.
        gpu_f2   = theano.function([], tcn.gpu_from_host(theano.dot(m,v)), mode = mode_with_gpu)

        # Assert they produce the same output
        assert numpy.allclose(no_gpu_f(), gpu_f(), atol=self.atol)
        assert numpy.allclose(no_gpu_f(), gpu_f2(), atol=self.atol)
        # Assert that the gpu version actually uses gpu
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f.maker.env.toposort() ]) == 1
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f2.maker.env.toposort() ]) == 1

    def test_gemv1(self):
        ''' test vector1+dot(matrix,vector2) '''
        v1 = theano.tensor._shared( numpy.array(numpy.random.rand(2)  , dtype='float32'))
        v2 = theano.tensor._shared( numpy.array(numpy.random.rand(5)  , dtype='float32'))
        m  = theano.tensor._shared( numpy.array(numpy.random.rand(5,2), dtype='float32'))

        no_gpu_f = theano.function([], v2+theano.dot(m,v1), mode = mode_without_gpu)
        gpu_f    = theano.function([], v2+theano.dot(m,v1), mode = mode_with_gpu)
        #gpu_f2 is needed to test the case when the input is not on the gpu
        #but the output is moved to the gpu.
        gpu_f2    = theano.function([], tcn.gpu_from_host(v2+theano.dot(m,v1)), mode = mode_with_gpu)

        # Assert they produce the same output
        assert numpy.allclose(no_gpu_f(), gpu_f(), atol=self.atol)
        assert numpy.allclose(no_gpu_f(), gpu_f2(), atol=self.atol)
        # Assert that the gpu version actually uses gpu
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f2.maker.env.toposort()]) == 1
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f.maker.env.toposort()]) == 1

    def test_gemv2(self):
        ''' test vector1+dot(vector2,matrix) '''
        v1 = theano.shared( numpy.array(numpy.random.rand(5)  , dtype='float32'))
        v2 = theano.shared( numpy.array(numpy.random.rand(2)  , dtype='float32'))
        m  = theano.shared( numpy.array(numpy.random.rand(5,2), dtype='float32'))

        no_gpu_f = theano.function([], v2+theano.dot(v1,m), mode = mode_without_gpu)
        gpu_f    = theano.function([], v2+theano.dot(v1,m), mode = mode_with_gpu)
        #gpu_f2 is needed to test the case when the input is not on the gpu
        #but the output is moved to the gpu.
        gpu_f2    = theano.function([], tcn.gpu_from_host(v2+theano.dot(v1,m)), mode = mode_with_gpu)

        # Assert they produce the same output
        assert numpy.allclose(no_gpu_f(), gpu_f(), atol=self.atol)
        assert numpy.allclose(no_gpu_f(), gpu_f2(), atol=self.atol)
        # Assert that the gpu version actually uses gpu
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f2.maker.env.toposort()]) == 1
        assert sum([node.op is gpu_gemv_inplace for node in
                    gpu_f.maker.env.toposort()]) == 1


class TestGpuGer(TestGer):
    def setUp(self):
        self.mode = mode_with_gpu
        dtype = self.dtype = 'float32'  # optimization isn't dtype-dependent
        self.A = tensor.tensor(dtype=dtype, broadcastable=(False, False))
        self.a = tensor.tensor(dtype=dtype, broadcastable=())
        self.x = tensor.tensor(dtype=dtype, broadcastable=(False,))
        self.y = tensor.tensor(dtype=dtype, broadcastable=(False,))
        self.ger = gpu_ger_no_inplace
        self.ger_destructive = gpu_ger_inplace
        self.gemm = tcn.blas.gpu_gemm_no_inplace

        # data on the gpu make the op always inplace
        self.ger = gpu_ger_inplace
        self.gemm = tcn.blas.gpu_gemm_inplace

class TestGpuGerNoTransfer(TestGer):
    @staticmethod
    def shared(val):
        try:
            return tcn.shared_constructor(val)
        except TypeError:
            return theano.shared(val)

    def setUp(self):
        self.mode = mode_with_gpu
        dtype = self.dtype = 'float32'  # optimization isn't dtype-dependent
        self.A = tensor.tensor(dtype=dtype, broadcastable=(False, False))
        self.a = tensor.tensor(dtype=dtype, broadcastable=())
        self.x = tensor.tensor(dtype=dtype, broadcastable=(False,))
        self.y = tensor.tensor(dtype=dtype, broadcastable=(False,))
        # data on the gpu make the op always inplace
        self.ger = gpu_ger_inplace
        self.ger_destructive = gpu_ger_inplace
        self.gemm = tcn.blas.gpu_gemm_inplace


class TestGpuGer_OpContract(TestCase, unittest_tools.T_OpContractMixin):
    def setUp(self):
        self.ops = [gpu_ger_no_inplace, gpu_ger_inplace]

    def clone(self, op):
        return tcn.blas.GpuGer(op.inplace)
