import re

import numpy as np
from numba import cuda, int32, int64, float32, float64
from numba.cuda.testing import unittest, CUDATestCase, skip_on_cudasim
from numba.cuda.compiler import compile_ptx
from numba.core import config


def useful_syncwarp(ary):
    i = cuda.grid(1)
    if i == 0:
        ary[0] = 42
    cuda.syncwarp(0xFFFFFFFF)
    ary[i] = ary[0]


def use_shfl_sync_idx(ary, idx):
    i = cuda.grid(1)
    val = cuda.shfl_sync(0xFFFFFFFF, i, idx)
    ary[i] = val


def use_shfl_sync_up(ary, delta):
    i = cuda.grid(1)
    val = cuda.shfl_up_sync(0xFFFFFFFF, i, delta)
    ary[i] = val


def use_shfl_sync_down(ary, delta):
    i = cuda.grid(1)
    val = cuda.shfl_down_sync(0xFFFFFFFF, i, delta)
    ary[i] = val


def use_shfl_sync_xor(ary, xor):
    i = cuda.grid(1)
    val = cuda.shfl_xor_sync(0xFFFFFFFF, i, xor)
    ary[i] = val


def use_shfl_sync_with_val(ary, into):
    i = cuda.grid(1)
    val = cuda.shfl_sync(0xFFFFFFFF, into, 0)
    ary[i] = val


def use_vote_sync_all(ary_in, ary_out):
    i = cuda.grid(1)
    pred = cuda.all_sync(0xFFFFFFFF, ary_in[i])
    ary_out[i] = pred


def use_vote_sync_any(ary_in, ary_out):
    i = cuda.grid(1)
    pred = cuda.any_sync(0xFFFFFFFF, ary_in[i])
    ary_out[i] = pred


def use_vote_sync_eq(ary_in, ary_out):
    i = cuda.grid(1)
    pred = cuda.eq_sync(0xFFFFFFFF, ary_in[i])
    ary_out[i] = pred


def use_vote_sync_ballot(ary):
    i = cuda.threadIdx.x
    ballot = cuda.ballot_sync(0xFFFFFFFF, True)
    ary[i] = ballot


def use_match_any_sync(ary_in, ary_out):
    i = cuda.grid(1)
    ballot = cuda.match_any_sync(0xFFFFFFFF, ary_in[i])
    ary_out[i] = ballot


def use_match_all_sync(ary_in, ary_out):
    i = cuda.grid(1)
    ballot, pred = cuda.match_all_sync(0xFFFFFFFF, ary_in[i])
    ary_out[i] = ballot if pred else 0


def use_independent_scheduling(arr):
    i = cuda.threadIdx.x
    if i % 4 == 0:
        ballot = cuda.ballot_sync(0x11111111, True)
    elif i % 4 == 1:
        ballot = cuda.ballot_sync(0x22222222, True)
    elif i % 4 == 2:
        ballot = cuda.ballot_sync(0x44444444, True)
    elif i % 4 == 3:
        ballot = cuda.ballot_sync(0x88888888, True)
    arr[i] = ballot


def _safe_cc_check(cc):
    if config.ENABLE_CUDASIM:
        return True
    else:
        return cuda.get_current_device().compute_capability >= cc


@skip_on_cudasim("Warp Operations are not yet implemented on cudasim")
class TestCudaWarpOperations(CUDATestCase):
    def test_useful_syncwarp(self):
        compiled = cuda.jit("void(int32[:])")(useful_syncwarp)
        nelem = 32
        ary = np.empty(nelem, dtype=np.int32)
        compiled[1, nelem](ary)
        self.assertTrue(np.all(ary == 42))

    def test_shfl_sync_idx(self):
        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_idx)
        nelem = 32
        idx = 4
        ary = np.empty(nelem, dtype=np.int32)
        compiled[1, nelem](ary, idx)
        self.assertTrue(np.all(ary == idx))

    def test_shfl_sync_up(self):
        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_up)
        nelem = 32
        delta = 4
        ary = np.empty(nelem, dtype=np.int32)
        exp = np.arange(nelem, dtype=np.int32)
        exp[delta:] -= delta
        compiled[1, nelem](ary, delta)
        self.assertTrue(np.all(ary == exp))

    def test_shfl_sync_down(self):
        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_down)
        nelem = 32
        delta = 4
        ary = np.empty(nelem, dtype=np.int32)
        exp = np.arange(nelem, dtype=np.int32)
        exp[:-delta] += delta
        compiled[1, nelem](ary, delta)
        self.assertTrue(np.all(ary == exp))

    def test_shfl_sync_xor(self):
        compiled = cuda.jit("void(int32[:], int32)")(use_shfl_sync_xor)
        nelem = 32
        xor = 16
        ary = np.empty(nelem, dtype=np.int32)
        exp = np.arange(nelem, dtype=np.int32) ^ xor
        compiled[1, nelem](ary, xor)
        self.assertTrue(np.all(ary == exp))

    def test_shfl_sync_const_mode_val(self):
        # Test `mode` argument is constant in shfl_sync calls.
        # Related to https://github.com/NVIDIA/numba-cuda/pull/231
        subtest = [
            (use_shfl_sync_idx, 4),
            (use_shfl_sync_up, 4),
            (use_shfl_sync_down, 4),
            (use_shfl_sync_xor, 16),
        ]

        args_re = r"\((.*)\)"
        m = re.compile(args_re)

        for func, value in subtest:
            with self.subTest(func=func.__name__):
                compiled = cuda.jit("void(int32[:], int32)")(func)
                nelem = 32
                ary = np.empty(nelem, dtype=np.int32)
                compiled[1, nelem](ary, value)
                irs = next(iter(compiled.inspect_llvm().values()))

                for ir in irs.split("\n"):
                    if "call" in ir and "llvm.nvvm.shfl.sync.i32" in ir:
                        args = m.search(ir).group(0)
                        arglist = args.split(",")
                        mode_arg = arglist[1]
                        self.assertNotIn("%", mode_arg)

    def test_shfl_sync_const_mode_val_sm100(self):
        # Test shfl_sync compiles with cc=(10, 0)
        subtest = [
            use_shfl_sync_idx,
            use_shfl_sync_up,
            use_shfl_sync_down,
            use_shfl_sync_xor,
        ]

        for func in subtest:
            with self.subTest(func=func.__name__):
                compile_ptx(func, (int32[:], int32), cc=(10, 0))

    def test_shfl_sync_types(self):
        types = int32, int64, float32, float64
        values = (
            np.int32(-1),
            np.int64(1 << 42),
            np.float32(np.pi),
            np.float64(np.pi),
        )
        for typ, val in zip(types, values):
            with self.subTest(typ=typ):
                compiled = cuda.jit((typ[:], typ))(use_shfl_sync_with_val)
                nelem = 32
                ary = np.empty(nelem, dtype=val.dtype)
                compiled[1, nelem](ary, val)
                self.assertTrue(np.all(ary == val))

    def test_vote_sync_all(self):
        compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_all)
        nelem = 32
        ary_in = np.ones(nelem, dtype=np.int32)
        ary_out = np.empty(nelem, dtype=np.int32)
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 1))
        ary_in[-1] = 0
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 0))

    def test_vote_sync_any(self):
        compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_any)
        nelem = 32
        ary_in = np.zeros(nelem, dtype=np.int32)
        ary_out = np.empty(nelem, dtype=np.int32)
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 0))
        ary_in[2] = 1
        ary_in[5] = 1
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 1))

    def test_vote_sync_eq(self):
        compiled = cuda.jit("void(int32[:], int32[:])")(use_vote_sync_eq)
        nelem = 32
        ary_in = np.zeros(nelem, dtype=np.int32)
        ary_out = np.empty(nelem, dtype=np.int32)
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 1))
        ary_in[1] = 1
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 0))
        ary_in[:] = 1
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 1))

    def test_vote_sync_ballot(self):
        compiled = cuda.jit("void(uint32[:])")(use_vote_sync_ballot)
        nelem = 32
        ary = np.empty(nelem, dtype=np.uint32)
        compiled[1, nelem](ary)
        self.assertTrue(np.all(ary == np.uint32(0xFFFFFFFF)))

    @unittest.skipUnless(
        _safe_cc_check((7, 0)), "Matching requires at least Volta Architecture"
    )
    def test_match_any_sync(self):
        compiled = cuda.jit("void(int32[:], int32[:])")(use_match_any_sync)
        nelem = 10
        ary_in = np.arange(nelem, dtype=np.int32) % 2
        ary_out = np.empty(nelem, dtype=np.int32)
        exp = np.tile((0b0101010101, 0b1010101010), 5)
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == exp))

    @unittest.skipUnless(
        _safe_cc_check((7, 0)), "Matching requires at least Volta Architecture"
    )
    def test_match_all_sync(self):
        compiled = cuda.jit("void(int32[:], int32[:])")(use_match_all_sync)
        nelem = 10
        ary_in = np.zeros(nelem, dtype=np.int32)
        ary_out = np.empty(nelem, dtype=np.int32)
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 0b1111111111))
        ary_in[1] = 4
        compiled[1, nelem](ary_in, ary_out)
        self.assertTrue(np.all(ary_out == 0))

    @unittest.skipUnless(
        _safe_cc_check((7, 0)),
        "Independent scheduling requires at least Volta Architecture",
    )
    def test_independent_scheduling(self):
        compiled = cuda.jit("void(uint32[:])")(use_independent_scheduling)
        arr = np.empty(32, dtype=np.uint32)
        exp = np.tile((0x11111111, 0x22222222, 0x44444444, 0x88888888), 8)
        compiled[1, 32](arr)
        self.assertTrue(np.all(arr == exp))

    def test_activemask(self):
        @cuda.jit
        def use_activemask(x):
            i = cuda.grid(1)
            if (i % 2) == 0:
                # Even numbered threads fill in even numbered array entries
                # with binary "...01010101"
                x[i] = cuda.activemask()
            else:
                # Odd numbered threads fill in odd numbered array entries
                # with binary "...10101010"
                x[i] = cuda.activemask()

        out = np.zeros(32, dtype=np.uint32)
        use_activemask[1, 32](out)

        # 0x5 = 0101: The pattern from even-numbered threads
        # 0xA = 1010: The pattern from odd-numbered threads
        expected = np.tile((0x55555555, 0xAAAAAAAA), 16)
        np.testing.assert_equal(expected, out)

    def test_lanemask_lt(self):
        @cuda.jit
        def use_lanemask_lt(x):
            i = cuda.grid(1)
            x[i] = cuda.lanemask_lt()

        out = np.zeros(32, dtype=np.uint32)
        use_lanemask_lt[1, 32](out)

        # A string of 1s that grows from the LSB for each entry:
        # 0, 1, 3, 7, F, 1F, 3F, 7F, FF, 1FF, etc.
        # or in binary:
        # ...0001, ....0011, ...0111, etc.
        expected = np.asarray([(2**i) - 1 for i in range(32)], dtype=np.uint32)
        np.testing.assert_equal(expected, out)


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