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

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2018-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr
import logging
import re
import ctypes
import numpy as np
from scipy.ndimage import zoom, gaussian_filter, map_coordinates, median_filter
from scipy.interpolate import interpn, RectBivariateSpline, RegularGridInterpolator
from scipy.fft import rfftn, irfftn, rfftfreq, fftfreq
from scipy.special import erfc
from scipy.signal import medfilt2d as medfilt2d_scipy
import fabio
import silx
from silx.image.tomography import compute_fourier_filter
from silx.math.medianfilter import medfilt2d
from tomoscan.esrf.scan.edfscan import EDFTomoScan

from ..utils.registration import phase_cross_correlation, phase_cross_correlation_paraboloid, shift
from ..utils.array import rebin, pad2
from ..utils.time import catch_time

# Module dictionary to hold shared multiprocessing.Array when using pools
shared_arrays = {}


def remove_spikes(img, threshold=0.04):
    """
    Remove 'spikes' from an images by comparing it to its median-filtered value.
    :param img: the image to filter
    :param threshold: the threshold (as a % of the image value) above which the pixel will be
        replaced by the filtered value.
    :return: nothing. The input image is modified in-place
    """
    if False:
        # In some cases, this hangs forever ??
        imgf = medfilt2d(img, 3)
    else:
        imgf = medfilt2d_scipy(img, 3)
        # Need to take care of corners, as filtering is done on a virtually zero-padded array,
        # which results in zeros in the 4 corners...
        imgf[0, 0] = 0.5 * (imgf[0, 1] + imgf[1, 0])
        imgf[0, -1] = 0.5 * (imgf[0, -2] + imgf[1, -1])
        imgf[-1, 0] = 0.5 * (imgf[-1, 1] + imgf[-2, 0])
        imgf[-1, -1] = 0.5 * (imgf[-1, -2] + imgf[-2, -1])
    img[:] = np.where(abs(img - imgf) > (threshold * imgf), imgf, img)


def simulate_probe(shape, dz, pixel_size=1e-6, wavelength=0.5e-10, nb_line_v=10,
                   nb_line_h=10, nb_spot=10, defect_amplitude=1, amplitude=100):
    """
    Create a simulated probe corresponding to a HoloTomo object size, with vertical and horizontal lines and spots
    coming from optics defects, which are then propagated.

    :param shape: the 2D shape (ny, nx) of the probe
    :param dz: array of propagation distances (m) (unrelated to holo-tomo distances)
    :param pixel_size: detector pixel size (m)
    :param wavelength: the wavelength (m)
    :param nb_line_v: number of vertical lines. Width = max 5% of the probe horizontal size.
    :param nb_line_h: number of horizontal lines. Width = max 5% of the probe vertical size
    :param nb_spot: number of spots. Radius= max 5% of the probe horizontal size
    :param defect_amplitude: the relative amplitude of the introduced optical defects (default:1).
    :param amplitude: the average amplitude of the calculated wavefront. The square amplitude shall
        correspond to the average number of incident photons per pixel.
    :return: the simulated probe, with shape (nz, nbmode=1, ny, nx)
    """
    from ..wavefront import Wavefront, PropagateNearField
    ny, nx = shape
    nz = len(dz)
    # take convenient dimensions for the wavefront
    d = np.zeros((nz, ny, nx), dtype=np.complex64)
    for j in range(nz):
        for i in range(nb_line_v):
            w = 1 + np.random.randint(0, nx * 0.05)
            ii = np.random.randint(0, nx - w)
            t = np.random.randint(10)
            d[j, :, ii:ii + w] = t

        for i in range(nb_line_h):
            w = 1 + np.random.randint(0, ny * 0.05)
            ii = np.random.randint(0, ny - w)
            t = np.random.randint(10)
            d[j, ii:ii + w] = t
        x, y = np.meshgrid(np.arange(0, nx), np.arange(0, ny))

        for i in range(nb_spot):
            w = 1 + np.random.randint(0, nx * 0.05)
            ix = np.random.randint(w, nx - w)
            iy = np.random.randint(w, ny - w)
            r = np.sqrt((x - ix) ** 2 + (y - iy) ** 2)
            t = np.random.uniform(0, 10)

            d[j] += t * (r < w)
        w = Wavefront(d=np.fft.fftshift(np.exp(1j * 1e-2 * d[j] * defect_amplitude)), pixel_size=pixel_size,
                      wavelength=wavelength)
        w = PropagateNearField(dz=dz[j]) * w
        d[j] = w.get(shift=True).reshape(ny, nx)
        d[j] *= amplitude / abs(d[j]).mean()
    return d.reshape((nz, 1, ny, nx))


def zoom_pad_images(x, magnification, padding, nz, pad_method='reflect_linear',
                    mask2neg=True):
    """
    Zoom & pad images

    :param x: the array of shape (nz, ny, nx) corresponding to one projection data
    :param magnification: the nz values for the magnification
    :param padding: the size of the padding of each side of the images
    :param nz: number of distances
    :param pad_method: the padding mode, either one from numpy, or among 'reflect_linear',
        'reflect_erf', 'reflect_sine'. The padded areas are equal to the reflected array,
        but with an interpolation allowing to avoid discontinuity at the array borders.
    :param mask2neg: parameter which will be passed to pad2() to flag the padded
        values as extrapolated (not observed).
    :return: the zoomed & padded data, of shape (nz, ny+2*padding, nx+2*padding)
    """
    d0 = x[0]
    pady0, padx0 = padding
    ny0, nx0 = d0.shape[0] + 2 * pady0, d0.shape[1] + 2 * padx0
    xz = np.empty((nz, ny0, nx0), dtype=np.float32)
    for iz in range(nz):
        # Zoom
        mg = magnification[0] / magnification[iz]
        if np.isclose(mg, 1):
            d = x[iz]
        else:
            d = zoom(x[iz], zoom=mg, order=1, mode='nearest', prefilter=False)
        ny, nx = d.shape
        # easier with even shape, symmetric padding...
        if (nx0 - nx) % 2:
            d = d[:, :-1]
            nx -= 1
        if (ny0 - ny) % 2:
            d = d[:-1]
            ny -= 1

        if ny0 <= ny and nx0 <= nx:
            xz[iz] = d[ny // 2 - ny0 // 2:ny // 2 - ny0 // 2 + ny0, nx // 2 - nx0 // 2:nx // 2 - nx0 // 2 + nx0]
        elif ny0 > ny and nx0 <= nx:
            pady = (ny0 - ny) // 2
            xz[iz] = pad2(d[:, nx // 2 - nx0 // 2:nx // 2 - nx0 // 2 + nx0], (pady, 0),
                          mode=pad_method, mask2neg=mask2neg)
        elif ny0 <= ny and nx0 > nx:
            padx = (nx0 - nx) // 2
            xz[iz, :, nx0 // 2 - nx // 2:nx0 // 2 - nx // 2 + nx] = \
                pad2(d[ny // 2 - ny0 // 2:ny // 2 - ny0 // 2 + ny0], (0, padx),
                     mode=pad_method, mask2neg=mask2neg)
        else:
            padx = (nx0 - nx) // 2
            pady = (ny0 - ny) // 2
            # print(d.shape, xz[iz].shape, pady, padx)
            xz[iz] = pad2(d, (pady, padx), mode=pad_method, mask2neg=mask2neg)
    return xz


def zoom_pad_images_init(iobs_array):
    """
    Initialiser for the pool, to pass the shared arrays
    :param iobs_array:
    :return:
    """
    shared_arrays["iobs_array"] = iobs_array


def zoom_pad_images_kw(kwargs):
    nx, ny, nz = kwargs['nx'], kwargs['ny'], kwargs['nz']
    nb_proj_sharray = kwargs['nb_proj_sharray']
    pady, padx = kwargs['padding']
    pad_method = kwargs['pad_method']
    iobs_array = shared_arrays['iobs_array']
    i = kwargs['i']
    iobs_shape = (nb_proj_sharray, nz, ny + 2 * pady, nx + 2 * padx)
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
    iobs[i] = zoom_pad_images(iobs[i, :, :ny, :nx], kwargs['magnification'], padding=(pady, padx),
                              nz=nz, pad_method=pad_method)


def vsnr_init(iobs_array, ref_array):
    """
    Initialiser for the pool, to pass the shared arrays
    :param iobs_array: observed data array
    :param ref_array: reference (direct beam) array
    :return:
    """
    shared_arrays["iobs_array"] = iobs_array
    shared_arrays["ref_array"] = ref_array


def vsnr_kw(kwargs):
    """
    Apply a VSNR correction on the stack of observed data (and optionally reference)
    :param kwargs:
    :return:
    """
    nx, ny, nz = kwargs['nx'], kwargs['ny'], kwargs['nz']
    pady, padx = kwargs['padding']
    nb_proj_sharray = kwargs['nb_proj_sharray']
    idx = kwargs['idx']
    filter_params = kwargs['filter_params']
    niter = kwargs['niter']
    beta = kwargs['beta']
    igpu = kwargs['igpu']
    do_ref = kwargs['do_ref']
    if kwargs['roi'] is not None:
        y0, y1, x0, x1 = kwargs['roi']
    else:
        y0, y1, x0, x1 = 0, ny, 0, nx
    # We want a multiple of 2 size along all axes
    x1 -= (x1 - x0) % 2
    y1 -= (y1 - x0) % 2

    # Init cuda environment
    from ..processing_unit import default_processing_unit as pu
    pu.select_gpu(language='cuda', ranking='order', gpu_rank=igpu)
    from pynx.utils.vsnr_cuda import vsnr_filter

    # Extract arrays
    iobs_array = shared_arrays['iobs_array']
    iobs_shape = (nb_proj_sharray, nz, ny + 2 * pady, nx + 2 * padx)
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
    iobs = iobs[..., y0:y1, x0:x1]  # This should return a slice
    tmp = np.ascontiguousarray(iobs[idx])  # iobs[idx] returns a copy
    tmp = vsnr_filter(tmp, filter_params=filter_params, niter=niter, beta=beta)
    for i in range(len(idx)):
        iobs[idx[i]] = tmp[i]
    if do_ref:
        ref_array = shared_arrays['ref_array']
        ref = np.frombuffer(ref_array.get_obj(), dtype=ctypes.c_float).reshape((nz, ny, nx))
        ref = ref[..., y0:y1, x0:x1]
        ref[:] = vsnr_filter(np.ascontiguousarray(ref), filter_params=filter_params, niter=niter, beta=beta)


def shift_images_kw(kwargs):
    """
    fft-shift a stack of nz 2D images
    :param kwargs: dictionary with i (projection index), padding (pady, padx tuple of
        padding on each side), nproj (number of projections), nz (number of planes),
        ny, nx (2D image size without padding), dx and dy (each a list of nz shifts).
    :return:
    """
    i = kwargs['i']
    pady, padx = kwargs['padding']
    nproj, nz, ny, nx = kwargs['nproj'], kwargs['nz'], kwargs['ny'], kwargs['nx']
    dx, dy = kwargs['dx'], kwargs['dy']
    iobs_array = shared_arrays['iobs_array']
    iobs_shape = (nproj, nz, ny + 2 * pady, nx + 2 * padx)
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
    for iz in range(nz):
        # We must convert padded values which are <0 (stored as -Ipad-1)
        # since the jump from >0 to <0 values creates ripples during the fftshift
        tmp1 = iobs[i, iz] < 0
        iobs[i, iz][tmp1] = -iobs[i, iz][tmp1] - 1
        iobs[i, iz] = shift(iobs[i, iz], (dy[iz], dx[iz]))
        iobs[i, iz][tmp1] = -iobs[i, iz][tmp1] - 1


def align_images(x, x0=None, nz=4, low_cutoff=None,
                 low_width=0.03, high_cutoff=None, high_width=0.03,
                 align_max_shift=None, binning=1):
    """
    Determine shifts between a series of images at different z values.

    :param x: the images to align, which can either have a shape (nz, ny, nx)
        or (3, nz, ny, nx). In the latter case, the 3 images will be averaged
        to use the 'difference' image with the average, and be less sensitive
        to the constant features.
    :param x0: the reference image (empty beam) by which the images will be
        normalised before alignment
    :param nz: number of planes to align
    :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 low_width: the width of the low cutoff filter, also as a percentage of the size
    :param high_cutoff: same as low_cutoff fot the high frequency filter, should be close below 0.5
    :param high_width: same as low_width
    :param align_max_shift: this is used to set to zero a margin on one of the
        images, so that the cross-correlation (CC) is correctly computed without wrapping
    :param binning: if >1, the images will be binned to determine the shift, to lower
        any noise effect (NB: in practice, does not seem to have much benefit)
    :return: dx, dy, each an array of nz values, with the shift relative to the
        first plane (so dx[0] = dy[0]=0).
    """
    if align_max_shift is None:
        align_max_shift = 0
    dx = [0]
    dy = [0]

    # Binning the image may be better for noisy images
    if binning > 1:
        if x.ndim == 3:
            x = rebin(x, (1, binning, binning))
        else:
            x = rebin(x, (1, 1, binning, binning))
        if x0 is not None:
            x0 = rebin(x0, (1, binning, binning))

    if x.ndim == 3:
        d0 = x[0].copy()
    else:
        if True:
            # Subtract mean
            d0 = x[1, 0] - x[:, 0].mean(axis=0)
        else:
            # Subtract median - faster version than np.median
            iz = 0
            c01, c02, c12 = x[0, iz] >= x[1, iz], x[0, iz] >= x[2, iz], x[1, iz] >= x[2, iz]
            c10, c20, c21 = np.logical_not(c01), np.logical_not(c02), np.logical_not(c12)
            d0 = np.select([(c21 & c10) | (c01 & c12), (c12 & c20) | (c02 & c21)],
                           (0, x[1, iz] - x[2, iz]), default=x[1, iz] - x[0, iz])

    if x0 is not None:
        d0 /= x0[0]
    d0 -= d0.mean()  # np.percentile(d0, 0.001)
    if align_max_shift:
        # We cut out the borders of the images, so that the CC will not wrap,
        # which in principle allows to find the shift where the two images are
        # fully overlapping. This is particularly useful for highly structured
        # images up to their boundaries.
        # Only one image should be zero-padded.
        d0[:align_max_shift] = 0
        d0[-align_max_shift:] = 0
        d0[:, :align_max_shift] = 0
        d0[:, -align_max_shift:] = 0
    # d0_mask = x[0] >= 0  # Ignore masked areas ?
    # d0 = rfftn(d0)
    for iz in range(1, nz):
        #         print(i,iz)
        # Align
        if x.ndim == 3:
            d = x[iz].copy()
        else:
            # Use differential alignment
            if True:
                # Subtract mean
                d = x[1, iz] - x[:, iz].mean(axis=0)
            else:
                # Subtract median - faster version than np.median
                c01, c02, c12 = x[0, iz] >= x[1, iz], x[0, iz] >= x[2, iz], x[1, iz] >= x[2, iz]
                c10, c20, c21 = np.logical_not(c01), np.logical_not(c02), np.logical_not(c12)
                d = np.select([(c21 & c10) | (c01 & c12), (c12 & c20) | (c02 & c21)],
                              (0, x[1, iz] - x[2, iz]), default=x[1, iz] - x[0, iz])
        if x0 is not None:
            d /= x0[iz]
        d -= d.mean()  # np.percentile(d, 0.001)
        # pixel_shift = phase_cross_correlation(d0, d1, upsample_factor=upsample_factor,
        #                                       reference_mask=None, moving_mask=None,
        #                                       return_error=False, normalization=None,
        #                                       low_cutoff=low_cutoff, low_width=low_width,
        #                                       high_cutoff=high_cutoff, high_width=high_width,
        #                                       space='fourier', r2c=True)
        pixel_shift = phase_cross_correlation_paraboloid(d0, d, low_cutoff=low_cutoff, low_width=low_width,
                                                         high_cutoff=high_cutoff, high_width=high_width)
        if iz < nz - 1:
            # Update d0 for next plane. Always compare to the previous distance
            d0 = d
            if align_max_shift:
                # See explanation above
                d0[:align_max_shift] = 0
                d0[-align_max_shift:] = 0
                d0[:, :align_max_shift] = 0
                d0[:, -align_max_shift:] = 0
        # d0 = rfftn(d0)
        dx.append(pixel_shift[1] * binning + dx[-1])
        dy.append(pixel_shift[0] * binning + dy[-1])
    return dx, dy


def align_images_pool_init(iobs_array, ref_array=None):
    """
    Initial shared arrays for the pool process
    :param iobs_array: the shared data array for observed intensities
    :param ref_array: the shared data array for the reference (empty beam) intensities.
    :return: nothing
    """
    shared_arrays["iobs_array"] = iobs_array
    if ref_array is not None:
        shared_arrays["ref_array"] = ref_array


def align_images_kw(kwargs):
    """
    Determine the relative shifts between images in a stack. This uses the
    shared iobs_array and ref_array
    :param kwargs: dictionary with the parameters, including:
        - padding: (pady, padx) padding margins on each side
        - nproj: number of projections in the shared array
        - nz: number of distances which must be aligned
        - ny, nx: 2D image size without the padding
        - normalise: if True, will normalise the images with the reference data
        - diff: if True, will align the difference between a projection and the rolling
            mean of the image with the previous and next.
    :return: the tuple (dx, dy) of shifts (nz values) relative to the first z plane.
    """
    try:
        i = kwargs['i']
        pady, padx = kwargs['padding']
        nproj, nz, ny, nx = kwargs['nproj'], kwargs['nz'], kwargs['ny'], kwargs['nx']
        normalise = kwargs['normalise']
        diff = kwargs['diff']
        del kwargs['i'], kwargs['padding'], kwargs['ny'], kwargs['nx'], kwargs['nproj'], \
            kwargs['normalise'], kwargs['diff']

        iobs_array = shared_arrays['iobs_array']
        iobs_shape = (nproj, nz, ny + 2 * pady, nx + 2 * padx)
        iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
        if diff:
            if True:
                if pady and padx:
                    x = iobs[i - 1 if i >= 1 else i: i + 2, :, pady:-pady, padx:-padx]
                elif pady:
                    x = iobs[i - 1 if i >= 1 else i: i + 2, :, pady:-pady]
                elif padx:
                    x = iobs[i - 1 if i >= 1 else i: i + 2, :, :, padx:-padx]
                else:
                    x = iobs[i - 1 if i >= 1 else i: i + 2]
            else:
                # Use frames before and after, or the two closest since we use a median filter
                if pady and padx:
                    tmp = iobs[..., pady:-pady, padx:-padx]
                elif pady:
                    tmp = iobs[..., pady:-pady, :]
                elif padx:
                    tmp = iobs[..., padx:-padx]
                else:
                    tmp = iobs
                if i == 0:
                    x = np.array((tmp[i + 2], tmp[i], tmp[i + 1]))
                elif i == nproj - 1:
                    x = np.array((tmp[i - 1], tmp[i], tmp[i - 2]))
                else:
                    x = tmp[i - 1:i + 2]
        else:
            if pady and padx:
                x = iobs[i, :, pady:-pady, padx:-padx]
            elif pady:
                x = iobs[i, :, pady:-pady]
            elif padx:
                x = iobs[i, :, :, padx:-padx]
            else:
                x = iobs[i]

        if normalise:
            ref_array = shared_arrays['ref_array']
            ref_shape = (nz, ny + 2 * pady, nx + 2 * padx)
            ref = np.frombuffer(ref_array.get_obj(), dtype=ctypes.c_float).reshape(ref_shape)
            if pady and padx:
                x0 = ref[:, pady:-pady, padx:-padx]
            elif pady:
                x0 = ref[:, pady:-pady]
            elif padx:
                x0 = ref[..., padx:-padx]
            else:
                x0 = ref
            return align_images(x, x0=x0, binning=1, **kwargs)
        else:
            return align_images(x, x0=None, binning=1, **kwargs)
    except Exception as ex:
        import traceback
        print(ex)
        print(traceback.format_exc())
        raise ex


def calc_distortion(imgz, refz, nz, tile_size=100, low_cutoff=None,
                    low_width=0.03, high_cutoff=None, high_width=0.03,
                    distort_img=True, return_shifts=False):
    """
    Distort projections to align them to reference images. The
    :param imgz: the array (nz, ny, nx) of projection images which will be distorted
    :param refz: the array (nz, ny, nx) of reference (empty beam) images
    :param nz: number of propagation distances.
    :param tile_size: size of the tiles into which the images will be divided to map
        the distortion
    :param low_cutoff: for the 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 low_width: the width of the low cutoff filter, also as a percentage of the size
    :param high_cutoff: same as low_cutoff fot the high frequency filter, should be close below 0.5
    :param high_width: same as low_width
    :param distort_img: if True, the images (projections) will be distorted to match the references.
    :param return_shifts: if True, return the computed shifts along with the center coordinates
        of the tiles used.
    :return: if return_shifts==True, return a tuple of arrays (cy, cx, sy, sx) with cy and cx
        the center coordinates of the tiles and sy, sx the shifts along the horizontal directions.
        The shape of the arrays will be (ny//tile_size, nx//tile_size) for cy and cx and
        (nz, ny//tile_size, nx//tile_size) for sy and sx.
        If distort_img==True, the projection images will be distorted.
    """
    # naive implementation, distort independently all distances
    ny, nx = imgz.shape[-2:]
    # Make sure we have 3d even if nz==1
    imgz = imgz.reshape((nz, ny, nx))
    refz = refz.reshape((nz, ny, nx))
    # First generate the tile coordinates, spanning as much as possible of the image
    nty, ntx = ny // tile_size, nx // tile_size
    cy = np.linspace(0, ny, nty + 1, dtype=np.int32)[:-1]
    cx = np.linspace(0, nx, ntx + 1, dtype=np.int32)[:-1]
    cyc, cxc = cy + tile_size / 2, cx + tile_size / 2
    vy0, vx0 = np.arange(ny, dtype=np.float32), np.arange(nx, dtype=np.float32)
    vy1, vx1 = np.meshgrid(vy0, vx0, indexing="ij")
    img = np.empty((nty, ntx, tile_size, tile_size), dtype=np.float32)
    ref = np.empty((nty, ntx, tile_size, tile_size), dtype=np.float32)
    shifts = []
    for iz in range(nz):
        for iy in range(nty):
            for ix in range(ntx):
                img[iy, ix] = imgz[iz, cy[iy]:cy[iy] + tile_size, cx[ix]:cx[ix] + tile_size]
                ref[iy, ix] = refz[iz, cy[iy]:cy[iy] + tile_size, cx[ix]:cx[ix] + tile_size]
        # # We are assuming a small distortion - set a small margin of zeros on the side
        # #  so the CC is more correctly computed.
        # # TODO: check this approach. Need to actually limit the shift in the phase_cc function ?
        # m = 3  # This allows to compute the CC correctly up to 6 pixels shift
        # ref[:, :, :m] = 0
        # ref[:, :, -m:] = 0
        # ref[:, :, :, :m] = 0
        # ref[:, :, :, -m:] = 0
        sy, sx = phase_cross_correlation_paraboloid(img, ref, low_cutoff=low_cutoff, low_width=low_width,
                                                    high_cutoff=high_cutoff, high_width=high_width)
        if distort_img:
            sy = median_filter(sy, 3, mode='nearest')
            sx = median_filter(sx, 3, mode='nearest')
            if True:
                spline_degree = 1
                interpolator = RectBivariateSpline(cyc, cxc, sy, kx=spline_degree, ky=spline_degree)
                dy = interpolator(vy0, vx0)
                interpolator = RectBivariateSpline(cyc, cxc, sx, kx=spline_degree, ky=spline_degree)
                dx = interpolator(vy0, vx0)
            else:
                dy = zoom(sy, (ny / nty, nx / ntx), mode='reflect', order=1, grid_mode=True)
                dx = zoom(sx, (ny / nty, nx / ntx), mode='reflect', order=1, grid_mode=True)

            vy2 = vy1 + dy
            vx2 = vx1 + dx

            # imgz[iz] = interpn((vy, vx), imgz[iz],, method = 'cubic')

            imgz[iz] = map_coordinates(imgz[iz], (vy2, vx2), output=np.float32, order=3, mode='reflect',
                                       prefilter=False)

            # coords = np.transpose(np.array([vy2, vx2]), axes=[1, 2, 0])
            # interpolator = RegularGridInterpolator((vy0, vx0), imgz[iz], bounds_error=False, method="linear", fill_value=None)
            # imgz[iz] = interpolator(coords)

            # Slow & wrong ?
            # interpolator = RectBivariateSpline(vy0, vx0, imgz[iz], kx=3, ky=3, s=0)
            # imgz[iz] = interpolator(vx2, vy2, grid=False)

        shifts.append((sy, sx))
    if return_shifts:
        return shifts


def load_data(proj_urlz, dark, planes, binning=None, ref=None, distorted=0, normalise=False, timing=False):
    """
    Load images from files.
    :param proj_urlz: the url (file path) for the nz projections at different distances
    :param dark: the dark image array, of size (nz, ny, nx)
    :param planes: a tuple of planes to use with nz values. Typically (1,2,3,4) but
        some planes can be omitted.
    :param binning: factor (integer>1) by which the images will be binned
    :param ref: reference array of shape (nz, ny, nx), by which the images will be divided
        if normalise==True
    :param distorted: if > 1, the distortion between the projections and the reference
        images will be computed, and the projections will be distorted to adapt, unless
        normalise==True, in which case the reference images will be distorted before
        normalisation. The given value (typically 64 or 100) corresponds to the tile
        size which will be used for the distortion.
    :param normalise: if Trye, the images will be normalised by the reference images.
    :param timing: if True, will also return a dictionary with the time spent on
        different parts (loading, normalise, distort)
    :return: the stack of loaded images with a shape (nz, ny, nx). If distorted>0, the
        shift is also returned.
    """
    d = None
    nz = len(planes)
    timings = {'load-bin-dark': 0, 'distort': 0, 'normalise': 0}
    with catch_time() as timer:
        for iz in range(nz):
            img = fabio.open(proj_urlz[iz]).data
            if binning is not None:
                if binning > 1:
                    img = rebin(img, binning)
            if d is None:
                ny, nx = img.shape
                d = np.empty((nz, ny, nx), dtype=np.float32)
            d[iz] = img
            if dark is not None:
                d[iz] -= dark[iz]
    timings['load-bin-dark'] = timer.dt
    if normalise:
        ref2 = ref.copy()
        if distorted:
            with catch_time() as timer:
                shifts = calc_distortion(ref2, d, nz=nz, tile_size=distorted, low_cutoff=0.03, return_shifts=True)
            timings['distort'] = timer.dt
        with catch_time() as timer:
            for iz in range(nz):
                d[iz] = d[iz] / ref2[iz] * ref2[iz].mean()
        timings['normalise'] = timer.dt
    else:
        if distorted:
            with catch_time() as timer:
                shifts = calc_distortion(d, ref, nz=nz, tile_size=distorted, low_cutoff=0.03, return_shifts=True)
        timings['distort'] = timer.dt
    if distorted:
        if timing:
            return d, shifts, timings
        else:
            return d, shifts
    if timing:
        return d, timings
    return d


def load_data_pool_init(iobs_array, dark_array, ref_array):
    """
    Initialiser for the pool, to pass the shared arrays
    :param iobs_array:
    :param dark_array:
    :return:
    """
    shared_arrays["iobs_array"] = iobs_array
    shared_arrays["dark_array"] = dark_array
    shared_arrays["ref_array"] = ref_array


def load_data_kw(kwargs):
    nproj, nx, ny, planes = kwargs['nproj'], kwargs['nx'], kwargs['ny'], kwargs['planes']
    nz = len(planes)
    distorted = kwargs['distorted']
    normalise = kwargs['normalise']
    pady, padx = kwargs['padding']
    iobs_array = shared_arrays['iobs_array']
    dark_array = shared_arrays['dark_array']
    ref_array = shared_arrays['ref_array']
    i = kwargs['i']
    proj_urlz = kwargs['proj_urlz']  # list of URL for the data file for each plane.
    spikes_threshold = kwargs['spikes_threshold']
    iobs_shape = (nproj, nz, ny + 2 * pady, nx + 2 * padx)
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
    dark = np.frombuffer(dark_array.get_obj(), dtype=ctypes.c_float).reshape((nz, ny, nx))
    ref = np.frombuffer(ref_array.get_obj(), dtype=ctypes.c_float).reshape((nz, ny, nx))
    if distorted:
        d, shifts, timings = load_data(proj_urlz, dark, planes, kwargs['binning'], ref, distorted=distorted,
                                       normalise=normalise, timing=True)
    else:
        d, timings = load_data(proj_urlz, dark, planes, kwargs['binning'], ref, distorted=distorted,
                               normalise=normalise, timing=True)
    if spikes_threshold > 0:
        for iz in range(nz):
            remove_spikes(d[iz], spikes_threshold)
    iobs[i, :, :ny, :nx] = d
    if distorted:
        return shifts, timings
    else:
        return timings


def calc_apply_double_flat_field(iz, nproj, nb_proj_sharray, nx, ny, nz, padx, pady, ref, freqs, nb_workers):
    """
    Compute a 'double flat-field', i.e. the average of all projections, and normalise them it.
    :param iz: z index of the projections to normalise
    :param nproj: number of projections
    :param nb_proj_sharray: number of projections in the shared array
    :param nx: x size of the images (before padding)
    :param ny: y size of the images (before padding)
    :param nz: number of planes
    :param padx: padding size on each size along x
    :param pady: padding size on each size along y
    :param ref: the reference array by which the images need to be divided before computing
        the double flat-field.
    :param freqs: tuple (high_cutoff, low_cutoff). high_cutoff applies a 2D high-pass
        filter (inpractive should be within [0.25;0.5[) to the 2D average. The
        images are then normalised by the average value + the filtered array.
        If low_cutoff is nonzero, then a low-pass filter is applied along the direction
        of the nproj projections - in order to filter out outliers. This is a 3D filter,
        and therefore quite slow.
    :param nb_workers: number of workers for the fft, when used
    :return: nothing - the iobs array is changed in-place.
    """
    iobs_array = shared_arrays['iobs_array']
    iobs_shape = (nb_proj_sharray, nz, ny + 2 * pady, nx + 2 * padx)
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
    if nb_workers < 1:
        nb_workers = 1
    if freqs[1] < 1e-3:
        # Not a 3D filter
        v = (iobs[:nproj, iz, :ny, :nx] / ref).mean(axis=0)
        if freqs[0] < 1e-3:
            # Remove any null values (?? ID16A ??)
            av = v.mean()
            v[v <= (av / 1000)] = av / 1000  # :KLUDGE ?
            # Normalise by the average <radios/ref>
            iobs[:nproj, iz, :ny, :nx] /= v
        else:
            # 2D high-pass filtering
            av = v.mean()
            v = rfftn(v.astype(np.float32, copy=False), workers=nb_workers)
            r = np.sqrt(rfftfreq(nx).astype(np.float32) ** 2 + fftfreq(ny).astype(np.float32)[:, np.newaxis] ** 2)
            high_cutoff, high_width = freqs[0], freqs[0] if freqs[0] < 0.03 else 0.03
            v *= 1 - 0.5 * erfc((r - high_cutoff) / high_width)
            v = irfftn(v, workers=nb_workers) + av
            v[v <= (av / 1000)] = av / 1000  # :KLUDGE ?
            iobs[:nproj, iz, :ny, :nx] /= v
    else:
        # 3D filtering (varying along the projections axis) - more time-consuming
        # First rebin data along projections for faster calculations
        rb = 2 if nproj < 144 else nproj // 72
        v = rebin(iobs[:nproj, iz, :ny, :nx] / ref, (rb, 1, 1))
        if len(v) % 2:
            # r2c fft uses an even size
            v = v[:-1]
        if freqs[0] >= 1e-3:
            # in-plane filtering
            av = v.mean()
            v = rfftn(v, workers=nb_workers, axes=(-2, -1))
            r = np.sqrt(rfftfreq(nx).astype(np.float32) ** 2 + fftfreq(ny).astype(np.float32)[:, np.newaxis] ** 2)
            high_cutoff, high_width = freqs[0], freqs[0] if freqs[0] < 0.03 else 0.03
            v *= 1 - 0.5 * erfc((r - high_cutoff) / high_width)
            v = irfftn(v, workers=nb_workers, axes=(-2, -1))
        else:
            av = 0
        n = len(v)
        v = rfftn(v, workers=nb_workers, axes=(0,))
        r = rfftfreq(n).astype(np.float32)[:, np.newaxis, np.newaxis]
        low_cutoff, low_width = freqs[1], 0.03 if freqs[1] > 0.03 else freqs[1]
        v *= 0.5 * erfc((r - low_cutoff) / low_width)
        v = irfftn(v, workers=nb_workers, axes=(0,))
        for i in range(n):
            iobs[i * rb:(i + 1) * rb, iz, :ny, :nx] /= av + v[i]
        if (rb * n) < nproj:
            iobs[rb * n:, iz, :ny, :nx] /= av + v[-1]


def calc_apply_double_flat_field_kw(kwargs):
    iz, nproj, nx, ny, nz = kwargs['iz'], kwargs['nproj'], kwargs['nx'], kwargs['ny'], kwargs['nz']
    freqs, ref, nb_proj_sharray = kwargs['freqs'], kwargs['ref'], kwargs['nb_proj_sharray']
    pady, padx = kwargs['padding']
    nb_workers = kwargs['nb_workers']
    calc_apply_double_flat_field(iz, nproj, nb_proj_sharray, nx, ny, nz, padx, pady, ref, freqs, nb_workers)


def load_holoct_data(i, img_name, binning=None):
    """
    Load edf images (phased projections) resulting from an holoct process.
    :param i: index of the projection
    :param img_name: generic name of the projections, with a field for the index.
    :param binning: binning factor
    :return: the loaded image
    """
    img = fabio.open(img_name % i).data
    if binning is not None:
        if binning > 1:
            img = rebin(img, binning)
    return img


def load_holoct_data_pool_init(ph_array):
    """
    Initialiser for the pool, to pass the shared arrays
    :param ph_array:
    :return:
    """
    shared_arrays["ph_array"] = ph_array


def load_holoct_data_kw(kwargs):
    nproj, nx, ny = kwargs['nproj'], kwargs['nx'], kwargs['ny']
    ph_array = shared_arrays['ph_array']
    ph_array_dtype = kwargs['ph_array_dtype']
    i, idx = kwargs['i'], kwargs['idx']
    ph_shape = (nproj, ny, nx)
    ph = np.frombuffer(ph_array.get_obj(), dtype=ph_array_dtype).reshape(ph_shape)
    d = load_holoct_data(idx, kwargs['img_name'], kwargs['binning'])
    # The data shape may be larger that the original frame size, depending
    # on mcrop/ncrop values
    if d.shape[0] > ny:
        d = d[d.shape[0] // 2 - ny // 2: d.shape[0] // 2 - ny // 2 + ny]
    if d.shape[1] > nx:
        d = d[:, d.shape[1] // 2 - nx // 2: d.shape[1] // 2 - nx // 2 + nx]
    ph[i] = d


def save_phase_edf(idx, ph, prefix_output):
    """
    Save phased projections to edf
    :param idx: index of the projection
    :param ph: array for the image
    :param prefix_output: prefix for the images
    :return: nothing
    """
    edf = fabio.edfimage.EdfImage(data=ph.astype(np.float32))
    edf.write("%s_%04d.edf" % (prefix_output, idx))


def save_phase_edf_kw(kwargs):
    return save_phase_edf(**kwargs)


def get_params_id16b(scan_path, planes=(1,), sx0=None, sx=None, cx=None, pixel_size_detector=None, logger=None):
    """
    Get the pixel size(s), magnified propagation distances and angular range from an ID16B dataset.
    By default, this will use the 'info' file parsed by EDFTomoScan, but if
    sx0 is given, the motor positions will be used instead. If sx or cx are given,
    the given values supersede those from the EDF files.

    :param scan_path: the path to the scan, either "path/to/sample_name_1_" or
        "path/to/sample_name_%d_" for multi-distance datasets.
    :param planes: the list of planes to consider, usually [1] or [1, 2, 3, 4]
    :param sx0: focus offset in mm (optional)
    :param sx: sample motor position in mm (optional). Should be iterable (list, array)
        if len(planes)>1.
    :param cx: detector motor position in mm (optional)
    :param pixel_size_detector: detector pixel size, if already known.
    :param logger: a logger object
    :return: a dictionary with keys including 'wavelength', 'pixel_size', 'distance',
        the latter two being arrays, after taking int account magnification.
    """
    if logger is None:
        logger = logging.getLogger()
    logger.info("Getting ID16B experiment parameters")
    px = []
    d = []
    result = {}
    for iz in planes:
        if re.search('%[0-9]*?d', scan_path) is None:  # Matches %d, %2d, %02d
            scans = scan_path[:-2] + "%d_" % iz
        else:
            scans = scan_path % iz
        s = EDFTomoScan(scans, n_frames=1)
        if iz == planes[0]:
            logger.info(f"Reading parameters (pixel and distances) from: {scans}")
            result['wavelength'] = 12.3984e-10 / EDFTomoScan.retrieve_information(scans, dataset_basename=None,
                                                                                  key="Energy", ref_file=None,
                                                                                  type_=float)
            result['nx'] = s.dim_1
            result['ny'] = s.dim_2
            w = result['wavelength']
            logger.info("Energy: %6.2fkeV   Wavelength: %8.5fnm" % (12.3984e-10 / w, w * 1e9))
            logger.info("Frame size: (%d, %d)" % (result['ny'], result['nx']))

        # Compute the distances using:
        # * motor positions in edf files (cx and sx)
        # * magnification from optic_used (non-magnified pixel size) and pixel_size
        #   also recorded in the edf files
        edfname = "%s/%s0050.edf" % (scans, s.get_dataset_basename())
        edf = silx.io.open(edfname)
        if cx is None:
            cx1 = edf["scan_0/instrument/positioners/cx"][0]
        else:
            cx1 = cx
        if sx is None:
            sx1 = edf["scan_0/instrument/positioners/sx"][0]
        else:
            if len(planes) == 1:
                sx1 = float(sx)
            else:
                sx1 = sx[iz]
        # Magnified pixel detector
        try:
            px_magnified_edf = edf["scan_0/instrument/detector_0/others/pixel_size"][0] / 1e6
        except KeyError:
            logger.info("No pixel_size field available from edf files")
            px_magnified_edf = None
        if pixel_size_detector is None:
            px1 = edf["scan_0/instrument/detector_0/others/optic_used"][0] / 1e6
        else:
            px1 = pixel_size_detector

        if sx0 is None:
            if px_magnified_edf is None:
                raise RuntimeError("Could not compute sx0 from metadata in edf files - missing "
                                   "magnified pixel size. Please supply sx0")
            # Compute sx0 from available fields
            mag = px1 / px_magnified_edf
            sx0 = sx1 - cx1 / mag

        # This should be redundant, we already know the magnification
        z1 = (sx1 - sx0) / 1000  # Focus-sample distance (m)
        z2 = d0 = (cx1 - sx1 + sx0) / 1000  # Sample-detector distance (m)
        magnification = (z1 + z2) / z1
        px.append(px1 / magnification)

        # print(px[-1], px_magnified_edf, px[-1]- px_magnified_edf)
        if px_magnified_edf is not None:
            if not np.isclose(px[-1], px_magnified_edf, rtol=5e-4, atol=0):
                logger.warning("Magnified pixel size (%6.2fnm) calculated from sx, cx, sx0 "
                               "differs from the one in edf file (%6.2fnm)" % (
                                   px[-1] * 1e9, px_magnified_edf * 1e9))
        # else:
        #     d0 = s.distance  # real distance in meters, from the .info Distance
        #     px.append(s.pixel_size)  # Magnified pixel size
        #     # real pixel size (taking into account on-detector optics)
        #     px1 = EDFTomoScan.retrieve_information(scans, dataset_basename=None, key="Optic_used",
        #                                            ref_file=None, type_=float) * 1e-6
        m = px1 / px[-1]  # Magnification
        logger.info(f"[iz={iz}]  Pixel size: {px[-1] * 1e9:8.2f}nm (magnified) {px1 * 1e9:8.2f}nm (detector), "
                    f"Distance: {d0 / m:8.6f}m (magnified) {d0:8.6f}m (real) [sx0={sx0:.3f}mm]")
        d.append(d0 / m)
    result['pixel_size'] = np.array(px)
    result['distance'] = np.array(d)

    # Tomo angular range
    if "%d" not in scan_path:
        scans = scan_path[:-2] + "%d_" % planes[0]
    else:
        scans = scan_path % planes[0]
    scan_range, tomo_n = EDFTomoScan.get_scan_range(scans), EDFTomoScan.get_tomo_n(scans)
    result['scan_range'] = scan_range
    result['tomo_n'] = tomo_n

    return result


def get_params_id16a(scan_path, planes=(1,), logger=None, **kwargs):
    """
    Get the pixel size(s), magnified propagation distances and angular range from an ID16A dataset.
    This will use the 'info' file parsed by EDFTomoScan.

    :param scan_path: the path to the scan, either "path/to/sample_name_1_" or
        "path/to/sample_name_%d_" for multi-distance datasets.
    :param planes: the list of planes to consider, usually [1] or [1, 2, 3, 4]
    :param logger: a logger object
    :return: a dictionary with keys including 'wavelength', 'pixel_size', 'distance',
        the latter two being arrays, after taking int account magnification.
    """
    if logger is None:
        logger = logging.getLogger()
    logger.info("Getting ID16A experiment parameters")
    px = []
    d = []
    result = {}
    for iz in planes:
        if re.search('%[0-9]*?d', scan_path) is None:
            scans = scan_path[:-2] + "%d_" % iz
        else:
            scans = scan_path % iz
        logger.info(f"Reading parameters (pixel and distances) for: {scans}")
        s = EDFTomoScan(scans, n_frames=1)
        if iz == planes[0]:
            result['wavelength'] = 12.3984e-10 / EDFTomoScan.retrieve_information(scans, dataset_basename=None,
                                                                                  key="Energy", ref_file=None,
                                                                                  type_=float)
            result['nx'] = s.dim_1
            result['ny'] = s.dim_2
            w = result['wavelength']
            logger.info("Energy: %6.2fkeV   Wavelength: %8.5fnm" % (12.3984e-10 / w, w * 1e9))
            logger.info("Frame size: (%d, %d)" % (result['ny'], result['nx']))

        # Convention is different from ID16B (magnified distance here)
        d0 = s.distance  # *magnified* distance in meters, from the .info Distance field (?)
        # Tomoscan assumes the distance is in mm
        if d0 < 0.001:
            logger.info("ID16A distance<.001 from EDFTomoScan.retrieve_information, apparently "
                        "it was already in meters so x1000 and hope for the best")
            d0 *= 1000
        px.append(s.pixel_size)  # Magnified pixel size
        # real pixel size (taking into account on-detector optics)
        px1 = EDFTomoScan.retrieve_information(scans, dataset_basename=None, key="Optic_used",
                                               ref_file=None, type_=float) * 1e-6
        m = px1 / px[-1]  # Magnification
        logger.info("  Pixel size: %8.2fnm (magnified) %8.2fnm (detector), "
                    "Distance: %8.6fm (magnified) %8.6fm (real)" % (px[-1] * 1e9, px1 * 1e9, d0, d0 * m))
        d.append(d0)
    result['pixel_size'] = np.array(px)
    result['distance'] = np.array(d)

    # Tomo angular range
    if "%d" not in scan_path:
        scans = scan_path[:-2] + "%d_" % planes[0]
    else:
        scans = scan_path % planes[0]
    scan_range, tomo_n = EDFTomoScan.get_scan_range(scans), EDFTomoScan.get_tomo_n(scans)
    result['scan_range'] = scan_range
    result['tomo_n'] = tomo_n

    return result


def sino_filter_pad(d, filter_name, padding):
    """
    Perform a 1D filtering along the x-axis, prior to filtered back-projection.
    This is only useful for padded arrays, in order to use the array values
    in the padded areas before cropping for the back-projection.
    In order to have the same normalisation as in nabu, the resulting array
    will need to be multiplied by pi/nb_proj.

    :param d: the array to filter, can have any dimensions, filtering will be
        done along the last (fastest) dimension
    :param filter_name: Available filters (from silx): Ram-Lak,
        Shepp-Logan, Cosine, Hamming, Hann, Tukey, Lanczos. Case-insensitive.
    :param padding: the number of padding pixels on each side of the last dimension
    :return: the filtered array, with the original array shape
    """
    nxd = d.shape[-1]
    # nx is the non-padded range. Filtering is done by padding to 2*nx
    nx = nxd - 2 * padding
    filter_f = compute_fourier_filter(2 * nx, filter_name=filter_name)[:nx + 1]  # R2C filter
    # filter_f *= np.pi / nproj  # normalisation as in nabu. Needs to be done elsewhere
    if nxd == 2 * nx:
        d1 = d
    else:
        # Need to pad d to 2*nx
        pad1 = (2 * nx - nxd) // 2
        vpad = [(0, 0) for i in range(d.ndim - 1)] + [(pad1, pad1)]
        d1 = np.pad(d, vpad, mode='edge')  # edge or reflect ?
    ixleft = (d1.shape[-1] - nxd) // 2
    return irfftn(rfftn(d1, axes=(-1,)) * filter_f, axes=(-1,))[..., ixleft:ixleft + nxd]
