import numpy as np

from numba import cuda, vectorize, guvectorize
from numba.np.numpy_support import from_dtype
from numba.cuda.testing import CUDATestCase, skip_on_cudasim
import unittest


class TestCudaDateTime(CUDATestCase):
    def test_basic_datetime_kernel(self):
        @cuda.jit
        def foo(start, end, delta):
            for i in range(cuda.grid(1), delta.size, cuda.gridsize(1)):
                delta[i] = end[i] - start[i]

        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
        arr2 = arr1 + np.random.randint(0, 10000, arr1.size)
        delta = np.zeros_like(arr1, dtype="timedelta64[D]")

        foo[1, 32](arr1, arr2, delta)

        self.assertPreciseEqual(delta, arr2 - arr1)

    def test_scalar_datetime_kernel(self):
        @cuda.jit
        def foo(dates, target, delta, matches, outdelta):
            for i in range(cuda.grid(1), matches.size, cuda.gridsize(1)):
                matches[i] = dates[i] == target
                outdelta[i] = dates[i] - delta

        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
        target = arr1[5]  # datetime
        delta = arr1[6] - arr1[5]  # timedelta
        matches = np.zeros_like(arr1, dtype=np.bool_)
        outdelta = np.zeros_like(arr1, dtype="datetime64[D]")

        foo[1, 32](arr1, target, delta, matches, outdelta)
        where = matches.nonzero()

        self.assertEqual(list(where), [5])
        self.assertPreciseEqual(outdelta, arr1 - delta)

    @skip_on_cudasim("ufunc API unsupported in the simulator")
    def test_ufunc(self):
        datetime_t = from_dtype(np.dtype("datetime64[D]"))

        @vectorize([(datetime_t, datetime_t)], target="cuda")
        def timediff(start, end):
            return end - start

        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
        arr2 = arr1 + np.random.randint(0, 10000, arr1.size)

        delta = timediff(arr1, arr2)

        self.assertPreciseEqual(delta, arr2 - arr1)

    @skip_on_cudasim("ufunc API unsupported in the simulator")
    def test_gufunc(self):
        datetime_t = from_dtype(np.dtype("datetime64[D]"))
        timedelta_t = from_dtype(np.dtype("timedelta64[D]"))

        @guvectorize(
            [(datetime_t, datetime_t, timedelta_t[:])],
            "(),()->()",
            target="cuda",
        )
        def timediff(start, end, out):
            out[0] = end - start

        arr1 = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
        arr2 = arr1 + np.random.randint(0, 10000, arr1.size)

        delta = timediff(arr1, arr2)

        self.assertPreciseEqual(delta, arr2 - arr1)

    @skip_on_cudasim("no .copy_to_host() in the simulator")
    def test_datetime_view_as_int64(self):
        arr = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
        darr = cuda.to_device(arr)
        viewed = darr.view(np.int64)
        self.assertPreciseEqual(arr.view(np.int64), viewed.copy_to_host())
        self.assertEqual(viewed.gpu_data, darr.gpu_data)

    @skip_on_cudasim("no .copy_to_host() in the simulator")
    def test_timedelta_view_as_int64(self):
        arr = np.arange("2005-02", "2006-02", dtype="datetime64[D]")
        arr = arr - (arr - 1)
        self.assertEqual(arr.dtype, np.dtype("timedelta64[D]"))
        darr = cuda.to_device(arr)
        viewed = darr.view(np.int64)
        self.assertPreciseEqual(arr.view(np.int64), viewed.copy_to_host())
        self.assertEqual(viewed.gpu_data, darr.gpu_data)


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