import numpy as np

from numba.roc.vectorizers import HsaGUFuncVectorize
from numba.roc.dispatch import HSAGenerializedUFunc
from numba import guvectorize
import unittest


def ufunc_add_core(a, b, c):
    for i in range(c.size):
        c[i] = a[i] + b[i]


class TestGUFuncBuilding(unittest.TestCase):
    def test_gufunc_building(self):
        ufbldr = HsaGUFuncVectorize(ufunc_add_core, "(x),(x)->(x)")
        ufbldr.add("(float32[:], float32[:], float32[:])")
        ufbldr.add("(intp[:], intp[:], intp[:])")
        ufunc = ufbldr.build_ufunc()
        self.assertIsInstance(ufunc, HSAGenerializedUFunc)

        # Test integer version
        A = np.arange(100, dtype=np.intp)
        B = np.arange(100, dtype=np.intp) + 1
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_equal(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.intp), got.dtype)

        # Test integer version with 2D inputs
        A = A.reshape(50, 2)
        B = B.reshape(50, 2)
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_equal(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.intp), got.dtype)

        # Test integer version with 3D inputs
        A = A.reshape(5, 10, 2)
        B = B.reshape(5, 10, 2)
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_equal(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.intp), got.dtype)

        # Test real version
        A = np.arange(100, dtype=np.float32)
        B = np.arange(100, dtype=np.float32) + 1
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_allclose(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.float32), got.dtype)

        # Test real version with 2D inputs
        A = A.reshape(50, 2)
        B = B.reshape(50, 2)
        expected = A + B
        got = ufunc(A, B)

        np.testing.assert_allclose(expected, got)
        self.assertEqual(expected.dtype, got.dtype)
        self.assertEqual(np.dtype(np.float32), got.dtype)

    def test_gufunc_building_scalar_output(self):
        def sum_row(inp, out):
            tmp = 0.
            for i in range(inp.shape[0]):
                tmp += inp[i]
            out[0] = tmp

        ufbldr = HsaGUFuncVectorize(sum_row, "(n)->()")
        ufbldr.add("void(int32[:], int32[:])")
        ufunc = ufbldr.build_ufunc()

        inp = np.arange(300, dtype=np.int32).reshape(100, 3)
        out = ufunc(inp)

        for i in range(inp.shape[0]):
            np.testing.assert_equal(inp[i].sum(), out[i])

    def test_gufunc_scalar_input_saxpy(self):
        def axpy(a, x, y, out):
            for i in range(out.shape[0]):
                out[i] = a * x[i] + y[i]

        ufbldr = HsaGUFuncVectorize(axpy, '(),(t),(t)->(t)')
        ufbldr.add("void(float32, float32[:], float32[:], float32[:])")
        saxpy = ufbldr.build_ufunc()

        A = np.float32(2)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i])

        X = np.arange(10, dtype=np.float32)
        Y = np.arange(10, dtype=np.float32)
        out = saxpy(A, X, Y)

        for j in range(10):
            exp = A * X[j] + Y[j]
            self.assertTrue(exp == out[j], (exp, out[j]))

        A = np.arange(5, dtype=np.float32)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A[j] * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i], (exp, out[j, i]))


class TestGUFuncDecor(unittest.TestCase):
    def test_gufunc_decorator(self):
        @guvectorize(["void(float32, float32[:], float32[:], float32[:])"],
                     '(),(t),(t)->(t)', target='roc')
        def saxpy(a, x, y, out):
            for i in range(out.shape[0]):
                out[i] = a * x[i] + y[i]

        A = np.float32(2)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i])

        X = np.arange(10, dtype=np.float32)
        Y = np.arange(10, dtype=np.float32)
        out = saxpy(A, X, Y)

        for j in range(10):
            exp = A * X[j] + Y[j]
            self.assertTrue(exp == out[j], (exp, out[j]))

        A = np.arange(5, dtype=np.float32)
        X = np.arange(10, dtype=np.float32).reshape(5, 2)
        Y = np.arange(10, dtype=np.float32).reshape(5, 2)
        out = saxpy(A, X, Y)

        for j in range(5):
            for i in range(2):
                exp = A[j] * X[j, i] + Y[j, i]
                self.assertTrue(exp == out[j, i], (exp, out[j, i]))


if __name__ == '__main__':
    unittest.main()

[(165, 'unittest.main', 'unittest.main', False, 'import unittest\n'), (16, 'numba.roc.vectorizers.HsaGUFuncVectorize', 'HsaGUFuncVectorize', False, 'from numba.roc.vectorizers import HsaGUFuncVectorize\n'), (23, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (28, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (38, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (48, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (53, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (58, 'numpy.testing.assert_allclose', 'np.testing.assert_allclose', True, 'import numpy as np\n'), (68, 'numpy.testing.assert_allclose', 'np.testing.assert_allclose', True, 'import numpy as np\n'), (79, 'numba.roc.vectorizers.HsaGUFuncVectorize', 'HsaGUFuncVectorize', False, 'from numba.roc.vectorizers import HsaGUFuncVectorize\n'), (94, 'numba.roc.vectorizers.HsaGUFuncVectorize', 'HsaGUFuncVectorize', False, 'from numba.roc.vectorizers import HsaGUFuncVectorize\n'), (98, 'numpy.float32', 'np.float32', True, 'import numpy as np\n'), (108, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (109, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (116, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (129, 'numba.guvectorize', 'guvectorize', False, 'from numba import guvectorize\n'), (135, 'numpy.float32', 'np.float32', True, 'import numpy as np\n'), (145, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (146, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (153, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (24, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (30, 'numpy.dtype', 'np.dtype', True, 'import numpy as np\n'), (40, 'numpy.dtype', 'np.dtype', True, 'import numpy as np\n'), (50, 'numpy.dtype', 'np.dtype', True, 'import numpy as np\n'), (54, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (60, 'numpy.dtype', 'np.dtype', True, 'import numpy as np\n'), (70, 'numpy.dtype', 'np.dtype', True, 'import numpy as np\n'), (83, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (99, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (100, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (117, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (118, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (136, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (137, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (154, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (155, 'numpy.arange', 'np.arange', True, 'import numpy as np\n')]