# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

import numpy as np
from numba import cuda
from numba.cuda.testing import unittest, CUDATestCase
from cuda.core.experimental import Device
from numba.cuda.testing import skip_on_cudasim


class TestCudaEvent(CUDATestCase):
    def test_event_elapsed(self):
        N = 32
        dary = cuda.device_array(N, dtype=np.double)
        evtstart = cuda.event()
        evtend = cuda.event()

        evtstart.record()
        cuda.to_device(np.arange(N, dtype=np.double), to=dary)
        evtend.record()
        evtend.wait()
        evtend.synchronize()
        # Exercise the code path
        evtstart.elapsed_time(evtend)

    def test_event_elapsed_stream(self):
        stream = cuda.stream()
        self.event_elapsed_inner(stream)

    @skip_on_cudasim("Testing cuda.core events requires driver")
    def test_event_elapsed_cuda_core_stream(self):
        dev = Device()
        dev.set_current()
        stream = dev.create_stream()
        self.event_elapsed_inner(stream)

    def event_elapsed_inner(self, stream):
        N = 32
        dary = cuda.device_array(N, dtype=np.double)
        evtstart = cuda.event()
        evtend = cuda.event()

        evtstart.record(stream=stream)
        cuda.to_device(np.arange(N, dtype=np.double), to=dary, stream=stream)
        evtend.record(stream=stream)
        evtend.wait(stream=stream)
        evtend.synchronize()
        # Exercise the code path
        evtstart.elapsed_time(evtend)


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