import numpy
import timeit
import sys
# import cv2
from scipy import signal, spatial

from ..backend_availability import (
    CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
    cp, cp_signal, _torque_kernel,
    cl, cla, cl_rfftn, cl_irfftn
)
from .timing import Timer

'''
# ------------------------------------------------------------------------------
 Poisson Solver for 1D, 2D, and 3D scalar fields.
    Data may be defined on a rectlinear grid or unstructured points.

  1) create an object
        PoissonSolver(args)
            'solver'      : type of solver
                            'S' for spatial convolution
                                    only option for unstructured points
                            'F' for convolution in frequency space
                                    default option for regular grids

         --> for regular grid
            'grid'        : shape of the rectlinear grid
                                tuple of dimensions (Z,Y,X)
            'spacings'    : spacings of the rectlinear grid
                                tuple of spacings (dz,dy,dx)

         --> for unstrucutred points
            'points'      : unstructured points
                                numpy array of shape (N, dim)
            'pvolumes'    : point volumes for each point
                                numpy array of shape (N, 1)

  2) call prepare()
          this initializes internal datastructures depending upon the mesh.

          Once initialized, solve() can be called any number of times to
          solve the Poisson equation on the same domain. If the domain
          changes (e.g., a Lagrangian mesh), create a new object.

  3) call solve(f), where, f is a numpy array defining a scalar field

          for rectlinear grid,     f.shape = grid.shape
          for unstructured points, f.shape = (points.shape[0], 1)

# ------------------------------------------------------------------------------
'''

class PoissonSolver(object):

    # constructor
    def __init__(self, **kwargs):

        args = kwargs.keys()

        if ('grid' in args) == ('points' in args):
            raise SyntaxError("Poisson Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid")

        if ('grid' in args) != ('spacings' in args):
            raise SyntaxError("Poisson Solver needs spacings for regular grid")

        if ('points' in args) != ('pvolumes' in args):
            raise SyntaxError("Poisson Solver needs point volumes for unstructured points")

        # (default) data type
        self.dtype = numpy.float32
        if 'dtype' in args:
            self.dtype = kwargs['dtype']

        if 'grid' in args:
            self.ptype = 'G'
            self.stype = 'F'             # default
            self.gdims = kwargs['grid']
            self.gdx = kwargs['spacings']
            self.dim = len(self.gdims)
            if self.dim != len(self.gdx):
                raise ValueError("Dimensions of spacings should match that of the grid")

        elif 'points' in args:
            self.ptype = 'P'
            self.stype = 'S'            # only option
            self.points = kwargs['points']
            self.pvolumes = kwargs['pvolumes']
            self.dim = len(self.points.shape)
            if self.pvolumes.shape[0] != self.points.shape[0]:
                raise ValueError("Number of pvolumes should match that of the points")

        if 'solver' in args:
            if (kwargs['solver'] != 'S' and kwargs['solver'] != 'F'):
                raise SyntaxError('Solver can be only of type S(patial) or F(requency)')

            if (kwargs['solver'] == 'F') and ('points' in args):
                raise SyntaxError('Only Spatial solver supported for unstructured points')

            if (kwargs['solver'] == 'S') and ('grid' in args):
                raise SyntaxError('Frequncy solver strongly suggested for structured grids')

            self.stype = kwargs['solver']
        if 'poisson_backend' in args:
            self.poisson_backend = kwargs['poisson_backend']
            if self.poisson_backend not in ('cpu', 'opencl', 'cupy'):
                raise ValueError("poisson_backend must be one of: 'cpu', 'opencl', 'cupy'")

        if (self.dim > 3):
            raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions")

        self.ready = False
        if self.poisson_backend == 'cupy':
            self._gpu_cache = None  # Cache for GPU arrays
        if self.poisson_backend == 'opencl':
            self._opencl_cache = None  # Cache for OpenCL FFT buffers
            self._opencl_queue = None  # OpenCL command queue
            self._opencl_context = None  # OpenCL context (set when queue is created)
            self._opencl_crop_kernel = None  # Cached OpenCL crop kernel
            self._opencl_scatter_kernel = None  # Cached scatter (reverse-crop) kernel
            self._opencl_fused_mul_kernel = None  # Cached fused complex multiply kernel
            self._opencl_scatter_batch_kernel = None  # Cached batch scatter kernel
            self._opencl_crop_batch_kernel = None  # Cached batch crop kernel
            self._opencl_batch_mul_kernel = None  # Cached batch complex multiply kernel
            self._opencl_last_timing = None  # Last per-stage OpenCL timing in ms
        # Cache for OpenCV FFT (CPU backend)
        self._opencv_cache = None  # Cache for OpenCV FFT buffers
        # if (self.ptype == 'G'):
        #     print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx)
        # elif (self.ptype == 'P'):
        #     print( 'PoissonSolver', self.dim, 'D points =', self.points.shape)

    # --------------------------------------------------------------------------
    # create the Green's function in appropriate representation
    # --------------------------------------------------------------------------

    # generate a grid -(N-1) to (N-1) to define distance function
    #     where N is the size of the original grid
    def _create_radialGrid(self, verbose):

        mtimer = Timer()
        if verbose:
            print( '  - creating distance kernel:', end ='')

        rdims = tuple(2*d - 1 for d in self.gdims)

        R = numpy.indices(rdims, dtype=self.dtype)
        for d in range(self.dim):
            R[d] += 1-self.gdims[d]
            R[d] *= self.gdx[d]

        if self.dim == 1:
            numpy.absolute(R[0], R[0])
            R = R[0]

        elif self.dim == 2:
            numpy.hypot(R[1], R[0], R[0])
            R = R[0]

            # half of the average grid spacing
            zval = 0.5 * (self.gdx[0]+self.gdx[1]) / 2.0
            R[self.gdims[0]-1, self.gdims[1]-1] = zval

        elif self.dim == 3:
            numpy.hypot(R[2], R[1], R[1])
            numpy.hypot(R[1], R[0], R[0])
            R = R[0]

            # half of the average grid spacing
            zval = 0.5 * (self.gdx[0]+self.gdx[1]+self.gdx[2]) / 3.0
            R[self.gdims[0]-1, self.gdims[1]-1, self.gdims[2]-1] = zval

        if verbose:
            print( R.shape, R.min(), R.max())
            mtimer.end()

        return R

    # --------------------------------------------------------------------------
    # scale a function by point volumes for correct integration
    def _scale(self, func):

        if self.ptype == 'P':
            return numpy.multiply(func, self.pvolumes)

        elif self.ptype == 'G':
            sfunc = func * numpy.prod(self.gdx)

            '''
            # volume of the boundary elements is only half the voxel
            if self.dim == 1:
                sfunc[ 0] *= 0.5
                sfunc[-1] *= 0.5
            elif self.dim == 2:
                sfunc[ 0, :] *= 0.5
                sfunc[-1, :] *= 0.5
                sfunc[ :, 0] *= 0.5
                sfunc[ :,-1] *= 0.5
            elif self.dim == 3:
                sfunc[ 0, :, :] *= 0.5
                sfunc[-1, :, :] *= 0.5
                sfunc[ :, 0, :] *= 0.5
                sfunc[ :,-1, :] *= 0.5
                sfunc[ :, :, 0] *= 0.5
                sfunc[ :, :,-1] *= 0.5
            '''
        return sfunc

    def prepare(self, verbose=False):

        # ----------------------------------------------------------------------
        gtimer = Timer()
        if verbose:
            print( '\nInitializing Poisson solver, type =', self.stype)

        # ----------------------------------------------------------------------
        # generate a grid to store pairwise distances to store Green's function
        self.G = None                      # definition depends upon solver type

        # ----------------------------------------------------------------------
        # for Fourier solution,
        if self.stype == 'F':

            if self.ptype == 'G':
                self.G = self._create_radialGrid(verbose)

            elif self.ptype == 'P':
                raise SyntaxError('Frequncy solver not supported for unstructured points')
                #self.G = create_distArray()

        # ----------------------------------------------------------------------
        # for spatial solution,
            # use the pairwise distance matrix, N x N
            # where N is the number of points
        elif self.stype == 'S':

            if self.ptype == 'G':
                raise SyntaxError('Spatial solver not suggested for structured grids')
                #self.G = self.create_points(verbose)

            elif self.ptype == 'P':
                # for 1D points, add an extra axis
                if (self.dim == 1) and (len(self.points.shape) == 1):
                    self.points = self.points[:,numpy.newaxis]

                # create pairwise distance matrix now
                ltimer = Timer()
                self.G = spatial.distance.cdist(self.points, self.points)

                # TODO: use a good approximation here
                zval = numpy.power(10.0, -10.0)
                self.G += zval * numpy.identity(self.G.shape[0])

                if verbose:
                    print('  - created pairwise distance matrix:', self.G.shape, self.G.min(), self.G.max(), end = "")
                    ltimer.end()

        # ----------------------------------------------------------------------
        # compute the Green's function
        ltimer = Timer()
        if verbose:
            print( '  - computing the Green\'s function:')

        if self.dim == 1:
            numpy.multiply(self.G, 0.5, self.G)

        elif self.dim == 2:
            numpy.log(self.G, self.G)
            numpy.multiply(self.G, (0.5 / numpy.pi), self.G)

        elif self.dim == 3:
            numpy.reciprocal(self.G, self.G)
            numpy.multiply(self.G, (-0.25 / numpy.pi), self.G)

        if verbose:
            print( self.G.shape, self.G.min(), self.G.max(), end = "")
            ltimer.end()

        # ----------------------------------------------------------------------
        self.ready = True

        if verbose:
            print( 'Poisson solver initialized', end ='')
            gtimer.end()

    # --------------------------------------------------------------------------
    # compute the integral solution
    # --------------------------------------------------------------------------
    def solve(self, f, verbose=False):

        if not self.ready:
            self.prepare(verbose)

        fshape = f.shape

        gtimer = Timer()
        if verbose:
            print( '\nSolving Poisson Eq.')

        # ----------------------------------------------------------------------
        # regular grid
        if self.ptype == 'G':

            #if (self.gdims - fshape).any():
            if self.gdims != fshape:
                print( self.gdims, fshape)
                raise ValueError("Shape of function should match shape of grid")

            # convolution in frequency domain using OpenCV FFT (2D) or numpy FFT (3D)
            if self.stype == 'F':
                if self.dim == 2:
                    # # Use OpenCV DFT for 2D (optimized with caching)
                    # # Convert to float32 for OpenCV (use ascontiguousarray for better performance)
                    # f_32 = numpy.ascontiguousarray(f, dtype=numpy.float32)
                    
                    # # For 'same' mode convolution, pad to (2*N-1) size
                    # full_shape = tuple(2 * s - 1 for s in f_32.shape)
                    
                    # # Initialize or validate cache
                    # cache_valid = (
                    #     self._opencv_cache is not None and
                    #     self._opencv_cache['fshape'] == fshape and
                    #     self._opencv_cache['full_shape'] == full_shape
                    # )
                    
                    # if not cache_valid:
                    #     # Cache G_fft (G doesn't change, so compute once)
                    #     G_32 = numpy.ascontiguousarray(self.G, dtype=numpy.float32)
                    #     G_fft = cv2.dft(G_32, flags=cv2.DFT_COMPLEX_OUTPUT)
                        
                    #     # Pre-compute crop slices for 'same' mode extraction
                    #     # For circular convolution at size (2*H-1, 2*W-1), the center corresponds to
                    #     # position [H-1, W-1], which matches scipy's fftconvolve mode='same' extraction
                    #     starts = [fshape[i] - 1 for i in range(len(fshape))]
                    #     slices = tuple(slice(starts[i], starts[i] + fshape[i]) for i in range(len(fshape)))
                        
                    #     # Pre-allocate reusable buffers
                    #     self._opencv_cache = {
                    #         'fshape': fshape,
                    #         'full_shape': full_shape,
                    #         'G_fft': G_fft,  # Cached: G doesn't change, so compute FFT once
                    #         'slices': slices,  # Cached: crop slices pre-computed
                    #         'gdx_prod': numpy.prod(self.gdx),  # Cached: grid spacing product
                    #         # Reusable buffers (avoid reallocation on each call)
                    #         'f_pad': numpy.zeros(full_shape, dtype=numpy.float32),  # Reused for padding
                    #         'conv_fft': numpy.empty((full_shape[0], full_shape[1], 2), dtype=numpy.float32),  # Output for mulSpectrums
                    #     }
                    
                    # cache = self._opencv_cache
                    
                    # # Zero-pad f into pre-allocated buffer (faster than creating new array)
                    # f_pad = cache['f_pad']
                    # f_pad.fill(0.0)  # Clear previous data efficiently
                    # f_slices = tuple(slice(0, s) for s in f_32.shape)
                    # f_pad[f_slices] = f_32
                    
                    # # Compute FFT of f (cv2.dft returns new array, but we can reuse for mulSpectrums)
                    # f_fft = cv2.dft(f_pad, flags=cv2.DFT_COMPLEX_OUTPUT)
                    
                    # # Multiply in frequency domain using OpenCV's optimized mulSpectrums
                    # # This is faster than manual complex multiplication and supports output buffer
                    # # Signature: mulSpectrums(a, b, flags, c=None, conjB=False)
                    # cv2.mulSpectrums(f_fft, cache['G_fft'], cv2.DFT_COMPLEX_OUTPUT, 
                    #                 cache['conv_fft'], False)
                    
                    # # Compute inverse FFT (cv2.idft returns new array)
                    # conv_full = cv2.idft(cache['conv_fft'], flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT)
                    
                    # # Extract 'same' mode result (use cached slices, scale in-place)
                    # p = conv_full[cache['slices']].copy()
                    # numpy.multiply(p, cache['gdx_prod'], p)
                    # from scipy import signal
                    p = signal.fftconvolve(f, self.G, mode='same')
                    # Scale by grid spacing
                    numpy.multiply(p, numpy.prod(self.gdx), p)
                else:
                    # For 3D, use numpy FFT (OpenCV DFT only supports 2D)
                    # from scipy import signal
                    p = signal.fftconvolve(f, self.G, mode='same')
                    # Scale by grid spacing
                    numpy.multiply(p, numpy.prod(self.gdx), p)

            # convolution in spatial domain
            elif self.stype == 'S':
                raise SyntaxError('Spatial solver not suggested for structured grids')
                #p = (self.G).dot(f.reshape(-1)).reshape(fshape)

        # ----------------------------------------------------------------------
        # unstructured points
        elif self.ptype == 'P':

            if (self.points.shape[0] != fshape[0]) or (1 != len(fshape)):
                raise ValueError("Shape of function should be (N,1) with N as the number of points")

            # convolution in spatial domain
                # Dec 13, 2017 -- pre scaling works for unstructured grids!
            if self.stype == 'S':
                p = numpy.multiply(f, self.pvolumes)
                numpy.dot(self.G, p, p)
                #numpy.multiply(p, self.pvolumes, p)

            elif self.stype == 'F':
                raise SyntaxError('Frequncy solver not supported for unstructured points')

        # ----------------------------------------------------------------------
        if verbose:
            print( 'Poisson solver finished')
            gtimer.end()

        return p

    # --------------------------------------------------------------------------
    # GPU-accelerated solve using CuPy
    # --------------------------------------------------------------------------
    def solve_gpu(self, f, verbose=False, return_gpu=False):
        """
        GPU-accelerated Poisson solver using CuPy.
        Falls back to CPU if CuPy is not available.
        
        Args:
            f: Input field (numpy array)
            verbose: Print timing information
            
        Returns:
            p: Solution (numpy array on CPU)
        """
        if not CUPY_AVAILABLE:
            if verbose:
                print("CuPy not available, falling back to CPU solver")
            return self.solve(f, verbose)
        
        if not self.ready:
            self.prepare(verbose)
        
        fshape = f.shape
        
        gtimer = Timer()
        if verbose:
            print('\nSolving Poisson Eq. on GPU')
        
        # ----------------------------------------------------------------------
        # regular grid
        if self.ptype == 'G':
            if self.gdims != fshape:
                print(self.gdims, fshape)
                raise ValueError("Shape of function should match shape of grid")
            
            # convolution in frequency domain using GPU
            if self.stype == 'F':
                # Initialize GPU cache if needed
                if self._gpu_cache is None:
                    if verbose:
                        print('  - Transferring Green\'s function to GPU...')
                    self._gpu_cache = {
                        'G_gpu': cp.asarray(self.G, dtype=cp.float32),
                        'gdx_prod': numpy.prod(self.gdx)
                    }
                
                # Transfer input to GPU
                f_gpu = cp.asarray(f, dtype=cp.float32)
                
                # Perform FFT convolution on GPU
                # CuPy's fftconvolve is GPU-accelerated
                p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
                
                # Scale by grid spacing (on GPU)
                p_gpu *= self._gpu_cache['gdx_prod']

                if verbose:
                    print('Poisson solver (GPU) finished')
                    gtimer.end()
                
                # Transfer result back to CPU
                if return_gpu:
                    return p_gpu  # Return CuPy array
                return cp.asnumpy(p_gpu)
                # p = cp.asnumpy(p_gpu)
                
            elif self.stype == 'S':
                raise SyntaxError('Spatial solver not suggested for structured grids')
        
    def _solve_gpu_internal(self, f_gpu):
        """Internal GPU solve that accepts and returns GPU arrays."""
        """
        GPU-accelerated Poisson solver using CuPy.
        Falls back to CPU if CuPy is not available.
        
        Args:
            f_gpu: Input field (numpy array)
            
        Returns:
            p_gpu: Solution (CuPy array)
        """     
        
        # ----------------------------------------------------------------------
        # regular grid
        if self.ptype == 'G':
            
            # convolution in frequency domain using GPU
            if self.stype == 'F':
                # Initialize GPU cache if needed
                if self._gpu_cache is None:
                    self._gpu_cache = {
                        'G_gpu': cp.asarray(self.G, dtype=cp.float32),
                        'gdx_prod': numpy.prod(self.gdx)
                    }
                
                # Perform FFT convolution on GPU
                # CuPy's fftconvolve is GPU-accelerated
                p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
                
                # Scale by grid spacing (on GPU)
                p_gpu *= self._gpu_cache['gdx_prod']               

                return p_gpu  # Return CuPy array

    def clear_gpu_cache(self):
        """Clear GPU cache to free GPU memory"""
        if self._gpu_cache is not None:
            self._gpu_cache = None
            if CUPY_AVAILABLE:
                cp.get_default_memory_pool().free_all_blocks()

        # ----------------------------------------------------------------------
        # unstructured points - not optimized for GPU
        # elif self.ptype == 'P':
        #     if verbose:
        #         print('  - GPU solver not optimized for unstructured points, using CPU')
        #     return self.solve(f, verbose)
        
        # ----------------------------------------------------------------------

        # return p

    def _get_opencl_queue(self):
        """Create and cache an OpenCL queue on first available GPU."""
        if self._opencl_queue is not None:
            return self._opencl_queue

        if not OPENCL_AVAILABLE:
            raise RuntimeError("OpenCL/pyvkfft is not available")

        ctx = None
        for platform in cl.get_platforms():
            for device in platform.get_devices():
                if device.type & cl.device_type.GPU == 0:
                    continue
                ctx = cl.Context(devices=(device,))
                break
            if ctx is not None:
                break

        if ctx is None:
            raise RuntimeError("No OpenCL GPU device found")

        self._opencl_queue = cl.CommandQueue(ctx)
        self._opencl_context = ctx
        return self._opencl_queue
    
    def _get_opencl_crop_kernel(self, queue):
        """Create and cache OpenCL kernel for cropping convolution result."""
        if hasattr(self, '_opencl_crop_kernel') and self._opencl_crop_kernel is not None:
            return self._opencl_crop_kernel
        
        # Simple 2D copy kernel
        kernel_code = """
        __kernel void crop_and_scale(
            __global const float* src,
            __global float* dst,
            const int src_width,
            const int dst_width,
            const int dst_height,
            const int start_x,
            const int start_y,
            const float scale
        ) {
            int gid_x = get_global_id(0);
            int gid_y = get_global_id(1);
            
            if (gid_x < dst_width && gid_y < dst_height) {
                int src_idx = (start_y + gid_y) * src_width + (start_x + gid_x);
                int dst_idx = gid_y * dst_width + gid_x;
                dst[dst_idx] = src[src_idx] * scale;
            }
        }
        """
        
        program = cl.Program(self._opencl_context, kernel_code).build()
        self._opencl_crop_kernel = program.crop_and_scale
        return self._opencl_crop_kernel

    def _get_opencl_scatter_kernel(self, queue):
        """Kernel that copies unpadded input into a zero-padded device buffer (reverse of crop)."""
        if self._opencl_scatter_kernel is not None:
            return self._opencl_scatter_kernel

        kernel_code = """
        __kernel void scatter_to_padded(
            __global const float* src,
            __global float* dst,
            const int src_width,
            const int src_height,
            const int dst_width
        ) {
            int gid_x = get_global_id(0);
            int gid_y = get_global_id(1);
            if (gid_x < src_width && gid_y < src_height) {
                dst[gid_y * dst_width + gid_x] = src[gid_y * src_width + gid_x];
            }
        }
        """
        program = cl.Program(self._opencl_context, kernel_code).build()
        self._opencl_scatter_kernel = program.scatter_to_padded
        return self._opencl_scatter_kernel

    def _get_opencl_fused_mul_kernel(self, queue):
        """Kernel that multiplies two complex spectra by one Green's spectrum in a single pass."""
        if self._opencl_fused_mul_kernel is not None:
            return self._opencl_fused_mul_kernel

        kernel_code = """
        __kernel void fused_complex_mul(
            __global float2* fft1,
            __global float2* fft2,
            __global const float2* g,
            const int n
        ) {
            int i = get_global_id(0);
            if (i < n) {
                float2 gv = g[i];
                float2 a = fft1[i];
                float2 b = fft2[i];
                fft1[i] = (float2)(a.x * gv.x - a.y * gv.y,
                                   a.x * gv.y + a.y * gv.x);
                fft2[i] = (float2)(b.x * gv.x - b.y * gv.y,
                                   b.x * gv.y + b.y * gv.x);
            }
        }
        """
        program = cl.Program(self._opencl_context, kernel_code).build()
        self._opencl_fused_mul_kernel = program.fused_complex_mul
        return self._opencl_fused_mul_kernel

    def _get_opencl_scatter_batch_kernel(self, queue):
        """Batched scatter: copies two unpadded inputs into a stacked zero-padded device buffer."""
        if self._opencl_scatter_batch_kernel is not None:
            return self._opencl_scatter_batch_kernel

        kernel_code = """
        __kernel void scatter_batch(
            __global const float* src,
            __global float* dst,
            const int src_width,
            const int src_height,
            const int dst_width,
            const int src_plane,
            const int dst_plane
        ) {
            int gid_x = get_global_id(0);
            int gid_y = get_global_id(1);
            int gid_z = get_global_id(2);
            if (gid_x < src_width && gid_y < src_height) {
                int src_idx = gid_z * src_plane + gid_y * src_width + gid_x;
                int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
                dst[dst_idx] = src[src_idx];
            }
        }
        """
        program = cl.Program(self._opencl_context, kernel_code).build()
        self._opencl_scatter_batch_kernel = program.scatter_batch
        return self._opencl_scatter_batch_kernel

    def _get_opencl_crop_batch_kernel(self, queue):
        """Batched crop+scale: extracts 'same' region from two stacked full-conv results."""
        if self._opencl_crop_batch_kernel is not None:
            return self._opencl_crop_batch_kernel

        kernel_code = """
        __kernel void crop_batch(
            __global const float* src,
            __global float* dst,
            const int src_width,
            const int dst_width,
            const int dst_height,
            const int start_x,
            const int start_y,
            const float scale,
            const int src_plane,
            const int dst_plane
        ) {
            int gid_x = get_global_id(0);
            int gid_y = get_global_id(1);
            int gid_z = get_global_id(2);
            if (gid_x < dst_width && gid_y < dst_height) {
                int src_idx = gid_z * src_plane + (start_y + gid_y) * src_width + (start_x + gid_x);
                int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
                dst[dst_idx] = src[src_idx] * scale;
            }
        }
        """
        program = cl.Program(self._opencl_context, kernel_code).build()
        self._opencl_crop_batch_kernel = program.crop_batch
        return self._opencl_crop_batch_kernel

    def _get_opencl_batch_mul_kernel(self, queue):
        """Batched complex multiply: multiplies each batch plane by the same Green's spectrum.
        Uses multiple elements per work item, __constant for broadcast-read Green's spectrum,
        and fma for better throughput/precision."""
        if self._opencl_batch_mul_kernel is not None:
            return self._opencl_batch_mul_kernel

        kernel_code = """
        __kernel void batch_complex_mul(
            __global float2* fft_batch,
            __global const float2* g,
            const int n
        ) {
            const int b = get_global_id(0);
            const int base = b * n;
            const int i0 = get_global_id(1) * 4;
            const int i_end = min(i0 + 4, n);
            for (int i = i0; i < i_end; i++) {
                float2 gv = g[i];
                float2 a = fft_batch[base + i];
                fft_batch[base + i] = (float2)(fma(a.x, gv.x, -(a.y * gv.y)),
                                               fma(a.x, gv.y,   a.y * gv.x));
            }
        }
        """
        program = cl.Program(self._opencl_context, kernel_code).build()
        self._opencl_batch_mul_kernel = program.batch_complex_mul
        return self._opencl_batch_mul_kernel

    def _get_opencl_alloc_host_ptr(self):
        """Return an allocator that creates buffers with ALLOC_HOST_PTR for pinned host access (reduces H2D/D2H transfer cost)."""
        ctx = self._opencl_context
        mf = cl.mem_flags

        def allocator(nbytes):
            return cl.Buffer(ctx, mf.READ_WRITE | mf.ALLOC_HOST_PTR, size=nbytes)

        return allocator

    @staticmethod
    def _center_crop_same(full_conv, target_shape):
        """Crop full convolution result to scipy 'same' output shape."""
        starts = [(full_conv.shape[i] - target_shape[i]) // 2 for i in range(len(target_shape))]
        slices = tuple(slice(starts[i], starts[i] + target_shape[i]) for i in range(len(target_shape)))
        return full_conv[slices]


    def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False):
        """
        Batched OpenCL Poisson solver for two 2D fields.
        Optimized: scatter H2D (~4x smaller transfer), fused multiply
        (reads g_fft once), async D2H (one sync instead of two).
        When profile=True, records and optionally prints per-stage timings.
        """
        if not OPENCL_AVAILABLE:
            if verbose:
                print("OpenCL/pyvkfft not available, falling back to CPU solver")
            return self.solve(f1, verbose), self.solve(f2, verbose)

        # if self.dim != 2:
        #     return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)

        if not self.ready:
            self.prepare(verbose)

        fshape = f1.shape
        if fshape != f2.shape:
            raise ValueError("solve_opencl_batch2 requires inputs with identical shapes")

        if self.ptype != 'G' or self.stype != 'F':
            return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)

        if self.gdims != fshape:
            print(self.gdims, fshape)
            raise ValueError("Shape of function should match shape of grid")

        timing = {'cache_prepare_ms': 0.0}
        t_total = timeit.default_timer()

        t0 = timeit.default_timer()
        queue = self._get_opencl_queue()
        timing['queue_ms'] = (timeit.default_timer() - t0) * 1000.0

        t0 = timeit.default_timer()
        fshape = f1.shape
        full_shape = tuple(2 * s - 1 for s in fshape)
        timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0

        cache_valid = (
            self._opencl_cache is not None and
            self._opencl_cache['fshape'] == fshape and
            self._opencl_cache['full_shape'] == full_shape and
            'f_batch_dev' in self._opencl_cache
        )

        if not cache_valid:
            t0 = timeit.default_timer()
            g_host = numpy.asarray(self.G, dtype=numpy.float32)
            g_pad = numpy.zeros(full_shape, dtype=numpy.float32)
            g_slices = tuple(slice(0, s) for s in g_host.shape)
            g_pad[g_slices] = g_host
            g_dev = cla.to_device(queue, g_pad)
            g_fft_dev = cl_rfftn(g_dev)
            if profile:
                queue.finish()

            # For circular conv at size (2s-1) with G of shape (2s-1) and f of shape s,
            # alias-free positions are [s-1, 2s-2], matching scipy 'same' extraction.
            same_starts = [fshape[i] - 1 for i in range(self.dim)]
            fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
            fft_n_elements = 1
            for d in fft_shape:
                fft_n_elements *= d
            src_plane = fshape[0] * fshape[1]
            dst_plane = full_shape[0] * full_shape[1]

            # Pinned allocator for staging (H2D) and output (D2H) buffers
            alloc_pinned = self._get_opencl_alloc_host_ptr()

            self._opencl_cache = {
                'fshape': fshape,
                'full_shape': full_shape,
                'fft_shape': fft_shape,
                'fft_n_elements': fft_n_elements,
                'g_fft_dev': g_fft_dev,
                'same_starts': same_starts,
                'gdx_prod': numpy.prod(self.gdx),
                'src_plane': src_plane,
                'dst_plane': dst_plane,
                'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
                'host_stage': numpy.empty((2,) + fshape, dtype=numpy.float32),
                'f_stage_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32),
                'f_fft_batch_dev': cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
                'conv_batch_dev': cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
                'p_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
                'p1_host': numpy.empty(fshape, dtype=numpy.float32),
                'p2_host': numpy.empty(fshape, dtype=numpy.float32),
            }
            timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0

        c = self._opencl_cache
        fshape_c = c['fshape']
        full_shape_c = c['full_shape']

        # H2D: compact transfer of unpadded data, then scatter to padded buffer on device
        t0 = timeit.default_timer()
        H, W = fshape_c
        f1_f32 = f1 if (f1.flags['C_CONTIGUOUS'] and f1.dtype == numpy.float32) else numpy.ascontiguousarray(f1, dtype=numpy.float32)
        f2_f32 = f2 if (f2.flags['C_CONTIGUOUS'] and f2.dtype == numpy.float32) else numpy.ascontiguousarray(f2, dtype=numpy.float32)
        host_stage = c['host_stage']
        numpy.copyto(host_stage[0], f1_f32)
        numpy.copyto(host_stage[1], f2_f32)
        f_stage_dev = c['f_stage_dev']
        cl.enqueue_copy(queue, f_stage_dev.data, host_stage, is_blocking=False)
        c['f_batch_dev'].fill(0, queue=queue)
        scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
        scatter_batch(
            queue, (W, H, 2), None,
            f_stage_dev.data, c['f_batch_dev'].data,
            numpy.int32(W), numpy.int32(H),
            numpy.int32(full_shape_c[1]),
            numpy.int32(H * W),
            numpy.int32(c['dst_plane']),
        )
        if profile:
            queue.finish()
        timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0

        # Single batched 2D FFT for both fields (ndim=2 treats dim-0 as batch)
        t0 = timeit.default_timer()
        fft_batch = cl_rfftn(c['f_batch_dev'], c['f_fft_batch_dev'], ndim=2)
        if profile:
            queue.finish()
        timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0

        # Batched multiply: each batch plane * same Green's spectrum (4 elements/work item, explicit local size)
        t0 = timeit.default_timer()
        n_fft = c['fft_n_elements']
        global_n = max(1, (n_fft + 3) // 4)
        local_n = min(256, global_n)
        global_n = ((global_n + local_n - 1) // local_n) * local_n
        batch_mul = self._get_opencl_batch_mul_kernel(queue)
        batch_mul(
            queue, (2, global_n), (1, local_n),
            fft_batch.data,
            c['g_fft_dev'].data,
            numpy.int32(n_fft)
        )
        if profile:
            queue.finish()
        timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0

        # Single batched 2D iFFT for both fields
        t0 = timeit.default_timer()
        conv_batch = cl_irfftn(fft_batch, c['conv_batch_dev'], ndim=2)
        if profile:
            queue.finish()
        timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0

        # Batch crop + scale on device
        t0 = timeit.default_timer()
        crop_batch = self._get_opencl_crop_batch_kernel(queue)
        same_starts = c['same_starts']
        crop_batch(
            queue, (fshape_c[1], fshape_c[0], 2), None,
            conv_batch.data, c['p_batch_dev'].data,
            numpy.int32(full_shape_c[1]),
            numpy.int32(fshape_c[1]),
            numpy.int32(fshape_c[0]),
            numpy.int32(same_starts[1]),
            numpy.int32(same_starts[0]),
            numpy.float32(c['gdx_prod']),
            numpy.int32(c['dst_plane']),
            numpy.int32(c['src_plane'])
        )
        if profile:
            queue.finish()
        timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0

        # D2H: map pinned buffer, copy into pre-allocated host arrays (zero allocation)
        t0 = timeit.default_timer()
        p1 = c['p1_host']
        p2 = c['p2_host']
        mapped_p = c['p_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True)
        numpy.copyto(p1, mapped_p[0])
        numpy.copyto(p2, mapped_p[1])
        mapped_p.base.release(queue=queue)
        timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0

        timing['total_ms'] = (timeit.default_timer() - t_total) * 1000.0
        self._opencl_last_timing = timing

        if verbose:
            print('Poisson solver (OpenCL FFT batch2) finished')
        if profile:
            print(
                "OpenCL batch2 timing [ms]: "
                f"total={timing['total_ms']:.3f}, "
                f"queue={timing['queue_ms']:.3f}, "
                f"cache={timing['cache_prepare_ms']:.3f}, "
                f"host={timing['host_prepare_ms']:.3f}, "
                f"h2d={timing['h2d_set_ms']:.3f}, "
                f"fft={timing['fft_f_ms']:.3f}, "
                f"mul={timing['mul_ms']:.3f}, "
                f"ifft={timing['ifft_ms']:.3f}, "
                f"d2h={timing['d2h_get_ms']:.3f}, "
                f"crop={timing['crop_scale_ms']:.3f}"
            )
        return p1, p2
    