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

class StructuredGrid(object):

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

        '''
        kwargs:
            grid:        ndarray of grid dimensions (Y,X) or (Z,Y,X)
            spacings:    ndarray of grid spacings (dy, dx) or (dz, dy, dx)
            verbose:     verbosity level
        '''

        args = kwargs.keys()

        if ('grid' not in args) or ('spacings' not in args):
            raise SyntaxError("Dimensions and spacings of the grid are required")

        verbose = 0
        if 'verbose' in args:
            verbose = kwargs['verbose']

        self.dims = kwargs['grid']
        self.dx = kwargs['spacings']
        self.dim = len(self.dims)

        if self.dim != 2 and self.dim != 3:
            raise ValueError("StructuredGrid works for 2D and 3D only")

        if self.dim != len(self.dx):
            raise ValueError("Dimensions of spacings should match that of the grid")

        # Initialize GPU cache
        self._gpu_divcurl_cache = None
        
        if verbose > 0:
            print('     Initializing', self.dim, 'D structured grid...', end ='')
            sys.stdout.flush()
            mtimer = Timer()
            print(' Done!', end ='')
            mtimer.end()

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    def divcurl(self, vfield, verbose=False):
        # compute div and curl of vfield (vector field)
        #if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any():
        if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
            raise ValueError("Dimensions of vector field should match that of the grid")

        if verbose:
            print('     Computing div and curl...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        if self.dim == 2:

            # self.dx = (dy,dx)
            dudy, dudx = numpy.gradient(vfield[:,:,0], self.dx[0], self.dx[1])
            dvdy, dvdx = numpy.gradient(vfield[:,:,1], self.dx[0], self.dx[1])

            numpy.add(dudx, dvdy, dudx)
            numpy.subtract(dvdx, dudy, dvdx)

            if verbose:
                print(' Done!', end ='')
                mtimer.end()

            return (dudx, dvdx)

        elif self.dim == 3:

            # self.dx = (dz,dy,dx)
            dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
            dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
            dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])

            numpy.add(dudx, dvdy, dudx)
            numpy.add(dudx, dvdz, dudx)

            numpy.subtract(dwdy, dvdz, dwdy)
            numpy.subtract(dudz, dwdx, dudz)
            numpy.subtract(dvdx, dudy, dvdx)

            if verbose:
                print(' Done!', end ='')
                mtimer.end()

            return (dudx, dwdy, dudz, dvdx)

    # def divcurl_gpu(self, vfield, verbose=False):
    #     """
    #     GPU-accelerated computation of divergence and curl of vector field.
    #     Uses CuPy for faster gradient computation on GPU.
    #     """
    #     if not CUPY_AVAILABLE:
    #         if verbose:
    #             print('CuPy not available, falling back to CPU')
    #         return self.divcurl(vfield, verbose)
        
    #     if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
    #         raise ValueError("Dimensions of vector field should match that of the grid")

    #     if verbose:
    #         print('     Computing div and curl on GPU...', end ='')
    #         sys.stdout.flush()
    #         mtimer = Timer()

    #     if self.dim == 2:
    #         # Transfer to GPU
    #         vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
            
    #         # Compute gradients on GPU
    #         dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1])
    #         dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1])

    #         # Compute div and curl on GPU
    #         div_gpu = dudx + dvdy
    #         curl_gpu = dvdx - dudy
            
    #         # Transfer back to CPU
    #         div = cp.asnumpy(div_gpu)
    #         curl = cp.asnumpy(curl_gpu)

    #         if verbose:
    #             print(' Done!', end ='')
    #             mtimer.end()

    #         return (div, curl)

    #     elif self.dim == 3:
    #         # Transfer to GPU
    #         vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
            
    #         # Compute gradients on GPU
    #         dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
    #         dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
    #         dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])

    #         # Compute div and curl components on GPU
    #         div_gpu = dudx + dvdy + dwdz
    #         curlu_gpu = dwdy - dvdz
    #         curlv_gpu = dudz - dwdx
    #         curlw_gpu = dvdx - dudy
            
    #         # Transfer back to CPU
    #         div = cp.asnumpy(div_gpu)
    #         curlu = cp.asnumpy(curlu_gpu)
    #         curlv = cp.asnumpy(curlv_gpu)
    #         curlw = cp.asnumpy(curlw_gpu)

    #         if verbose:
    #             print(' Done!', end ='')
    #             mtimer.end()

    #         return (div, curlu, curlv, curlw)

    def curl3D(self, vfield, verbose=False):

        #if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any():
        if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
            raise ValueError("Dimensions of vector field should match that of the grid")

        if self.dim != 3:
            raise ValueError("curl3D works only for 2D")

        if verbose:
            print('     Computing curl...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        # self.dx = (dz,dy,dx)
        dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
        dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
        dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])

        numpy.subtract(dwdy, dvdz, dwdy)
        numpy.subtract(dudz, dwdx, dudz)
        numpy.subtract(dvdx, dudy, dvdx)

        if verbose:
            print(' Done!', end ='')
            mtimer.end()

        return (dwdy, dudz, dvdx)

    def curl3D_gpu(self, vfield, verbose=False):
        """
        GPU-accelerated computation of curl for 3D vector field.
        Uses CuPy for faster gradient computation on GPU.
        """
        if not CUPY_AVAILABLE:
            if verbose:
                print('CuPy not available, falling back to CPU')
            return self.curl3D(vfield, verbose)
        
        if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
            raise ValueError("Dimensions of vector field should match that of the grid")

        if self.dim != 3:
            raise ValueError("curl3D works only for 3D")

        if verbose:
            print('     Computing curl on GPU...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        # Transfer to GPU
        vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
        
        # Compute gradients on GPU
        dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
        dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
        dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])

        # Compute curl components on GPU
        curl_u = dwdy - dvdz
        curl_v = dudz - dwdx
        curl_w = dvdx - dudy
        
        # Transfer back to CPU
        curlu = cp.asnumpy(curl_u)
        curlv = cp.asnumpy(curl_v)
        curlw = cp.asnumpy(curl_w)

        if verbose:
            print(' Done!', end ='')
            mtimer.end()

        return (curlu, curlv, curlw)

    def rotated_gradient(self, sfield, verbose=False):

        if (sfield.shape != self.dims):
	    #if (sfield.shape - self.dims).any():
            raise ValueError("Dimensions of scalar field should match that of the grid")

        if self.dim != 2:
            raise ValueError("rotated_gradient works only for 2D")

        if verbose:
            print('     Computing rotated gradient...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
        ddy *= -1.0

        grad = numpy.stack((ddy, ddx), axis=-1)

        if verbose:
            print(' Done!', end ='')
            mtimer.end()

        return grad

    def rotated_gradient_gpu(self, sfield, verbose=False):
        """
        GPU-accelerated computation of rotated gradient for 2D scalar field.
        Uses CuPy for faster gradient computation on GPU.
        """
        if not CUPY_AVAILABLE:
            if verbose:
                print('CuPy not available, falling back to CPU')
            return self.rotated_gradient(sfield, verbose)
        
        if (sfield.shape != self.dims):
            raise ValueError("Dimensions of scalar field should match that of the grid")

        if self.dim != 2:
            raise ValueError("rotated_gradient works only for 2D")

        if verbose:
            print('     Computing rotated gradient on GPU...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        # Transfer to GPU
        sfield_gpu = cp.asarray(sfield, dtype=cp.float32)
        
        # Compute gradients on GPU
        ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1])
        ddy *= -1.0

        # Stack and transfer back to CPU
        grad_gpu = cp.stack((ddy, ddx), axis=-1)
        grad = cp.asnumpy(grad_gpu)

        if verbose:
            print(' Done!', end ='')
            mtimer.end()

        return grad

    def gradient(self, sfield, verbose=False):

        if (sfield.shape != self.dims):
	    #if (sfield.shape - self.dims).any():
            raise ValueError("Dimensions of scalar field should match that of the grid")

        if verbose:
            print('     Computing gradient...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        if self.dim == 2:

            # self.dx = (dy,dx)
            ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
            grad = numpy.stack((ddx, ddy), axis = -1)

        elif self.dim == 3:

            # self.dx = (dz,dy,dx)
            ddz, ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1], self.dx[2])
            grad = numpy.stack((ddx, ddy, ddz), axis = -1)

        if verbose:
            print(' Done!', end ='')
            mtimer.end()

        return grad

    def gradient_gpu(self, sfield, verbose=False):
        """
        GPU-accelerated computation of gradient for scalar field.
        Uses CuPy for faster gradient computation on GPU.
        """
        if not CUPY_AVAILABLE:
            if verbose:
                print('CuPy not available, falling back to CPU')
            return self.gradient(sfield, verbose)
        
        if (sfield.shape != self.dims):
            raise ValueError("Dimensions of scalar field should match that of the grid")

        if verbose:
            print('     Computing gradient on GPU...', end ='')
            sys.stdout.flush()
            mtimer = Timer()

        # Transfer to GPU
        sfield_gpu = cp.asarray(sfield, dtype=cp.float32)

        if self.dim == 2:
            # Compute gradients on GPU
            ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1])
            grad_gpu = cp.stack((ddx, ddy), axis=-1)

        elif self.dim == 3:
            # Compute gradients on GPU
            ddz, ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1], self.dx[2])
            grad_gpu = cp.stack((ddx, ddy, ddz), axis=-1)

        # Transfer back to CPU
        grad = cp.asnumpy(grad_gpu)

        if verbose:
            print(' Done!', end ='')
            mtimer.end()

        return grad

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

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
