"""
test the tensor and sparse type. The CudaNdarray type is tested in sandbox/cuda/tests/test_tensor_op.py.test_may_share_memory_cuda
"""

import numpy

import theano

try:
    import scipy.sparse
    scipy_imported = True
except ImportError:
    scipy_imported = False

from theano.misc.may_share_memory import may_share_memory


def test_may_share_memory():
    a = numpy.random.rand(5, 4)
    b = numpy.random.rand(5, 4)
    va = a.view()
    vb = b.view()
    ra = a.reshape((4, 5))
    rb = b.reshape((4, 5))
    ta = a.T
    tb = b.T

    for a_, b_, rep in [(a, a, True), (b, b, True), (a, b, False),
                      (a, a[0], True), (a, a[:, 0], True), (a, a.T, True),
                      (a, (0,), False), (a, 1, False), (a, None, False),
                      (a, va, True), (b, vb, True), (va, b, False), (a, vb, False),
                      (a, ra, True), (b, rb, True), (ra, b, False), (a, rb, False),
                      (a, ta, True), (b, tb, True), (ta, b, False), (a, tb, False),
                      ]:

        assert may_share_memory(a_, b_, False) == rep
        assert may_share_memory(b_, a_, False) == rep

    # test that it raise error when needed.
    for a_, b_, rep in [(a, (0,), False), (a, 1, False), (a, None, False), ]:
        assert may_share_memory(a_, b_, False) == rep
        assert may_share_memory(b_, a_, False) == rep
        try:
            may_share_memory(a_, b_)
            raise Exception("An error was expected")
        except TypeError:
            pass
        try:
            may_share_memory(b_, a_)
            raise Exception("An error was expected")
        except TypeError:
            pass

if scipy_imported:
    def test_may_share_memory_scipy():
        a = scipy.sparse.csc_matrix(scipy.sparse.eye(5, 3))
        b = scipy.sparse.csc_matrix(scipy.sparse.eye(4, 3))
        as_ar = lambda a: theano._asarray(a, dtype='int32')
        for a_, b_, rep in [(a, a, True), (b, b, True), (a, b, False),
                          (a, a.data, True), (a, a.indptr, True), (a, a.indices, True), (a, as_ar(a.shape), False),
                          (a.data, a, True), (a.indptr, a, True), (a.indices, a, True), (as_ar(a.shape), a, False),
                          (b, b.data, True), (b, b.indptr, True), (b, b.indices, True), (b, as_ar(b.shape), False),
                          (b.data, b, True), (b.indptr, b, True), (b.indices, b, True), (as_ar(b.shape), b, False),
                          (b.data, a, False), (b.indptr, a, False), (b.indices, a, False), (as_ar(b.shape), a, False),
                          ]:

            assert may_share_memory(a_, b_) == rep
            assert may_share_memory(b_, a_) == rep

        # test that it raise error when needed.
        for a_, b_, rep in [(a, (0,), False), (a, 1, False), (a, None, False)]:
            assert may_share_memory(a_, b_, False) == rep
            assert may_share_memory(b_, a_, False) == rep
            try:
                may_share_memory(a_, b_)
                raise Exception("An error was expected")
            except TypeError:
                pass
            try:
                may_share_memory(b_, a_)
                raise Exception("An error was expected")
            except TypeError:
                pass
