# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2016-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import warnings
import numpy as np
from ..processing_unit import has_cuda

if has_cuda:
    from ..processing_unit.kernel_source import get_kernel_source as getks
    from ..processing_unit.cu_processing_unit import CUProcessingUnit
    import pycuda.gpuarray as cua
    from pycuda.reduction import ReductionKernel as CU_RedK
    from pycuda.elementwise import ElementwiseKernel as CU_ElK
    import pycuda.tools as cu_tools
    from pycuda.compiler import SourceModule

    argmax_dtype = np.dtype([("idx", np.int32), ("cur_max", np.float32)])
    cu_tools.get_or_register_dtype("idx_max", argmax_dtype)

    from pyvkfft.fft import rfftn as pyvkfft_rfftn, irfftn as pyvkfft_irfftn


# if has_opencl:  # TODO: OpenCL registration


def register_translation_cuda(ref_img, img, upsampling=10, processing_unit=None, overwrite=False, blocksize=None):
    """
    CUDA image registration, with or without subpixel precision. Sub-pixel
    implementation is currently very slow (slower than scikit-image).

    :param ref_img, img: the images to be registered, either as numpy array or as
        a pycuda.gpuarray.GPUArray.
    :param upsampling: the upsampling factor (integer >=1), for subpixel registration
    :param processing_unit: the processing unit to be used for the calculation.
        Should already be initialised. If None, the default one will be used.
    :param overwrite: if True and the input images are pycuda GPUArrays, they will be
        overwritten.
    :param blocksize: the CUDA blocksize for the subpixel registration.
        If None, will be automatically set to 32. Larger values
        do not seem to bring improvements. Must be smaller than upsampling**2
    :return: the computed shift as a tuple (dy, dx)
    """
    warnings.warn("register_translation_cuda() is deprecated. Use "
                  "register_translation_2d_paraboloid_cuda() instead.", DeprecationWarning)
    pu = processing_unit
    if pu is None:
        pu = CUProcessingUnit()
        pu.init_cuda(test_fft=False, verbose=False)

    if blocksize is None:
        blocksize = 32
    if blocksize > upsampling ** 2:
        blocksize = upsampling ** 2

    if not isinstance(ref_img, cua.GPUArray):
        ref_img = cua.to_gpu(ref_img.astype(np.complex64))
    else:
        ref_img = ref_img.astype(np.complex64)

    if not isinstance(img, cua.GPUArray):
        img = cua.to_gpu(img.astype(np.complex64))
    else:
        img = img.astype(np.complex64)

    if register_translation_cuda.cu_argmax_c_red is None:
        register_translation_cuda.cu_argmax_c_red = \
            CU_RedK(argmax_dtype, neutral="idx_max(0,0.0f)", name='argmax_c',
                    reduce_expr="argmax_reduce(a,b)",
                    map_expr="idx_max(i, abs(d[i]))",
                    preamble=getks("cuda/argmax.cu"),
                    options=["-use_fast_math"],
                    arguments="pycuda::complex<float> *d")

    if register_translation_cuda.cu_argmax_f_red is None:
        register_translation_cuda.cu_argmax_f_red = \
            CU_RedK(argmax_dtype, neutral="idx_max(0,0.0f)", name='argmax_f',
                    reduce_expr="argmax_reduce(a,b)",
                    map_expr="idx_max(i, d[i])",
                    preamble=getks("cuda/argmax.cu"),
                    options=["-use_fast_math"],
                    arguments="float *d")

    if overwrite is False:
        d1f, d2f = cua.empty_like(ref_img), cua.empty_like(ref_img)
        # out-of-place FFT of the two images
        pu.fft(ref_img, d1f, norm=True)
        pu.fft(img, d2f, norm=True)
    else:
        d1f, d2f = ref_img, img
        pu.fft(d1f, d1f, norm=True)
        pu.fft(d2f, d2f, norm=True)

    # Pixel registration
    cc0 = d1f * d2f.conj()
    pu.ifft(cc0, d1f, norm=False)
    idx = register_translation_cuda.cu_argmax_c_red(d1f).get()["idx"]

    ny, nx = d1f.shape
    shift0 = [idx // nx, idx % nx]
    shift0[0] -= (shift0[0] > ny / 2) * ny
    shift0[1] -= (shift0[1] > nx / 2) * nx
    if upsampling == 1:
        return shift0

    # Sub-pixel registration
    if blocksize not in register_translation_cuda.cu_cc_zoom:
        reg_mod = SourceModule(getks("cuda/argmax.cu") +
                               getks("utils/cuda/registration.cu") % {"blocksize": blocksize},
                               options=["-use_fast_math"])
        register_translation_cuda.cu_cc_zoom[blocksize] = reg_mod.get_function("cc_zoom")
    cu_cc_zoom = register_translation_cuda.cu_cc_zoom[blocksize]

    upsample_range = 1.5

    nxy1 = np.int32(upsampling * upsample_range)
    y0 = np.float32(shift0[0] - (nxy1 // 2) / upsampling)
    x0 = np.float32(shift0[1] - (nxy1 // 2) / upsampling)
    z0 = np.float32(0)
    dxy = np.float32(1 / upsampling)
    nyu, nxu = np.int32(d1f.shape[-2]), np.int32(d1f.shape[-1])
    nzu = np.int32(1)
    cc1_cu = cua.empty((nxy1, nxy1), dtype=np.float32)

    cu_cc_zoom(cc1_cu, cc0, x0, y0, z0, dxy, dxy, dxy, nxy1, nxy1, nxu, nyu, nzu,
               block=(int(blocksize), 1, 1), grid=(int(nxy1), int(nxy1), 1))
    idx1 = register_translation_cuda.cu_argmax_f_red(cc1_cu).get()["idx"]

    ix1 = idx1 % nxy1
    iy1 = idx1 // nxy1
    x1 = x0 + dxy * ix1
    y1 = y0 + dxy * iy1

    return y1, x1


register_translation_cuda.cu_argmax_c_red = None
register_translation_cuda.cu_argmax_f_red = None
register_translation_cuda.cu_cc_zoom = {}


def register_translation_cuda_n(ref_img, img, upsampling=10, processing_unit=None, overwrite=False,
                                blocksize=None):
    """
    CUDA image registration, with or without subpixel precision. Sub-pixel
    implementation is currently very slow (slower than scikit-image).
    This version allows to register a stack of N*M images against a stack of N reference images.
    Performance improvement compared with looping individual registration is
    mostly obtained for subpixel accuracy, with ideally N*M a multiple of the
    number of GPU multi-processors.

    :param ref_img, img: the stack of images to be registered, either as numpy arrays or as
        a pycuda.gpuarray.GPUArray.
    :param upsampling: the upsampling factor (integer >=1), for subpixel registration
    :param processing_unit: the processing unit to be used for the calculation.
        Should already be initialised. If None, the default one will be used.
    :param overwrite: if True and the input images are pycuda GPUArrays, they will be
        overwritten.
    :param blocksize: the CUDA blocksize for the subpixel registration.
        If None, will be automatically set to min(upsampling**2, 64). Larger values
        do not seem to bring improvements.
        Check that register_translation_cuda_n.cu_cc_zoomN[blocksize].shared_size_bytes
        does not exceed 48kb.
    :return: the computed shift as a tuple of arrays (dy, dx)
    """
    pu = processing_unit
    if pu is None:
        pu = CUProcessingUnit()
        pu.init_cuda(test_fft=False, verbose=False)

    if blocksize is None:
        blocksize = min(upsampling ** 2, 512)

    if not isinstance(ref_img, cua.GPUArray):
        ref_img = cua.to_gpu(ref_img.astype(np.complex64))
    else:
        ref_img = ref_img.astype(np.complex64)

    if not isinstance(img, cua.GPUArray):
        img = cua.to_gpu(img.astype(np.complex64))
    else:
        img = img.astype(np.complex64)

    if img.ndim == 2:
        img = img.reshape((1, img.shape[0], img.shape[1]))
    if ref_img.ndim == 2:
        ref_img = ref_img.reshape((1, ref_img.shape[0], ref_img.shape[1]))

    if register_translation_cuda.cu_argmax_c_red is None:
        register_translation_cuda.cu_argmax_c_red = \
            CU_RedK(argmax_dtype, neutral="idx_max(0,0.0f)", name='argmax_c',
                    reduce_expr="argmax_reduce(a,b)",
                    map_expr="idx_max(i, abs(d[i]))",
                    preamble=getks("cuda/argmax.cu"),
                    options=["-use_fast_math"],
                    arguments="pycuda::complex<float> *d")

    if overwrite is False:
        d1f, d2f = cua.empty_like(ref_img), cua.empty_like(img)
        # out-of-place FFT of the two images
        pu.fft(ref_img, d1f, ndim=2, norm=True)
        pu.fft(img, d2f, ndim=2, norm=True)
    else:
        d1f, d2f = ref_img, img
        pu.fft(d1f, d1f, ndim=2, norm=True)
        pu.fft(d2f, d2f, ndim=2, norm=True)

    # Pixel registration
    nzref, ny, nx = ref_img.shape
    nz, ny, nx = img.shape
    dy, dx = np.empty(nz, dtype=np.float32), np.empty(nz, dtype=np.float32)
    cutmp = cua.empty_like(d2f[0])
    # This may go faster by treating all images simultaneously ?
    for i in range(nz):
        i1 = i * nzref // nz
        d2f[i] = d1f[i1] * d2f[i].conj()
        pu.ifft(d2f[i], cutmp, ndim=2, norm=False)
        idx = register_translation_cuda.cu_argmax_c_red(cutmp).get()["idx"]
        dy0, dx0 = idx // nx, idx % nx
        dy[i] = dy0 - (dy0 > ny / 2) * ny
        dx[i] = dx0 - (dx0 > nx / 2) * nx
    if upsampling == 1:
        return dy, dx

    # Sub-pixel registration
    if blocksize not in register_translation_cuda_n.cu_cc_zoomN:
        reg_mod = SourceModule(getks("cuda/argmax.cu") +
                               getks("utils/cuda/registration.cu") % {"blocksize": blocksize},
                               options=["-use_fast_math"])
        register_translation_cuda_n.cu_cc_zoomN[blocksize] = reg_mod.get_function("cc_zoomN")
    cu_cc_zoom_n = register_translation_cuda_n.cu_cc_zoomN[blocksize]

    upsample_range = 1.5

    nxy1 = np.int32(upsampling * upsample_range)
    y0 = np.float32(dy - (nxy1 // 2) / upsampling)
    x0 = np.float32(dx - (nxy1 // 2) / upsampling)
    dxy = np.float32(1 / upsampling)
    nyu, nxu = np.int32(ny), np.int32(nx)
    x0_cu = cua.to_gpu(x0.astype(np.float32))
    y0_cu = cua.to_gpu(y0.astype(np.float32))
    cu_cc_zoom_n(d2f, x0_cu, y0_cu, dxy, dxy, nxy1, nxy1, nxu, nyu,
                 block=(int(blocksize), 1, 1), grid=(nz, 1, 1))

    return y0_cu.get(), x0_cu.get()


register_translation_cuda_n.cu_cc_zoomN = {}


def register_translation_2d_paraboloid_cuda(ref_img, img, low_cutoff=None, high_cutoff=None,
                                            low_width=0.03, high_width=0.03,
                                            return_cc=False, return_gpu_arrays=False):
    """
    CUDA image registration. Sub-pixel accuracy is provided by a paraboloid
    fit of the CC maximum.
    :param ref_img, img: the images to be registered, either as numpy array or as
        a pycuda.gpuarray.GPUArray. Type should be float32.
        These can also be a stack of 2D images of size - in that case
        the registration is done in // for the all images, and the shift are returned
        as arrays which have the shape ref_img.shape[:-2].
    :param low_cutoff: a 0<value<<0.5 can be given (typically it should be a few 0.01),
        an erfc filter with a cutoff at low_cutoff*N (where N is the size along each dimension)
        will be applied, after the images have been FT'd.
    :param high_cutoff: same as low_cutoff fot the high frequency filter, should be close below 0.5.
        This is less useful than the low_cutoff, as the high frequencis are the most useful for
        the registration.
    :param low_width: the width of the low cutoff filter, also as a percentage of the size. Default=0.03
    :param high_width: same as low_width
    :param return_cc: if True, also return the CC map (default=False)
    :param return_gpu_arrays: if True, return the results as gpu arrays (default=False)
    :return: the computed shift as a tuple (dy, dx)
    """
    if not isinstance(ref_img, cua.GPUArray):
        ref_img = cua.to_gpu(ref_img.astype(np.float32))
    else:
        ref_img = ref_img.astype(np.float32)

    if not isinstance(img, cua.GPUArray):
        img = cua.to_gpu(img.astype(np.float32))
    else:
        img = img.astype(np.float32)

    if register_translation_2d_paraboloid_cuda.cu_argmax_f_red is None:
        register_translation_2d_paraboloid_cuda.cu_argmax_f_red = \
            CU_RedK(argmax_dtype, neutral="idx_max(0,0.0f)", name='argmax_f',
                    reduce_expr="argmax_reduce(a,b)",
                    map_expr="idx_max(i, d[i])",
                    preamble=getks("cuda/argmax.cu"),
                    options=["-use_fast_math"],
                    arguments="float *d")

    if register_translation_2d_paraboloid_cuda.cu_paraboloid9 is None:
        cu_paraboloid_src = """
        // parabolic 3-point fit along x and y, computed for a stack of CC maps
        // result is stored in cy and cx
        __device__ void paraboloid9fit(const int i, float* cy, float* cx, idx_max *im, float *cc, const int ny, const int nx)
        {
            const int idx = im[i].idx;
            const int iy = idx / nx;
            const int ix = idx - iy*nx;
            const int ixm = ix-1 + nx *(ix==0);
            const int ixp = ix+1 - nx *(ix==(nx-1));
            const int iym = iy-1 + ny *(iy==0);
            const int iyp = iy+1 - ny *(iy==(ny-1));

            // Paraboloid fit
            // CC(x,y)=Ax^2 + By^2 + Cxy + Dx + Ey + F

            const float F= cc[ix  + iy  * nx];
            const float A= (cc[ixp  + iy  * nx] + cc[ixm  + iy  * nx])/2.0f - F;
            const float B= (cc[ix   + iyp * nx] + cc[ix   + iym * nx])/2.0f - F;
            const float D= (cc[ixp  + iy  * nx] - cc[ixm  + iy  * nx])/2.0f;
            const float E= (cc[ix   + iyp * nx] - cc[ix   + iym * nx])/2.0f;
            const float C= (cc[ixp  + iyp * nx] - cc[ixp  + iym * nx] - cc[ixm  + iyp * nx] + cc[ixm  + iym * nx])/4.0f;

            const float dx = (float)ix + fmaxf(fminf( (2*B*D-C*E) / (C*C-4*A*B), 0.5f),-0.5f);
            const float dy = (float)iy + fmaxf(fminf( (2*A*E-C*D) / (C*C-4*A*B), 0.5f),-0.5f);
            cx[i] = dx - nx * (dx>(nx/2)) + nx * (dx<(-nx/2));
            cy[i] = dy - ny * (dy>(ny/2)) + ny * (dy<(-ny/2));
        }

        """
        register_translation_2d_paraboloid_cuda.cu_paraboloid9 = \
            CU_ElK(name='cu_paraboloid9',
                   operation="paraboloid9fit(i, vy, vx, im, cc, ny, nx)",
                   preamble=getks('cuda/complex.cu') + getks("cuda/argmax.cu") + cu_paraboloid_src,
                   options=["-use_fast_math"],
                   arguments="float* vy, float* vx, idx_max * im, float *cc, const int ny, const int nx")

    if register_translation_2d_paraboloid_cuda.cu_cc_fourier_conj_filter is None:
        cu_cc_fourier_conj_filter_src = """
        // Multiply two array by complex conjugate of other & filter
        // This is a half-hermitian array
        // nx should be the length of the complex array
        __device__ void cc_fourier_conj_filter(const int i, complexf* d1, complexf* d2, const float low_cutoff, 
                                      const float low_width, const float high_cutoff, const float high_width,
                                      const int ny, const int nx)
        {
            const int iz = i / (nx*ny);
            const int iy = (i-iz*nx*ny)/nx;
            const int ix = i-nx*(iy + ny*iz);

            const complexf c1 = d1[i];
            const complexf c2 = d2[i];

            float r = 1;
            const float fx = (float)ix/(float)(2*(nx-1));
            const float fy = (float)(iy - ny * (iy>=ny/2)) / (float)(ny);
            const float f = sqrtf(fx*fx + fy*fy);
            if(low_cutoff>=0.0f) r *=1.0f - 0.5f * erfcf((f - low_cutoff) / low_width);
            if(high_cutoff<=0.5f) r *= 0.5f * erfcf((f - high_cutoff) / high_width);
            d1[i] = complexf(c1.real()*c2.real()+c1.imag()*c2.imag(), c1.imag()*c2.real()-c1.real()*c2.imag()) * r;
        }

        """
        register_translation_2d_paraboloid_cuda.cu_cc_fourier_conj_filter = \
            CU_ElK(name='cc_fourier_conj_filter',
                   operation="cc_fourier_conj_filter(i, d1, d2, low_cutoff, low_width, high_cutoff, high_width, ny, nx)",
                   preamble=getks('cuda/complex.cu') + cu_cc_fourier_conj_filter_src,
                   options=["-use_fast_math"],
                   arguments="pycuda::complex<float>* d1, pycuda::complex<float>* d2,"
                             "const float low_cutoff, const float low_width,"
                             "const float high_cutoff, const float high_width, const int ny, const int nx")
    # out-of-place r2c
    d1f = pyvkfft_rfftn(ref_img, ndim=2)
    d2f = pyvkfft_rfftn(img, ndim=2)

    # Filter and d1f * d2f.conj()
    # cc0 = irfftn(d1f * d2f.conj(out=d2f), ndim=2)
    ny, nx = np.int32(img.shape[0]), np.int32(img.shape[1])
    low_cutoff = np.float32(-1) if low_cutoff is None else np.float32(low_cutoff)
    low_width = np.float32(low_width)
    high_cutoff = np.float32(1) if high_cutoff is None else np.float32(high_cutoff)
    high_width = np.float32(high_width)
    register_translation_2d_paraboloid_cuda.cu_cc_fourier_conj_filter(d1f, d2f, low_cutoff, low_width, high_cutoff,
                                                                      high_width, ny, nx)

    # c2r fft to get CC map
    cc0 = pyvkfft_irfftn(d1f, ndim=2)

    # Pixel registration
    idx = register_translation_2d_paraboloid_cuda.cu_argmax_f_red(cc0)

    # how many images ?
    if ref_img.ndim == 2:
        is_2d = True
        sh = 1
    else:
        is_2d = False
        sh = ref_img.shape[:-2]  # Keep the shift shape same as input images extra dimensions

    # Paraboloid 3x3 pixels fit
    cy, cx = cua.empty(sh, dtype=np.float32), cua.empty(sh, dtype=np.float32)
    ny, nx = np.int32(cc0.shape[0]), np.int32(cc0.shape[1])

    register_translation_2d_paraboloid_cuda.cu_paraboloid9(cy, cx, idx, cc0, ny, nx)
    if return_cc:
        if return_gpu_arrays:
            return cy, cx, cc0
        if is_2d:
            return cy.get()[0], cx.get()[0], cc0.get()
        else:
            return cy.get(), cx.get(), cc0.get()
    if return_gpu_arrays:
        return cy, cx
    if is_2d:
        return cy.get()[0], cx.get()[0]
    else:
        return cy.get(), cx.get()


register_translation_2d_paraboloid_cuda.cu_paraboloid9 = None
register_translation_2d_paraboloid_cuda.cu_argmax_f_red = None
register_translation_2d_paraboloid_cuda.cu_cc_fourier_conj_filter = None
