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

from typing import List
from dataclasses import dataclass, field
from numba import cuda, float32
from numba.cuda.compiler import compile_ptx_for_current_device, compile_ptx
from math import cos, sin, tan, exp, log, log10, log2, pow, tanh
from operator import truediv
import numpy as np
from numba.cuda.testing import CUDATestCase, skip_on_cudasim, skip_unless_cc_75
import unittest


@dataclass
class FastMathCriterion:
    fast_expected: List[str] = field(default_factory=list)
    fast_unexpected: List[str] = field(default_factory=list)
    prec_expected: List[str] = field(default_factory=list)
    prec_unexpected: List[str] = field(default_factory=list)

    def check(self, test: CUDATestCase, fast: str, prec: str):
        test.assertTrue(all(i in fast for i in self.fast_expected))
        test.assertTrue(all(i not in fast for i in self.fast_unexpected))
        test.assertTrue(all(i in prec for i in self.prec_expected))
        test.assertTrue(all(i not in prec for i in self.prec_unexpected))


@skip_on_cudasim("Fastmath and PTX inspection not available on cudasim")
class TestFastMathOption(CUDATestCase):
    def _test_fast_math_common(self, pyfunc, sig, device, criterion):
        # Test jit code path
        fastver = cuda.jit(sig, device=device, fastmath=True)(pyfunc)
        precver = cuda.jit(sig, device=device)(pyfunc)

        criterion.check(
            self, fastver.inspect_asm(sig), precver.inspect_asm(sig)
        )

        # Test compile_ptx code path
        fastptx, _ = compile_ptx_for_current_device(
            pyfunc, sig, device=device, fastmath=True
        )
        precptx, _ = compile_ptx_for_current_device(pyfunc, sig, device=device)

        criterion.check(self, fastptx, precptx)

    def _test_fast_math_unary(self, op, criterion: FastMathCriterion):
        def kernel(r, x):
            r[0] = op(x)

        def device_function(x):
            return op(x)

        self._test_fast_math_common(
            kernel, (float32[::1], float32), device=False, criterion=criterion
        )
        self._test_fast_math_common(
            device_function, (float32,), device=True, criterion=criterion
        )

    def _test_fast_math_binary(self, op, criterion: FastMathCriterion):
        def kernel(r, x, y):
            r[0] = op(x, y)

        def device(x, y):
            return op(x, y)

        self._test_fast_math_common(
            kernel,
            (float32[::1], float32, float32),
            device=False,
            criterion=criterion,
        )
        self._test_fast_math_common(
            device, (float32, float32), device=True, criterion=criterion
        )

    def test_cosf(self):
        self._test_fast_math_unary(
            cos,
            FastMathCriterion(
                fast_expected=["cos.approx.ftz.f32 "],
                prec_unexpected=["cos.approx.ftz.f32 "],
            ),
        )

    def test_sinf(self):
        self._test_fast_math_unary(
            sin,
            FastMathCriterion(
                fast_expected=["sin.approx.ftz.f32 "],
                prec_unexpected=["sin.approx.ftz.f32 "],
            ),
        )

    def test_tanf(self):
        self._test_fast_math_unary(
            tan,
            FastMathCriterion(
                fast_expected=[
                    "sin.approx.ftz.f32 ",
                    "cos.approx.ftz.f32 ",
                    "div.approx.ftz.f32 ",
                ],
                prec_unexpected=["sin.approx.ftz.f32 "],
            ),
        )

    @skip_unless_cc_75
    def test_tanhf(self):
        self._test_fast_math_unary(
            tanh,
            FastMathCriterion(
                fast_expected=["tanh.approx.f32 "],
                prec_unexpected=["tanh.approx.f32 "],
            ),
        )

    def test_tanhf_compile_ptx(self):
        def tanh_kernel(r, x):
            r[0] = tanh(x)

        fastptx, _ = compile_ptx(
            tanh_kernel, (float32[::1], float32), fastmath=True
        )
        precptx, _ = compile_ptx(tanh_kernel, (float32[::1], float32))

        criterion = FastMathCriterion(
            fast_expected=["tanh.approx.f32 "],
            prec_unexpected=["tanh.approx.f32 "],
        )

        criterion.check(self, fastptx, precptx)

    def test_expf(self):
        self._test_fast_math_unary(
            exp,
            FastMathCriterion(
                fast_unexpected=["fma.rn.f32 "], prec_expected=["fma.rn.f32 "]
            ),
        )

    def test_logf(self):
        # Look for constant used to convert from log base 2 to log base e
        self._test_fast_math_unary(
            log,
            FastMathCriterion(
                fast_expected=["lg2.approx.ftz.f32 ", "0f3F317218"],
                prec_unexpected=["lg2.approx.ftz.f32 "],
            ),
        )

    def test_log10f(self):
        # Look for constant used to convert from log base 2 to log base 10
        self._test_fast_math_unary(
            log10,
            FastMathCriterion(
                fast_expected=["lg2.approx.ftz.f32 ", "0f3E9A209B"],
                prec_unexpected=["lg2.approx.ftz.f32 "],
            ),
        )

    def test_log2f(self):
        self._test_fast_math_unary(
            log2,
            FastMathCriterion(
                fast_expected=["lg2.approx.ftz.f32 "],
                prec_unexpected=["lg2.approx.ftz.f32 "],
            ),
        )

    def test_powf(self):
        self._test_fast_math_binary(
            pow,
            FastMathCriterion(
                fast_expected=["lg2.approx.ftz.f32 "],
                prec_unexpected=["lg2.approx.ftz.f32 "],
            ),
        )

    def test_divf(self):
        self._test_fast_math_binary(
            truediv,
            FastMathCriterion(
                fast_expected=["div.approx.ftz.f32 "],
                fast_unexpected=["div.rn.f32"],
                prec_expected=["div.rn.f32"],
                prec_unexpected=["div.approx.ftz.f32 "],
            ),
        )

    def test_divf_exception(self):
        def f10(r, x, y):
            r[0] = x / y

        sig = (float32[::1], float32, float32)
        fastver = cuda.jit(sig, fastmath=True, debug=True, opt=False)(f10)
        precver = cuda.jit(sig, debug=True, opt=False)(f10)
        nelem = 10
        ary = np.empty(nelem, dtype=np.float32)
        with self.assertRaises(ZeroDivisionError):
            precver[1, nelem](ary, 10.0, 0.0)

        try:
            fastver[1, nelem](ary, 10.0, 0.0)
        except ZeroDivisionError:
            self.fail("Divide in fastmath should not throw ZeroDivisionError")

    @unittest.expectedFailure
    def test_device_fastmath_propagation(self):
        # The fastmath option doesn't presently propagate to device functions
        # from their callees - arguably it should do, so this test is presently
        # an xfail.
        @cuda.jit("float32(float32, float32)", device=True)
        def foo(a, b):
            return a / b

        def bar(arr, val):
            i = cuda.grid(1)
            if i < arr.size:
                arr[i] = foo(i, val)

        sig = (float32[::1], float32)
        fastver = cuda.jit(sig, fastmath=True)(bar)
        precver = cuda.jit(sig)(bar)

        # Variants of the div instruction are further documented at:
        # https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#floating-point-instructions-div

        # The fast version should use the "fast, approximate divide" variant
        self.assertIn("div.approx.f32", fastver.inspect_asm(sig))
        # The precise version should use the "IEEE 754 compliant rounding"
        # variant, and neither of the "approximate divide" variants.
        self.assertIn("div.rn.f32", precver.inspect_asm(sig))
        self.assertNotIn("div.approx.f32", precver.inspect_asm(sig))
        self.assertNotIn("div.full.f32", precver.inspect_asm(sig))


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