import numpy

from ..backend_availability import (
    CUPY_AVAILABLE,
    OPENCL_AVAILABLE,
    KERNEL_AVAILABLE,
    cp,
    _torque_kernel,
)
from .grid import StructuredGrid
from .poisson import PoissonSolver


class nHHD(object):

    # constructor
    def __init__(self, cuda=True, **kwargs):
        # Select Poisson backend: 'auto' (default), 'cupy', 'opencl', or 'cpu'
        backend = kwargs.get('poisson_backend', 'auto').lower()
        if backend not in ('auto', 'cupy', 'opencl', 'cpu'):
            raise ValueError("poisson_backend must be one of: 'auto', 'cupy', 'opencl', 'cpu'")

        # Resolve 'auto' to a specific backend
        if backend == 'auto':
            backend = 'cupy' if CUPY_AVAILABLE else ('opencl' if OPENCL_AVAILABLE else 'cpu')

        # Validate backend availability and fallback to 'cpu' if needed
        if backend == 'cupy' and not CUPY_AVAILABLE:
            backend = 'cpu'
        elif backend == 'opencl' and not OPENCL_AVAILABLE:
            backend = 'cpu'

        self.debug = kwargs.get('debug', False)
        self.poisson_backend = backend

        # Initialize CUDA streams for cupy backend
        if self.poisson_backend == 'cupy':
            self.stream1 = cp.cuda.Stream()
            self.stream2 = cp.cuda.Stream()

        self.cuda = cuda

        args = kwargs.keys()

        if ('grid' in args) == ('points' in args):
            raise SyntaxError("nHHD 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("nHHD needs spacings for regular grid")

        if 'grid' in args:
            self.mesh = StructuredGrid(grid=kwargs['grid'], spacings=kwargs['spacings'])
            self.psolver = PoissonSolver(grid=self.mesh.dims, spacings=self.mesh.dx, poisson_backend=self.poisson_backend)

        self.dim = self.mesh.dim
        if (self.dim != 2) and (self.dim != 3):
            raise ValueError("nHHD Solver works only for 2 and 3 dimensions")

        self._force_cpu_cache = None
        self.psolver.prepare(False)

    def get_force(self, vfield, verbose=False):
        if self.poisson_backend == 'cupy':
            return self.get_force_cuda(vfield, verbose)
        else:
            return self.get_force_cpu(vfield, verbose)

    def get_shear(self):
        if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
            return self.h
        elif self.poisson_backend == 'cupy':
            # Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
            if hasattr(self, 'h') and self.h is not None:
                return self.h
            return cp.asnumpy(self._h_gpu)

    def get_normal(self):
        if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
            return self.d
        elif self.poisson_backend == 'cupy':
            # Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
            if hasattr(self, 'd') and self.d is not None:
                return self.d
            return cp.asnumpy(self._d_gpu)

    def _solve_poisson_backend(self, field, verbose=False):
        """
        Dispatch Poisson solve using the selected backend.
        """
        if self.poisson_backend == 'opencl':
            return self.psolver.solve_opencl(field, verbose, profile=self.debug)

        if self.poisson_backend == 'cupy':
            return self.psolver.solve_gpu(field, verbose)

        if self.poisson_backend == 'cpu':
            return self.psolver.solve(field, verbose)

    # create the 3-component decomposition
    def decompose(self, vfield, verbose=0):
        # decompose for cpu and opencl backend

        if vfield.shape[-1] != self.dim:
            raise ValueError("nHHD.decompose requires a valid-dimensional vector field")

        if verbose > 0:
            print('\n ==> nhhd.decompose', vfield.shape)

        if self.dim == 2:

            # compute div curl
            (self.div, self.curlw) = self.mesh.divcurl(vfield)

            if self.poisson_backend == 'opencl':
                self.div = numpy.ascontiguousarray(self.div, dtype=numpy.float32)
                self.curlw = numpy.ascontiguousarray(self.curlw, dtype=numpy.float32)

            if verbose > 0:
                print('div =', self.div.shape, self.div.min(), self.div.max())
                print('curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())

            # compute potentials
            if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
                self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=self.debug)
            else:
                self.nD = self._solve_poisson_backend(self.div, verbose > 0)
                self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
            if verbose > 0:
                print('nD =', self.nD.shape, self.nD.min(), self.nD.max())
                print('nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())

            # compute fields as gradients of potentials
            self.d = self.mesh.gradient(self.nD)
            if verbose > 0:
                print('d =', self.d.shape, self.d.min(), self.d.max())

            self.r = self.mesh.rotated_gradient(self.nRu)
            if verbose > 0:
                print('r =', self.r.shape, self.r.min(), self.r.max())
        else:

            # compute div curl
            (self.div, self.curlu, self.curlv, self.curlw) = self.mesh.divcurl(vfield)

            if verbose > 0:
                print('div =', self.div.shape, self.div.min(), self.div.max())
                print('curlu =', self.curlu.shape, self.curlu.min(), self.curlu.max())
                print('curlv =', self.curlv.shape, self.curlv.min(), self.curlv.max())
                print('curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())

            # compute potentials
            self.nD = self._solve_poisson_backend(self.div, verbose > 0)
            if verbose > 0:
                print('nD =', self.nD.shape, self.nD.min(), self.nD.max())

            self.nRu = self._solve_poisson_backend(self.curlu, verbose > 0)
            self.nRu *= -1.0
            if verbose > 0:
                print('nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max())

            self.nRv = self._solve_poisson_backend(self.curlv, verbose > 0)
            self.nRv *= -1.0
            if verbose > 0:
                print('nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max())

            self.nRw = self._solve_poisson_backend(self.curlw, verbose > 0)
            self.nRw *= -1.0

            if verbose > 0:
                print('nRw =', self.nRw.shape, self.nRw.min(), self.nRw.max())

            # compute divergent field as gradient of D potential
            self.d = self.mesh.gradient(self.nD)
            if verbose > 0:
                print('d =', self.d.shape, self.d.min(), self.d.max())

            # temporary use of r
            self.r = numpy.stack((self.nRu, self.nRv, self.nRw), axis=-1)

            # compute rotational field as curl of R potential
            (cu, cv, cw) = self.mesh.curl3D(self.r)
            self.r = numpy.stack((cu, cv, cw), axis=-1)
            if verbose > 0:
                print('r =', self.r.shape, self.r.min(), self.r.max())

        self.h = numpy.add(self.d, self.r)
        numpy.subtract(vfield, self.h, self.h)

        if verbose > 0:
            print('h =', self.h.shape, self.h.min(), self.h.max())

    def decompose_cuda(self, vfield, verbose=0):
        """
        GPU-optimized decomposition that keeps intermediate results on GPU.
        Returns GPU arrays to avoid CPU-GPU transfers.

        Args:
            vfield: Vector field (numpy array or CuPy array)
            verbose: Verbosity level (default: 0)

        Returns:
            Tuple of GPU arrays: (d_gpu, h_gpu, r_gpu, nRu_gpu) for 2D fields
        """
        if not CUPY_AVAILABLE:
            if verbose > 0:
                print("CuPy not available, falling back to CPU decompose")
            self.decompose(vfield, verbose)
            return None

        if vfield.shape[-1] != self.dim:
            raise ValueError("nHHD.decompose requires a valid-dimensional vector field")

        if verbose > 0:
            print('\n ==> nhhd.decompose_cuda', vfield.shape)

        # Convert input to GPU if needed
        if isinstance(vfield, numpy.ndarray):
            vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
        else:
            vfield_gpu = vfield

        if self.dim == 2:
            # Compute div curl on GPU
            dx0, dx1 = self.mesh.dx[0], self.mesh.dx[1]

            dudy, dudx = cp.gradient(vfield_gpu[:, :, 0], dx0, dx1)
            dvdy, dvdx = cp.gradient(vfield_gpu[:, :, 1], dx0, dx1)

            div_gpu = cp.add(dudx, dvdy)
            curlw_gpu = cp.subtract(dvdx, dudy)

            if verbose > 0:
                print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max()))
                print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))

            # Use a single stream for both solves. Running concurrent FFT-based
            # solves on separate streams can cause occasional long-run jitter.
            nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
            nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
            cp.cuda.get_current_stream().synchronize()

            if verbose > 0:
                print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
                print('nR =', nRu_gpu.shape, float(nRu_gpu.min()), float(nRu_gpu.max()))

            # Compute gradients on GPU
            ddy, ddx = cp.gradient(nD_gpu, dx0, dx1)
            d_gpu = cp.stack((ddx, ddy), axis=-1)

            if verbose > 0:
                print('d =', d_gpu.shape, float(d_gpu.min()), float(d_gpu.max()))

            # Compute rotated gradient on GPU
            rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1)
            rdy = cp.multiply(rdy, -1.0)
            r_gpu = cp.stack((rdy, rdx), axis=-1)

            if verbose > 0:
                print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max()))

            # Compute h on GPU
            h_gpu = cp.add(d_gpu, r_gpu)
            h_gpu = cp.subtract(vfield_gpu, h_gpu)

            if verbose > 0:
                print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))

            # Ensure all default-stream work (gradient, etc.) is complete before use
            cp.cuda.get_current_stream().synchronize()

            # Store GPU arrays
            self._d_gpu = d_gpu
            self._h_gpu = h_gpu
            self._r_gpu = r_gpu
            self._nRu_gpu = nRu_gpu

            return (d_gpu, h_gpu, r_gpu, nRu_gpu)

        else:
            # For 3D, fall back to CPU for now
            if verbose > 0:
                print("3D GPU decompose not yet implemented, using CPU")
            self.decompose(vfield, verbose)
            return None

    def get_force_cpu(self, vfield, verbose=False):
        shape = (vfield.shape[0], vfield.shape[1])
        if (not hasattr(self, '_force_cpu_cache') or
                self._force_cpu_cache is None or
                self._force_cpu_cache['shape'] != shape):
            self._init_force_cpu_cache(shape)
        grid = self._force_cpu_cache['grid']

        self.decompose(vfield, verbose)

        # Optimized: use abs instead of absolute (slightly faster)
        self.Sn = numpy.sum(numpy.abs(self.d))

        # Optimized: direct indexing and sum (already efficient)
        self.Stx = numpy.sum(self.h[:, :, 0])
        self.Sty = numpy.sum(self.h[:, :, 1])

        # Optimized: find both max and min indices in single pass through flattened array
        nRu_flat = self.nRu.ravel()
        nR_max_idx = numpy.argmax(nRu_flat)
        nR_min_idx = numpy.argmin(nRu_flat)
        # Convert flat indices to 2D coordinates (faster than unravel_index)
        nR_max_loc = (nR_max_idx // shape[1], nR_max_idx % shape[1])
        nR_min_loc = (nR_min_idx // shape[1], nR_min_idx % shape[1])

        # Optimized: compute cross products directly without creating intermediate arrays
        # For 2D: cross(a, b) = a[0]*b[1] - a[1]*b[0]
        # Extract components once for reuse
        r_x = self.r[:, :, 0]
        r_y = self.r[:, :, 1]
        grid_x = grid[:, :, 0]
        grid_y = grid[:, :, 1]

        # Compute cross products directly (faster than numpy.cross for 2D)
        # cross(r_p, r) where r_p = grid - nR_max_loc
        max_y, max_x = nR_max_loc[0], nR_max_loc[1]
        cross_p = (grid_x - max_x) * r_y - (grid_y - max_y) * r_x

        # cross(r_n, r) where r_n = grid - nR_min_loc
        min_y, min_x = nR_min_loc[0], nR_min_loc[1]
        cross_n = (grid_x - min_x) * r_y - (grid_y - min_y) * r_x

        self.Stor = numpy.sum(cross_p) + numpy.sum(cross_n)

        # Optimized: pre-allocate output array with explicit dtype
        output = numpy.array([self.Sn, self.Stx, self.Sty, self.Stor], dtype=numpy.float32)

        # Optimized: combine threshold checks (short-circuit evaluation)
        flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)

        return flag, output

    def _init_force_cpu_cache(self, shape):
        """
        Initialize cached CPU grid arrays for force computation.
        Avoids per-frame mgrid/stack allocations and reduces jitter.
        """
        mgrid = numpy.mgrid[0:shape[0], 0:shape[1]]
        self._force_cpu_cache = {
            'shape': shape,
            'grid': numpy.stack((mgrid[0], mgrid[1]), axis=2),
        }

    def _init_force_gpu_cache(self, shape):
        """
        Initialize cached GPU arrays for force computation.
        Call once when shape is known to avoid repeated allocations.

        Args:
            shape: Tuple (height, width) of the vector field
        """
        if not CUPY_AVAILABLE:
            return

        # Cache grid arrays (only need to create once per shape)
        self._force_gpu_cache = {
            'shape': shape,
            'mgrid_0': cp.arange(shape[0], dtype=cp.float32)[:, None],
            'mgrid_1': cp.arange(shape[1], dtype=cp.float32)[None, :],
        }
        # Pre-broadcast for efficiency
        self._force_gpu_cache['mgrid_0_bc'] = cp.broadcast_to(
            self._force_gpu_cache['mgrid_0'], shape
        ).copy()  # .copy() to make contiguous
        self._force_gpu_cache['mgrid_1_bc'] = cp.broadcast_to(
            self._force_gpu_cache['mgrid_1'], shape
        ).copy()

    def get_force_cuda(self, vfield, verbose=False):
        """
        Further optimized GPU force computation using CuPy ElementwiseKernel.
        Fuses all operations into a single kernel for maximum performance.

        Args:
            vfield: Vector field (numpy array)
            verbose: Print debug information (default: False)

        Returns:
            tuple: (flag, output) where:
                - flag: Boolean indicating if force thresholds are met
                - output: Array [Sn, Stx, Sty, Stor] containing force components
        """
        if not CUPY_AVAILABLE:
            if verbose:
                print("CuPy not available, falling back to CPU get_force")
            return self.get_force(vfield, verbose)

        # Decompose the vector field
        self.decompose_cuda(vfield, verbose)

        d_gpu = self._d_gpu
        h_gpu = self._h_gpu
        r_gpu = self._r_gpu
        nRu_gpu = self._nRu_gpu

        # Use reduction kernels for better performance
        # Compute Sn as L1 sum directly on GPU (faster than linalg.norm for this case)
        Sn = float(cp.sum(cp.abs(d_gpu)))

        # Compute Stx, Sty
        Stx = float(cp.sum(h_gpu[:, :, 0]))
        Sty = float(cp.sum(h_gpu[:, :, 1]))

        # Find extrema indices
        nR_max_idx = int(cp.argmax(nRu_gpu))
        nR_min_idx = int(cp.argmin(nRu_gpu))

        shape = nRu_gpu.shape
        max_0, max_1 = nR_max_idx // shape[1], nR_max_idx % shape[1]
        min_0, min_1 = nR_min_idx // shape[1], nR_min_idx % shape[1]

        # Initialize cache if needed
        if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape:
            self._init_force_gpu_cache(shape)

        mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
        mgrid_1 = self._force_gpu_cache['mgrid_1_bc']

        r_x = r_gpu[:, :, 0]
        r_y = r_gpu[:, :, 1]

        if KERNEL_AVAILABLE and _torque_kernel is not None:
            torque = _torque_kernel(
                r_x, r_y, mgrid_0, mgrid_1,
                float(max_0), float(max_1), float(min_0), float(min_1)
            )
            Stor = float(cp.sum(torque))
        else:
            Stor = float(cp.sum(
                (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
                (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
            ))

        self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
        output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)

        # Keep stable CPU copies for get_normal()/get_shear() so callers get this frame's
        # result instead of reading GPU memory that can change (avoids CuPy-only jumps)
        self.d = cp.asnumpy(self._d_gpu)
        self.h = cp.asnumpy(self._h_gpu)

        flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000)

        return flag, output

    def clear_force_gpu_cache(self):
        """Clear the GPU cache for force computation to free memory."""
        if hasattr(self, '_force_gpu_cache'):
            self._force_gpu_cache = None
            if CUPY_AVAILABLE:
                cp.get_default_memory_pool().free_all_blocks()
