"""
Custom decorators for JIT-compilation with Numba

Warning: complex decorators calling Numba may cause segmentation faults when profiled.
https://github.com/numba/numba/issues/3625
"""

import functools
import inspect
import logging
import typing as tp

import numba
import numpy as np

from . import options
from . import utils

logger = logging.getLogger(__name__)


def njit(func: tp.Callable | None = None, **kwargs):
    """Wrapper for numba.njit.

    May cause segmentation faults with profilers.
    """
    def _njit(func2):
        return numba.njit(func2, **options.NUMBA_OPTS, **kwargs)
    if func is None:
        return _njit
    return _njit(func)


def njit_if_numba_integrate(func: tp.Callable | None = None, **kwargs) -> tp.Callable:
    """Njit compile a function if Numba integration is enabled"""
    if func:
        return utils.conditional_decorator(numba.njit, options.NUMBA_INTEGRATE, **kwargs)(func)
    return utils.conditional_decorator(numba.njit, options.NUMBA_INTEGRATE, **kwargs)


def njit_module(**kwargs):
    """Adapted from numba.jit_module.

    May cause segmentation faults with profilers.
    """
    # Get the module jit_module is being called from
    frame = inspect.stack()[1]
    module = inspect.getmodule(frame[0])
    # Replace functions in module with jit-wrapped versions
    for name, obj in module.__dict__.items():
        if inspect.isfunction(obj) and inspect.getmodule(obj) == module:
            logger.debug(
                "Auto decorating function %s from module %s with jit and options: %s",
                obj, module.__name__, kwargs
            )
            module.__dict__[name] = numba.njit(obj, **options.NUMBA_OPTS, **kwargs)


def vectorize(**kwargs):
    """Extended version of numba.vectorize with support for NUMBA_DISABLE_JIT"""
    def vectorize_inner(func: tp.Callable):
        if options.NUMBA_DISABLE_JIT:
            # Using functools.wraps() ensures that docstrings etc. are preserved
            @functools.wraps(func)
            def wrapper(*func_args, **func_kwargs):
                # If called with scalars
                if not \
                        next((isinstance(arg, np.ndarray) for arg in func_args), False) or \
                        next((isinstance(arg, np.ndarray) for arg in func_kwargs.values()), False):
                    return func(*func_args, **func_kwargs)
                # If called with 0D arrays
                if not np.all([arg.ndim for arg in func_args] + [arg.ndim for arg in func_kwargs.values()]):
                    return func(
                        *[arg.item() for arg in func_args],
                        **{name: value.item() for name, value in func_kwargs.items()}
                    )
                # If called with arrays
                return np.array([
                    func(*i_args, **{name: value[i] for name, value in func_kwargs.items()})
                    for i, i_args in enumerate(zip(*func_args))
                ])
            return wrapper
        return functools.wraps(func)(numba.vectorize(**kwargs)(func))
    return vectorize_inner
