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

"""
Separate module with function samples for serialization tests,
to avoid issues with __main__.
"""

import math
from math import sqrt
import numpy as np
import numpy.random as nprand

from numba import jit
from numba.core import types


@jit((types.int32, types.int32))
def add_with_sig(a, b):
    return a + b


@jit
def add_without_sig(a, b):
    return a + b


@jit(nopython=True)
def add_nopython(a, b):
    return a + b


@jit(nopython=True)
def add_nopython_fail(a, b):
    object()
    return a + b


def closure(a):
    @jit(nopython=True)
    def inner(b, c):
        return a + b + c

    return inner


K = 3.0


def closure_with_globals(x, **jit_args):
    @jit(**jit_args)
    def inner(y):
        # Exercise a builtin function and a module-level constant
        k = max(K, K + 1)
        # Exercise two functions from another module, one accessed with
        # dotted notation, one imported explicitly.
        return math.hypot(x, y) + sqrt(k)

    return inner


@jit(nopython=True)
def other_function(x, y):
    return math.hypot(x, y)


@jit(forceobj=True)
def get_global_objmode(x):
    return K * x


@jit(nopython=True)
def get_renamed_module(x):
    nprand.seed(42)
    return np.cos(x), nprand.random()


def closure_calling_other_function(x):
    @jit(nopython=True)
    def inner(y, z):
        return other_function(x, y) + z

    return inner


def closure_calling_other_closure(x):
    @jit(nopython=True)
    def other_inner(y):
        return math.hypot(x, y)

    @jit(nopython=True)
    def inner(y):
        return other_inner(y) + x

    return inner


# A dynamic function calling a builtin function
def _get_dyn_func(**jit_args):
    code = """
        def dyn_func(x):
            res = 0
            for i in range(x):
                res += x
            return res
        """
    ns = {}
    exec(code.strip(), ns)
    return jit(**jit_args)(ns["dyn_func"])


dyn_func = _get_dyn_func(nopython=True)
dyn_func_objmode = _get_dyn_func(forceobj=True)
