#!/home/esrf/favre/dev/pynx-py38-power9-env/bin/python
# -*- coding: utf-8 -*-

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

import numpy as np
import pycuda.driver as cu_drv
import pycuda.gpuarray as cua
import pycuda.tools as cu_tools
from pycuda.elementwise import ElementwiseKernel as CU_ElK
from pynx.processing_unit.kernel_source import get_kernel_source as getks
from pyvkfft.fft import rfftn as vkrfftn, irfftn as vkirfftn, fftn as vkfftn, ifftn as vkifftn
from numpy.fft import rfftfreq, fftfreq


def gaussian_filter_rfft(sh, sigma_r=0.04, qr0=0.1, sigma_t=0.005, theta=np.deg2rad(30), max=4):
    """
    Compute a Gaussian filter directly in Fourier space. It serves the same
    role as the Gabor filter in VSNR filtering.

    :param sh: array shape
    :param sigma_r: sigma in the radial direction
    :param qr0: relative frequency centre in the radial direction
    :param sigma_t: sigma in the tangential direction
    :param theta: angle of the streak in reciprocal space
    :param max: the maximum value for the filter
    :return: the half-Hermitian array for the filter, real values only
    """
    qx = rfftfreq(sh[1]).astype(np.float32)
    qy = fftfreq(sh[0]).astype(np.float32)
    qy = np.expand_dims(qy, axis=1)
    qr = qx * np.cos(theta) + qy * np.sin(theta)
    qt = -qx * np.sin(theta) + qy * np.cos(theta)

    g = max * np.exp(-.5 * ((qr - qr0) ** 2 / sigma_r ** 2 + qt ** 2 / sigma_t ** 2))
    return g.astype(np.float32)


def vsnr_filter1(gu0, gpsif, beta, niter, apod_width: int = 0, stream=None):
    """
    Perform a VSNR ADMM filtering given a stack of 2D images and a filter

    :param gu0: the cuda array to be filtered - should have been allocated
        using a pycuda memory pool for more efficient allocation of the
        temporary arrays. The array will be normalised automatically
        during the process.
    :param gpsif: the cuda array of the filter (real-valued). It should have
        been normalised with a maximum value of 0.01
    :param niter: number of iterations to perform
    :param apod_width: if >0, use an apodisation window of a given number
        of pixels over which the change in intensity will be tapered using
        a cosine function
    :param stream: the cuda stream to use
    :return: nothing. The d cuda array is modified in-place
    """
    if vsnr_filter1.cu_vsnr1 is None:
        # GPU kernels
        vsnr_filter1.cu_vsnr1 = CU_ElK(name='vsnr1',
                                       operation="y[i] = -l[i]+beta*y[i];",
                                       options=["-use_fast_math"],
                                       arguments="float *y, float *l, const float beta")

        vsnr_filter1.cu_vsnr2 = CU_ElK(name='vsnr2',
                                       operation="vsnr2(i, yx, yy, psif, ny, nxh, beta)",
                                       preamble=getks('cuda/complex.cu') + getks('utils/cuda/vsnr.cu'),
                                       options=["-use_fast_math"],
                                       arguments="pycuda::complex<float> *yx, pycuda::complex<float> *yy,"
                                                 "float* psif, const int ny, const int nxh, const float beta")

        vsnr_filter1.cu_vsnr3 = CU_ElK(name='vsnr3',
                                       operation="vsnr3(i, u0, lx, ly, yx, yy, ny, nx, beta)",
                                       preamble=getks('cuda/complex.cu') + getks('utils/cuda/vsnr.cu'),
                                       options=["-use_fast_math"],
                                       arguments="float *u0, float *lx, float *ly, float *yx, float *yy,"
                                                 "const int ny, const int nx, const float beta")

        # Same as vsnr2, just store FT(X)*FT(Psi) in gyx
        vsnr_filter1.cu_vsnr4 = CU_ElK(name='vsnr4',
                                       operation="vsnr4(i, yx, yy, psif, ny, nxh, beta)",
                                       preamble=getks('cuda/complex.cu') + getks('utils/cuda/vsnr.cu'),
                                       options=["-use_fast_math"],
                                       arguments="pycuda::complex<float> *yx, pycuda::complex<float> *yy,"
                                                 "float* psif, const int ny, const int nxh, const float beta")

        vsnr_filter1.cu_scale = CU_ElK(name='scale',
                                       operation="y[i] /= s[0]/nb;",
                                       options=["-use_fast_math"],
                                       arguments="float *y, float *s, const int nb")

        vsnr_filter1.cu_descale = CU_ElK(name='descale',
                                         operation="y[i] *= s[0]/nb;",
                                         options=["-use_fast_math"],
                                         arguments="float *y, float *s, const int nb")
        vsnr_filter1.cu_apod = CU_ElK(name='apod',
                                      operation="apod(i, y, w, ny, nx)",
                                      preamble=getks('cuda/complex.cu') + getks('utils/cuda/vsnr.cu'),
                                      options=["-use_fast_math"],
                                      arguments="float *y, const int w, const int ny, const int nx")

    gyx = cua.zeros_like(gu0)  # Y
    gyy = cua.zeros_like(gu0)
    glx = cua.zeros_like(gu0)  # lambda
    gly = cua.zeros_like(gu0)
    beta = np.float32(beta)

    ny = np.int32(gu0.shape[-2])
    nx = np.int32(gu0.shape[-1])
    nxh = nx // 2 + 1

    # Scale array
    s = cua.sum(gu0, stream=stream)
    vsnr_filter1.cu_scale(gu0, s, np.int32(gu0.size), stream=stream)

    for i in range(niter):
        # -lambda_x+beta*y_x in y_x array (same for y)
        vsnr_filter1.cu_vsnr1(gyx, glx, beta, stream=stream)
        vsnr_filter1.cu_vsnr1(gyy, gly, beta, stream=stream)
        # in-place rfft ; resulting array is half-hermitian
        gyx = vkrfftn(gyx, ndim=2, cuda_stream=stream)
        gyy = vkrfftn(gyy, ndim=2, cuda_stream=stream)
        # Compute FT(X)*psif_x and FT(X)*psif_y (temporarily stored in gyx,gyy
        # viewed as half-hermitian arrays) and their inverse FT
        vsnr_filter1.cu_vsnr2(gyx, gyy, gpsif, ny, nxh, beta, stream=stream)
        gyx = vkirfftn(gyx, ndim=2, cuda_stream=stream)
        gyy = vkirfftn(gyy, ndim=2, cuda_stream=stream)
        # Update Y and lambda
        vsnr_filter1.cu_vsnr3(gu0, glx, gly, gyx, gyy, ny, nx, beta, stream=stream)
    # Finish
    vsnr_filter1.cu_vsnr1(gyx, glx, beta, stream=stream)
    vsnr_filter1.cu_vsnr1(gyy, gly, beta, stream=stream)
    gyx = vkrfftn(gyx, ndim=2, cuda_stream=stream)
    gyy = vkrfftn(gyy, ndim=2, cuda_stream=stream)
    vsnr_filter1.cu_vsnr4(gyx, gyy, gpsif, ny, nxh, beta, stream=stream)
    gyx = vkirfftn(gyx, ndim=2, cuda_stream=stream)
    if apod_width:
        vsnr_filter1.cu_apod(gyx, np.int32(apod_width), ny, nx, stream=stream)
    gu0 -= gyx
    vsnr_filter1.cu_descale(gu0, s, np.int32(gu0.size), stream=stream)


# Cheap way to store kernels
vsnr_filter1.cu_vsnr1 = None


def vsnr_filter(u: np.ndarray, filter_params,
                beta: float = 10, niter: int = 5):
    """
    Perform a VSNR ADMM filtering on a stack of 2D images, given a filter.
    This function assumes that the cuda context has already been initialised.

    :param u: the stack of 2D images to filter, or a single 2D image
    :param filter_params: either a single dictionary of parameters
        for the Gaussian filter, or a list of dictionaries to apply
        multiple filters. Dict keys are: sigma_r, qr0, sigma_t, theta, max.
        See gaussian_filter_rfft doc for the parameters.
    :param beta: the beta ADMM parameter
    :param niter: number of cycles for the filtering
    :return: the filtered u array
    """
    # reshape array to 3D
    sh0 = list(u.shape)
    if u.ndim == 2:
        u = u.reshape([1] + sh0)
    elif u.ndim > 3:
        u = u.reshape([np.prod(sh0[:-2])] + sh0[-2:])
    nb = len(u)

    # Gaussian filter
    if type(filter_params) is dict:
        filter_params = [filter_params]

    # Use streams for // transfer of data to/from GPU and computing
    cu_stream = cu_drv.Stream()
    cu_stream_swap = cu_drv.Stream()
    cu_event_calc = cu_drv.Event()
    cu_event_swap = cu_drv.Event()

    # Move first array to GPU. Use two streams to compute & I/O to/from GPU
    cu_mem_pool = cu_tools.DeviceMemoryPool()
    gu_swap = cua.to_gpu_async(u[0], allocator=cu_mem_pool.allocate, stream=cu_stream_swap)
    gu0 = cua.zeros_like(gu_swap)
    vgpsif = [cua.to_gpu_async(gaussian_filter_rfft(sh0[-2:], **kw),
                               allocator=cu_mem_pool.allocate, stream=cu_stream_swap)
              for kw in filter_params]
    for i in range(nb):
        # Compute stream waits for swap, swap waits for compute
        cu_event_swap.record(cu_stream_swap)
        cu_event_calc.record(cu_stream)
        cu_stream.wait_for_event(cu_event_swap)
        cu_stream_swap.wait_for_event(cu_event_calc)
        # Switch arrays
        gu0, gu_swap = gu_swap, gu0
        for gpsif in vgpsif:
            vsnr_filter1(gu0, gpsif, beta, niter, stream=cu_stream)
        if i:
            cu_drv.memcpy_dtoh_async(src=gu_swap.gpudata, dest=u[i - 1], stream=cu_stream_swap)
        if i < nb - 1:
            cu_drv.memcpy_htod_async(dest=gu_swap.gpudata, src=u[i + 1], stream=cu_stream_swap)
    cu_drv.memcpy_dtoh_async(src=gu_swap.gpudata, dest=u[-1], stream=cu_stream)
    cu_stream.synchronize()
    return u.reshape(sh0)
