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

import numpy as np

from numba import vectorize, cuda
from numba.tests.npyufunc.test_vectorize_decor import (
    BaseVectorizeDecor,
    BaseVectorizeNopythonArg,
    BaseVectorizeUnrecognizedArg,
)
from numba.cuda.testing import skip_on_cudasim, CUDATestCase
import unittest


@skip_on_cudasim("ufunc API unsupported in the simulator")
class TestVectorizeDecor(CUDATestCase, BaseVectorizeDecor):
    """
    Runs the tests from BaseVectorizeDecor with the CUDA target.
    """

    target = "cuda"


@skip_on_cudasim("ufunc API unsupported in the simulator")
class TestGPUVectorizeBroadcast(CUDATestCase):
    def test_broadcast(self):
        a = np.random.randn(100, 3, 1)
        b = a.transpose(2, 1, 0)

        def fn(a, b):
            return a - b

        @vectorize(["float64(float64,float64)"], target="cuda")
        def fngpu(a, b):
            return a - b

        expect = fn(a, b)
        got = fngpu(a, b)
        np.testing.assert_almost_equal(expect, got)

    def test_device_broadcast(self):
        """
        Same test as .test_broadcast() but with device array as inputs
        """

        a = np.random.randn(100, 3, 1)
        b = a.transpose(2, 1, 0)

        def fn(a, b):
            return a - b

        @vectorize(["float64(float64,float64)"], target="cuda")
        def fngpu(a, b):
            return a - b

        expect = fn(a, b)
        got = fngpu(cuda.to_device(a), cuda.to_device(b))
        np.testing.assert_almost_equal(expect, got.copy_to_host())


@skip_on_cudasim("ufunc API unsupported in the simulator")
class TestVectorizeNopythonArg(BaseVectorizeNopythonArg, CUDATestCase):
    def test_target_cuda_nopython(self):
        warnings = ["nopython kwarg for cuda target is redundant"]
        self._test_target_nopython("cuda", warnings)


@skip_on_cudasim("ufunc API unsupported in the simulator")
class TestVectorizeUnrecognizedArg(BaseVectorizeUnrecognizedArg, CUDATestCase):
    def test_target_cuda_unrecognized_arg(self):
        self._test_target_unrecognized_arg("cuda")


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