import numpy as np
from numba import cuda
from numba.cuda.kernels.transpose import transpose
from numba.cuda.testing import unittest
from numba.cuda.testing import skip_on_cudasim, CUDATestCase


recordwith2darray = np.dtype([("i", np.int32), ("j", np.float32, (3, 2))])


@skip_on_cudasim("Device Array API unsupported in the simulator")
class TestTranspose(CUDATestCase):
    def test_transpose(self):
        variants = (
            (5, 6, np.float64),
            (128, 128, np.complex128),
            (1025, 512, np.float64),
        )

        for rows, cols, dtype in variants:
            with self.subTest(rows=rows, cols=cols, dtype=dtype):
                x = np.arange(rows * cols, dtype=dtype).reshape(cols, rows)
                y = np.zeros(rows * cols, dtype=dtype).reshape(rows, cols)
                dx = cuda.to_device(x)
                dy = cuda.cudadrv.devicearray.from_array_like(y)
                transpose(dx, dy)
                dy.copy_to_host(y)
                np.testing.assert_array_equal(x.transpose(), y)

    small_variants = (
        (2, 3),
        (16, 16),
        (16, 17),
        (17, 16),
        (14, 15),
        (15, 14),
        (14, 14),
    )

    def test_transpose_record(self):
        for rows, cols in self.small_variants:
            with self.subTest(rows=rows, cols=cols):
                arr = np.recarray((rows, cols), dtype=recordwith2darray)
                for x in range(rows):
                    for y in range(cols):
                        arr[x, y].i = x**2 + y
                        j = np.arange(3 * 2, dtype=np.float32)
                        arr[x, y].j = j.reshape(3, 2) * x + y

                transposed = arr.T
                d_arr = cuda.to_device(arr)
                d_transposed = cuda.device_array_like(transposed)
                transpose(d_arr, d_transposed)
                host_transposed = d_transposed.copy_to_host()
                np.testing.assert_array_equal(transposed, host_transposed)

    def test_transpose_bool(self):
        for rows, cols in self.small_variants:
            with self.subTest(rows=rows, cols=cols):
                arr = np.random.randint(2, size=(rows, cols), dtype=np.bool_)
                transposed = arr.T

                d_arr = cuda.to_device(arr)
                d_transposed = cuda.device_array_like(transposed)
                transpose(d_arr, d_transposed)

                host_transposed = d_transposed.copy_to_host()
                np.testing.assert_array_equal(transposed, host_transposed)

    def test_transpose_view(self):
        # Because the strides of transposes of views differ to those in NumPy
        # (see issue #4974), we test the shape and strides of a transpose.
        a = np.arange(120, dtype=np.int64).reshape((10, 12))
        a_view_t = a[::2, ::2].T

        d_a = cuda.to_device(a)
        d_a_view_t = d_a[::2, ::2].T

        self.assertEqual(d_a_view_t.shape, (6, 5))
        self.assertEqual(d_a_view_t.strides, (40, 8))

        h_a_view_t = d_a_view_t.copy_to_host()
        np.testing.assert_array_equal(a_view_t, h_a_view_t)


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