'''
Copyright (c) 2015, Harsh Bhatia (bhatia4@llnl.gov)
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this
list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

import numpy
from scipy import signal, spatial
from .timer import Timer

# Try to import CuPy for GPU acceleration
try:
    import cupy as cp
    import cupyx.scipy.signal as cp_signal
    CUPY_AVAILABLE = True
except ImportError:
    cp = None
    cp_signal = None
    CUPY_AVAILABLE = False


'''
# ------------------------------------------------------------------------------
 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 (self.dim > 3):
            raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions")

        self.ready = False
        self._gpu_cache = None  # Cache for GPU arrays
        # 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:
            init(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
            if self.stype == 'F':
                p = signal.fftconvolve(f, self.G, mode='same')
                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):
        """
        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']
                
                # Transfer result back to CPU
                p = cp.asnumpy(p_gpu)
                
            elif self.stype == 'S':
                raise SyntaxError('Spatial solver not suggested for structured grids')
        
        # ----------------------------------------------------------------------
        # 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)
        
        # ----------------------------------------------------------------------
        if verbose:
            print('Poisson solver (GPU) finished')
            gtimer.end()
        
        return p
    
    def solve_gpu_batch(self, f_list, verbose=False):
        """
        Batch GPU-accelerated Poisson solver for multiple inputs.
        Solves multiple Poisson equations in a single GPU call.
        
        Args:
            f_list: List of input fields (numpy arrays)
            verbose: Print timing information
            
        Returns:
            List of solutions (numpy arrays on CPU)
        """
        if not CUPY_AVAILABLE:
            if verbose:
                print("CuPy not available, falling back to CPU solver")
            return [self.solve(f, verbose) for f in f_list]
        
        if not self.ready:
            self.prepare(verbose)
        
        gtimer = Timer()
        if verbose:
            print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)')
        
        # Verify all inputs have same shape
        fshape = f_list[0].shape
        for f in f_list:
            if f.shape != fshape:
                raise ValueError("All input fields must have the same shape for batch processing")
        
        # ----------------------------------------------------------------------
        # regular grid
        if self.ptype == 'G':
            if 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)
                    }
                
                results = []
                
                # Process all inputs in batch
                for i, f in enumerate(f_list):
                    # Transfer input to GPU
                    f_gpu = cp.asarray(f, dtype=cp.float32)
                    
                    # Perform FFT convolution on GPU
                    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']
                    
                    # Transfer result back to CPU
                    results.append(cp.asnumpy(p_gpu))
                    
                    if verbose:
                        print(f'  - Solved {i+1}/{len(f_list)}')
                
                if verbose:
                    print('Poisson solver (GPU batch) finished')
                    gtimer.end()
                
                return results
                
            elif self.stype == 'S':
                raise SyntaxError('Spatial solver not suggested for structured grids')
        
        # ----------------------------------------------------------------------
        # 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) for f in f_list]
        
        return results
    
    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()