"""
Centralized backend availability checking for GPU acceleration libraries.

This module provides a single source of truth for checking whether CuPy and OpenCL
are available, avoiding duplicate import checks across the codebase.
"""

# CuPy availability and related objects
_cupy_available = False
_cupy_cp = None
_cupy_cp_signal = None
_cupy_streams_available = False
_cupy_stream = None
_cupy_torque_kernel = None
_cupy_kernel_available = False

# OpenCL availability and related objects
_opencl_available = False
_opencl_cl = None
_opencl_cla = None
_opencl_rfftn = None
_opencl_irfftn = None

# Try to import CuPy
try:
    import cupy as cp
    import cupyx.scipy.signal as cp_signal
    _cupy_cp = cp
    _cupy_cp_signal = cp_signal
    _cupy_available = True
    
    # Try to import CuPy streams for parallel execution
    try:
        from cupy.cuda import Stream
        _cupy_stream = Stream
        _cupy_streams_available = True
    except ImportError:
        _cupy_stream = None
        _cupy_streams_available = False
    
    # Create fused elementwise kernel for torque computation
    try:
        _cupy_torque_kernel = cp.ElementwiseKernel(
            'float32 r_x, float32 r_y, float32 grid_0, float32 grid_1, '
            'float32 max_0, float32 max_1, float32 min_0, float32 min_1',
            'float32 torque',
            '''
            float cross_p = (grid_0 - max_0) * r_y - (grid_1 - max_1) * r_x;
            float cross_n = (grid_0 - min_0) * r_y - (grid_1 - min_1) * r_x;
            torque = cross_p + cross_n;
            ''',
            'torque_kernel'
        )
        _cupy_kernel_available = True
    except Exception:
        _cupy_torque_kernel = None
        _cupy_kernel_available = False
        
except ImportError:
    pass

# Try to import OpenCL
try:
    import pyopencl as cl
    import pyopencl.array as cla
    from pyvkfft.fft import irfftn as cl_irfftn
    from pyvkfft.fft import rfftn as cl_rfftn
    _opencl_cl = cl
    _opencl_cla = cla
    _opencl_rfftn = cl_rfftn
    _opencl_irfftn = cl_irfftn
    _opencl_available = True
except ImportError:
    pass


def check_cupy_available():
    """Check if CuPy is available."""
    return _cupy_available


def check_opencl_available():
    """Check if OpenCL/pyvkfft is available."""
    return _opencl_available


def get_cupy():
    """Get CuPy module if available, None otherwise."""
    return _cupy_cp


def get_cupy_signal():
    """Get cupyx.scipy.signal module if available, None otherwise."""
    return _cupy_cp_signal


def get_cupy_stream():
    """Get CuPy Stream class if available, None otherwise."""
    return _cupy_stream


def get_cupy_torque_kernel():
    """Get CuPy torque kernel if available, None otherwise."""
    return _cupy_torque_kernel


def is_cupy_streams_available():
    """Check if CuPy streams are available."""
    return _cupy_streams_available


def is_cupy_kernel_available():
    """Check if CuPy torque kernel is available."""
    return _cupy_kernel_available


def get_opencl():
    """Get pyopencl module if available, None otherwise."""
    return _opencl_cl


def get_opencl_array():
    """Get pyopencl.array module if available, None otherwise."""
    return _opencl_cla


def get_opencl_rfftn():
    """Get pyvkfft.fft.rfftn function if available, None otherwise."""
    return _opencl_rfftn


def get_opencl_irfftn():
    """Get pyvkfft.fft.irfftn function if available, None otherwise."""
    return _opencl_irfftn


# For backward compatibility, export as module-level constants
CUPY_AVAILABLE = _cupy_available
OPENCL_AVAILABLE = _opencl_available
STREAMS_AVAILABLE = _cupy_streams_available
KERNEL_AVAILABLE = _cupy_kernel_available

# Export commonly used objects
cp = _cupy_cp
cp_signal = _cupy_cp_signal
Stream = _cupy_stream
_torque_kernel = _cupy_torque_kernel
cl = _opencl_cl
cla = _opencl_cla
cl_rfftn = _opencl_rfftn
cl_irfftn = _opencl_irfftn