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

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

# TODO: This is a tentative multiprocessing-based implementation of the holotomo script.
#  Still needs:
#  - a LOT of cleanup and OO organisation [https://xkcd.com/2054/]
#  - use a parallel process to monitor CPU, I/O and GPU usage !
#  - handle parallel process failures
#  - etc...

"""
Script for single or multi-distance holo-tomography reconstruction,
tuned for the ESRF ID16A and B beamlines.
"""
import timeit
import time
import sys
import os
import re
import platform
import argparse
import logging
import traceback
import warnings
from io import StringIO, TextIOBase
import gc
from multiprocessing import Pool, Process, Queue, JoinableQueue, Array, active_children
import ctypes
import psutil
import numpy as np
import h5py
from scipy.ndimage import median_filter, gaussian_filter
import fabio
from tifffile import imwrite
from ...utils.array import rebin, pad2
from ...utils.math import primes, test_smaller_primes
from ...utils.string import longest_common_prefix, replace_dir
from pynx.version import get_git_version

_pynx_version = get_git_version()

try:
    # Get the real number of processor cores available
    # os.sched_getaffinity is only available on some *nix platforms
    nproc = len(os.sched_getaffinity(0)) * psutil.cpu_count(logical=False) // psutil.cpu_count(logical=True)
except AttributeError:
    nproc = os.cpu_count()


def _get_gpu_total_mem():
    import pycuda.autoprimaryctx
    import pycuda.driver as drv
    return drv.mem_get_info()[-1] / 1024 ** 3


def get_gpu_total_mem():
    # Get the available GPU memory in a separate process
    # to avoid initialising cuda in the main process
    with Pool(1) as pool:
        return pool.apply(_get_gpu_total_mem)


def _get_vkfft_buffer_nbytes(ny, nx):
    import pycuda.autoprimaryctx
    from pyvkfft.cuda import VkFFTApp
    import numpy as np
    return VkFFTApp((ny, nx), dtype=np.complex64).get_tmp_buffer_nbytes()


def get_vkfft_buffer_nbytes(ny, nx):
    """
    Return the size of the temporary buffer required by VkFFT, if any
    :param ny:
    :param nx:
    :return: the buffer size
    """
    with Pool(1) as pool:
        return pool.apply(_get_vkfft_buffer_nbytes, (ny, nx))


class DeltaTimeFormatter(logging.Formatter):
    start_time = time.time()

    def format(self, record):
        record.delta = "%+8.2fs" % (time.time() - self.start_time)
        return super().format(record)


def new_logger_stream(logger_name: str, level: int = logging.INFO, flush_to=None, stdout=True):
    """Create / access a logger with a new StringIO output, allowing
    to share log between multiple processes using stream.getvalue().
    :param logger_name: the logger name
    :param level: the logging level
    :param flush_to: an opened file object to fluch the previous stream to, if any
    :param stdout: if True, also output log to stdout
    :return: a tuple (logger, stream)
    """
    logger = logging.getLogger(logger_name)
    if isinstance(flush_to, TextIOBase) and len(logger.handlers):
        h = logger.handlers[-1]
        h.flush()
        flush_to.write(h.stream.getvalue())
        flush_to.flush()
    logger.handlers.clear()
    LOGFORMAT = '%(asctime)s - %(delta)s - %(levelname)8s - %(name)16s - %(message)s'
    fmt = DeltaTimeFormatter(LOGFORMAT)

    if stdout:
        h = logging.StreamHandler(sys.stdout)
        h.setFormatter(fmt)
        logger.addHandler(h)
    logger.setLevel(level)
    stream = StringIO()
    handler = logging.StreamHandler(stream)
    handler.setFormatter(fmt)
    logger.addHandler(handler)
    return logger, stream


def load_align_zoom(ichunk, data_src, prefix, proj_idx, planes, binning, magnification,
                    padding, pad_method, rhapp, reference_plane, spikes_threshold,
                    align_method, align_interval, align_fourier_low_cutoff,
                    align_fourier_high_cutoff, prefix_output, nb_chunk,
                    motion, logger, iobs_array, proj_urls, distorted, normalise,
                    align_diff, align_max_shift, vsnr, params=None):
    import traceback
    from scipy.interpolate import interp1d, LSQUnivariateSpline
    from tomoscan.esrf.scan.edfscan import EDFTomoScan
    from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
    from pynx.holotomo.utils import load_data_pool_init, load_data_kw, zoom_pad_images, zoom_pad_images_init, \
        zoom_pad_images_kw, align_images_kw, align_images_pool_init, remove_spikes, calc_apply_double_flat_field_kw, \
        shift_images_kw, vsnr_kw, vsnr_init
    nz = len(planes)
    t0 = timeit.default_timer()
    logger.info("#" * 48)
    logger.info("# Loading data & align / zoom images for chunk %d/%d:" % (ichunk + 1, nb_chunk))
    logger.info("# [%d,%d...%d] - %d projections" % (proj_idx[0], proj_idx[1],
                                                     proj_idx[-1], len(proj_idx)))
    logger.info("#" * 48)
    dark = None
    if data_src.endswith('.nx'):
        for iz in range(nz):
            nx_name = data_src % (planes[iz] - 1)
            logger.info(f"Loading dark from: {nx_name}")
            s = NXtomoScan(nx_name)
            for k, v in s.darks.items():
                with h5py.File(v.file_path()) as h:
                    d = h[v.data_path()][v.data_slice()]
                if binning is not None:
                    if binning > 1:
                        d = rebin(d, binning)
                if dark is None:
                    ny, nx = d.shape
                    dark = np.zeros((nz, ny, nx), dtype=np.float32)
                dark[iz] += d
            dark[iz] /= len(s.darks)
            v5, v95 = np.percentile(dark[iz], (5, 95))
            logger.info(f"Dark[{iz}]: <dark>={dark[iz].mean():.1f} 5-95% = {v5:.1f}-{v95:.1f}]")
    else:
        for iz in range(nz):
            dark_url = EDFTomoScan.get_darks_url(data_src % planes[iz])
            if len(dark_url) == 0:
                # Try other names, if merging was not yet done
                dark_url = EDFTomoScan.get_darks_url(data_src % planes[iz], prefix='darkend')
                has_merged_dark = False
                logger.warning(f"No dark found with default location for plane {planes[iz]}"
                               f" - trying unmerged darkend*.edf")
            else:
                has_merged_dark = True
            if len(dark_url):
                vurls = [v.file_path() for v in dark_url.values()]
                logger.info(f"Loading {len(dark_url)} darks (name/prefix={longest_common_prefix(vurls)})")
                for url in vurls:
                    a = fabio.open(url)
                    d = a.data.astype(np.float32)
                    if not has_merged_dark:
                        # TODO: check this approach is correct for unmerged darks ??
                        h = a.getheader()
                        if "count_time" in h and "acq_expo_time" in h and "acq_mode" in h:
                            try:
                                ct = float(h['count_time'])
                                expo = float(h['acq_expo_time'])
                                if "accum" in h["acq_mode"].lower() and ct > 0 and expo > 0:
                                    logger.warn(f"Unmerged dark with count_time={ct:.3f} "
                                                f"acq_expo_time={expo:.3f} => "
                                                f"Dividing read dark by {expo / ct:.2f} ")
                                    d *= ct / expo
                            except Exception:
                                logger.warning("Tried to fix dark scale from expo_time and count_time"
                                               "but failed - hope no scale is needed.")
                                print(traceback.format_exc())

                    if binning is not None:
                        if binning > 1:
                            d = rebin(d, binning)
                    if dark is None:
                        ny, nx = d.shape
                        dark = np.zeros((nz, ny, nx), dtype=np.float32)
                    dark[iz] += d
                dark[iz] /= len(vurls)
            else:
                if iz == 0:
                    logger.error(f"No dark found for first plane {planes[iz]} - giving up")
                    raise RuntimeError(f"No dark found for first plane {planes[iz]}")
                else:
                    logger.warning(f"No dark found for plane {planes[iz]} - using previous dark")
                    dark[iz] = dark[iz - 1]
            v5, v95 = np.percentile(dark[iz], (5, 95))
            logger.info(f"Dark[{iz}]: <dark>={dark[iz].mean():.1f} 5-95% = {v5:.1f}-{v95:.1f}]")

    ny, nx = dark.shape[-2:]
    logger.info("Frame size: %d x %d" % (ny, nx))

    nb_proj = len(proj_idx)  # Number of projections to load
    pady, padx = padding
    nb_proj_sharray = len(iobs_array) // (nz * (ny + 2 * pady) * (nx + 2 * padx))  # Actual len of shared array

    ref = np.zeros_like(dark)
    if data_src.endswith('.nx'):
        for iz in range(nz):
            nx_name = data_src % (planes[iz] - 1)
            s = NXtomoScan(nx_name)
            vv = list(s.flats.values())
            for v in vv:  # TODO: take into accound beginning/end flats ?
                with h5py.File(v.file_path()) as h:
                    d = h[v.data_path()][v.data_slice()]
                if binning is not None:
                    if binning > 1:
                        d = rebin(d, binning)
                ref[iz] += d
            ref[iz] /= len(vv)
            ref[iz] -= dark[iz]
            v5, v95 = np.percentile(ref[iz], (5, 95))
            logger.info(f"Ref[{iz}]: <ref>={ref[iz].mean():.1f} 5-95% = {v5:.1f}-{v95:.1f} [dark subtracted]")
    else:
        for iz in range(nz):
            flats_url = EDFTomoScan.get_flats_url(data_src % planes[iz])
            if len(flats_url) == 0:
                flats_url = EDFTomoScan.get_flats_url(data_src % planes[iz], prefix='ref')
                if len(flats_url) == 0:
                    logger.error("Did not find any flat/empty beam reference frame")
                    raise RuntimeError(f"No flat/empty beam reference frame found for plane {planes[iz]}")
            # If there are multiple references at begin/end, these should have keys like 0.0 0.0001, 0.0002,
            # 2000.000, 2000.0001, 2000.0002, where the fractional parts indicate the order of the frame
            # and the integer when the image was taken (at beginning or end of scan)
            # Reduce these, so we can separate begin/end flats
            vflats = {}
            for k, v in flats_url.items():
                if int(k) in vflats:
                    vflats[int(k)].append(v.file_path())
                else:
                    vflats[int(k)] = [v.file_path()]
            # For now, use the first series of flats
            # TODO: load multiple empty beam images (begin/end)
            k0 = min(vflats.keys())
            logger.info(f"Loading {len(vflats[k0])} empty reference image(s) [iz={planes[iz]}]:"
                        f"{longest_common_prefix(vflats[k0])}")
            for v in vflats[k0]:
                d = fabio.open(v).data
                if binning is not None:
                    if binning > 1:
                        d = rebin(d, binning)
                ref[iz] += d.astype(np.float32)
            ref[iz] /= len(vflats[k0])
            ref[iz] -= dark[iz]
            v5, v95 = np.percentile(ref[iz], (5, 95))
            logger.info(f"Ref[{iz}]: <ref>={ref[iz].mean():.1f} 5-95% = {v5:.1f}-{v95:.1f} [dark subtracted]")

    if spikes_threshold and not normalise:
        # Also remove spikes on reference/empty beam images - Unless we use normalised
        # images, in which case the spike correction is done after normalisation.
        for r in ref:
            remove_spikes(r, threshold=spikes_threshold)

    # Check we don't have zero (or close to) valued pixels in reference
    for iz in range(nz):
        av = ref[iz].mean()
        tmp = ref[iz] < (av / 1000)
        nb0 = tmp.sum()
        if nb0:
            logger.warn(f"Empty reference image (iz={iz}): found {nb0} pixels below "
                        f"<Iref-dark>/1000 - setting to {av / 1000:.2f}")
            ref[iz][tmp] = av / 1000  # :KLUDGE ?

    # Prepare shared arrays for loading & zoom/pad
    dark_old = dark
    dark_array = Array(ctypes.c_float, dark.size)
    dark = np.frombuffer(dark_array.get_obj(), dtype=ctypes.c_float).reshape(dark_old.shape)
    dark[:] = dark_old
    del dark_old
    ref_old = ref
    ref_array = Array(ctypes.c_float, ref.size)
    ref = np.frombuffer(ref_array.get_obj(), dtype=ctypes.c_float).reshape(ref_old.shape)
    ref[:] = ref_old
    del ref_old
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape((nb_proj_sharray, nz,
                                                                              ny + 2 * pady, nx + 2 * padx))
    iobs = iobs[:nb_proj]

    logger.info("Loading projections images: [%d processes]" % nproc)
    vkw = [{'i': i, 'proj_urlz': [proj_urls[p][proj_idx[i]] for p in planes],
            'nproj': nb_proj_sharray, 'nx': nx, 'ny': ny, 'planes': planes,
            'padding': padding, 'binning': binning, 'spikes_threshold': spikes_threshold,
            'distorted': distorted, 'normalise': normalise}
           for i in range(len(proj_idx))]
    timings = {'load-bin-dark': 0, 'distort': 0, 'normalise': 0}
    with Pool(nproc, initializer=load_data_pool_init, initargs=(iobs_array, dark_array, ref_array)) as pool:
        res = pool.map(load_data_kw, vkw)  # , chunksize=1
        distortion_shifts = []
        for r in res:
            if distorted:
                dis, dt = r
                distortion_shifts.append(dis)
            else:
                dt = r
            for k, v in dt.items():
                timings[k] += v / nproc

    dt = timeit.default_timer() - t0
    nbytes = nz * nx * ny * (nb_proj + 2) * 2 * binning ** 2  # number of bytes actually read from disk (uint16)
    if distorted:
        s = " & distort"
    else:
        s = ""
    if normalise:
        s += " & normalise"
        for iz in range(nz):
            ref[iz] = ref[iz].mean()  # For statistics purposes
    logger.info(f"Time to load & uncompress{s} data: {dt:4.1f}s")
    mbps = nbytes / timings['load-bin-dark'] / 1024 ** 2
    logger.info(f"         Load/bin/dark correction: {timings['load-bin-dark']:4.1f}s [{mbps:6.2f} Mbytes/s]")
    if distorted:
        logger.info(f"                   Distortion: {timings['distort']:4.1f}s")

    if distorted:
        s = prefix_output + '_shifts.h5'
        logger.info(f"Saving distortion to: {s}")
        with h5py.File(s, 'a') as h:
            # Save parameters used for the script.
            if 'parameters' not in h and params is not None:
                config = h.create_group("parameters")
                config.attrs['NX_class'] = 'NXcollection'
                for k, v in vars(params).items():
                    if v is not None:
                        if type(v) is dict:
                            # This can happen if complex configuration is passed on
                            if len(v):
                                kd = config.create_group(k)
                                kd.attrs['NX_class'] = 'NXcollection'
                                for kk, vv in v.items():
                                    kd.create_dataset(kk, data=vv)
                        else:
                            config.create_dataset(k, data=v)
                # Also save magnification
                config.create_dataset("magnification", data=magnification)

            proj_idx1 = np.array(proj_idx, dtype=int)
            distortion_shifts = np.array(distortion_shifts)
            if 'distortion' not in h:
                entry_distortion = h.create_group("distortion")
                entry_distortion.attrs['NX_class'] = 'NXdata'
                entry_distortion.attrs['signal'] = 'data'
                entry_distortion.attrs['interpretation'] = 'image'
                h.attrs['default'] = 'distortion'  # This will be overwritten for shifts present
            else:
                entry_distortion = h['distortion']
            if 'data' in entry_distortion:
                distortion_shifts = np.append(distortion_shifts, entry_distortion['data'][()], axis=0)
                proj_idx1 = np.append(proj_idx1, entry_distortion['idx'][()])
                idx = np.argsort(proj_idx1)
                proj_idx1 = proj_idx1[idx]
                distortion_shifts = np.take(distortion_shifts, idx, axis=0)
                del entry_distortion['data']
                del entry_distortion['idx']
            entry_distortion.create_dataset('data', data=distortion_shifts, chunks=True, compression="gzip")
            entry_distortion.create_dataset('idx', data=proj_idx1)

    ################################################################
    # VSNR filter ?
    ################################################################
    if vsnr is not None:
        t1 = timeit.default_timer()
        filter_params = [{'qr0': vsnr[4 * i], 'sigma_r': vsnr[4 * i + 1], 'sigma_t': vsnr[4 * i + 2],
                          'theta': np.deg2rad(vsnr[4 * i + 3])} for i in range(len(vsnr) // 4)]

        msg = "Applying VSNR filtering: "
        for p in filter_params:
            msg += "("
            for k, v in p.items():
                msg += f"{k}={v:.3f}  " if k != 'theta' else f"{k}={np.rad2deg(v):.1f}°"
            msg += ") "
        msg += f" beta={vsnr[-2]} niter={vsnr[-1]}"
        logger.info(msg)
        roi = params.vsnr_roi
        if binning > 1:
            roi = [v // binning for v in roi]
        vkw = [{'i': i, 'nb_proj_sharray': nb_proj_sharray, 'nx': nx, 'ny': ny, 'nz': nz,
                'padding': padding, 'idx': range(0, nb_proj, params.ngpu),
                'filter_params': filter_params, 'beta': vsnr[-2], 'do_ref': not normalise,
                'niter': vsnr[-1], 'roi': roi,
                'igpu': i} for i in range(params.ngpu)]

        with Pool(nproc, initializer=vsnr_init, initargs=(iobs_array, ref_array)) as pool:
            pool.map(vsnr_kw, vkw)
        dt = timeit.default_timer() - t1
        logger.info(f"Finished VSNR filtering [dt={dt:.2f}s]")

    ################################################################
    # Apply double flat-field ?
    ################################################################
    if params.double_flat_field is not None:
        if params.double_flat_field[1] >= 1e-3:
            logger.info(f"Applying double flat-field correction with 2D+1D filter "
                        f"(in-plane high-pass f>{params.double_flat_field[0]} and "
                        f"low-pass f<{params.double_flat_field[1]} along projections)")
        elif params.double_flat_field[0] >= 1e-3:
            logger.info(f"Applying double flat-field correction with high-pass "
                        f"in-plane filter (f>{params.double_flat_field[0]})")
        else:
            logger.info("Applying double flat-field correction")
        if True:
            vkw = [{'iz': iz, 'nproj': nb_proj, 'nb_proj_sharray': nb_proj_sharray, 'nx': nx, 'ny': ny, 'nz': nz,
                    'padding': padding, 'freqs': params.double_flat_field, 'ref': ref[iz],
                    'nb_workers': nproc // nz} for iz in range(nz)]
            with Pool(nz, initializer=load_data_pool_init, initargs=(iobs_array, dark_array, ref_array)) as pool:
                pool.map(calc_apply_double_flat_field_kw, vkw)
        logger.info("Finished applying double-flat-field")

    ################################################################
    # Zoom & register images, keep first distance pixel size and size
    # Do this using multiple process to speedup
    ################################################################
    if nz > 1:
        logger.info("Magnification relative to iz=0: %s" % str(magnification / magnification[0]))
        t0 = timeit.default_timer()
        ref_zoom = zoom_pad_images(ref, magnification, padding, nz)

        vkw = [{'magnification': magnification, 'padding': padding, 'nb_proj_sharray': nb_proj_sharray,
                'pad_method': pad_method, 'nx': nx, 'ny': ny, 'nz': nz, 'i': i} for i in range(nb_proj)]
        with Pool(nproc, initializer=zoom_pad_images_init, initargs=(iobs_array,)) as pool:
            pool.map(zoom_pad_images_kw, vkw)

        logger.info("Zoom & pad images: dt = %6.2fs [padding: (%d,%d), %s]" %
                    (timeit.default_timer() - t0, pady, padx, pad_method))

        # Align images
        t0 = timeit.default_timer()
        if rhapp is not None:
            logger.info("Aligning images: loading shifts from holoCT (rhapp)")
            nb = np.loadtxt(rhapp, skiprows=4, max_rows=1, dtype=int)[2]
            m = np.loadtxt(rhapp, skiprows=5, max_rows=nb * 8, dtype=np.float32).reshape((nb, 4, 2))
            dx_holoct = m[..., 1] / binning
            dy_holoct = m[..., 0] / binning

            dx_holoct = np.take(dx_holoct, proj_idx, axis=0)
            dy_holoct = np.take(dy_holoct, proj_idx, axis=0)
        else:
            dx_holoct, dy_holoct = None, None

        # Take into account random motion displacements if available ?
        # When preset, the smooth displacements can be obtained
        # by comparing corrected displacements dmx[iz]=(dx+mx)[iz] between distances.
        # Once we have dmx_fit[iz], the real positions to use
        # are dmx_fit[iz] - mx[iz]
        mx = np.zeros((nb_proj, nz), dtype=np.float32)
        my = np.zeros((nb_proj, nz), dtype=np.float32)
        izref = list(planes).index(reference_plane)

        if motion:
            logger.info(f"Loading motion correction files: {data_src}/correct.txt")
            for iz in range(nz):
                m = np.loadtxt(f"{data_src % planes[iz]}/correct.txt", dtype=np.float32)
                m *= magnification[0] / magnification[iz]
                my[:, iz] = m[:, 1][proj_idx] / binning
                mx[:, iz] = m[:, 0][proj_idx] / binning
            # Look for a specific motion correction for the reference plane (exploiting the quali series)
            if os.path.isfile(f"{data_src % reference_plane}/correct_motion.txt"):
                logger.info(f"Using extra drift correction from: {data_src % reference_plane}"
                            f"/correct_motion.txt from reference plane")
                m_ref = np.loadtxt(f"{data_src % reference_plane}/correct_motion.txt", dtype=np.float32)
                m_ref *= magnification[0] / magnification[izref]
                m_ref_y = m_ref[:, 1][proj_idx] / binning - my[:, izref]
                m_ref_x = m_ref[:, 0][proj_idx] / binning - mx[:, izref]
                # The shift is applied to all planes since this corrects for the reference plane drift,
                # and we do not want to change the relative shifts between planes
                for iz in range(nz):
                    my[:, iz] += m_ref_y
                    mx[:, iz] += m_ref_x

        if align_method != 'rhapp':
            logger.info(f"Aligning images, method: {align_method}, diff={align_diff}")
            if align_interval > 1:
                idx = list(range(0, nb_proj, align_interval))
                proj_idx_interval = list(proj_idx[::align_interval])
                if proj_idx_interval[-1] != proj_idx[-1]:
                    proj_idx_interval.append(proj_idx[-1])
                    idx.append(nb_proj - 1)
                proj_idx_interval = np.array(proj_idx_interval)
            else:
                idx = range(nb_proj)
                proj_idx_interval = proj_idx
            # This can sometimes (<1 in 10) fail (hang indefinitely). Why ?
            # res = pool.map(align_images, range(nb_proj))
            ref_old = ref_zoom
            ref_zoom_array = Array(ctypes.c_float, ref_zoom.size)
            ref_zoom = np.frombuffer(ref_zoom_array.get_obj(), dtype=ctypes.c_float).reshape(ref_zoom.shape)
            ref_zoom[:] = ref_old
            del ref_old

            low_width = 0.03
            if align_fourier_low_cutoff is not None:
                low_width = align_fourier_low_cutoff if align_fourier_low_cutoff < 0.03 else 0.03
            high_width = 0.03
            if align_fourier_high_cutoff is not None:
                high_width = 0.5 - align_fourier_high_cutoff if align_fourier_high_cutoff > 0.47 else 0.03
            vkw = [{'i': i, 'nproj': nb_proj_sharray, 'nz': nz, 'ny': ny, 'nx': nx, 'padding': padding,
                    'low_cutoff': align_fourier_low_cutoff,
                    'low_width': low_width,
                    'high_cutoff': align_fourier_high_cutoff,
                    'high_width': high_width,
                    'normalise': not normalise,
                    'diff': align_diff,
                    'align_max_shift': align_max_shift // binning} for i in idx]
            align_ok = False
            nb_nok = 0
            while not align_ok:
                if nb_nok >= 4:
                    logger.critical("Alignment: 4 alignment failures, bailing out")
                    raise RuntimeError("Alignment: 4 alignment failures, bailing out")
                try:
                    res = []
                    with Pool(nproc, initializer=align_images_pool_init,
                              initargs=(iobs_array, ref_zoom_array)) as pool:
                        results = pool.imap(align_images_kw, vkw, chunksize=1)
                        for i in range(len(vkw)):
                            r = results.next(timeout=20)
                            res.append(r)
                    align_ok = True
                    logger.info("align OK")
                except Exception as ex:
                    import traceback
                    print(traceback.format_exc())
                    logger.info("Timeout, re-trying")
                    nb_nok += 1

            # print(res)

            dx = np.zeros((len(vkw), nz), dtype=np.float32)
            dy = np.zeros((len(vkw), nz), dtype=np.float32)
            for i in range(len(vkw)):
                dx[i] = res[i][0]
                dy[i] = res[i][1]

            # dx_debug = dx.copy()
            # dy_debug = dy.copy()

            # Correct for the random motion (if any) to get smooth displacements vs the projection number.
            # The random motion will be added back at the end.
            dx += np.take(mx, idx, axis=0)
            dy += np.take(my, idx, axis=0)

            # Replace first and last values which are less good (differential approach)
            dx[0] = dx[1]
            dx[-1] = dx[-2]
            dy[0] = dy[1]
            dy[-1] = dy[-2]

            # Switch reference plane
            for iz in range(nz):
                if iz != izref:
                    dx[:, iz] -= dx[:, izref]
                    dy[:, iz] -= dy[:, izref]
            dx[:, izref] = 0
            dy[:, izref] = 0

            # Keep copy of raw positions before fit/filtering
            dxraw = dx.copy()
            dyraw = dy.copy()

            if 'poly' in align_method:
                # Use polyfit to smooth shifts
                # First use a median filter to remove outliers
                dx = median_filter(dx, (3, 1), mode='nearest')
                dy = median_filter(dy, (3, 1), mode='nearest')
                k = int(align_method[-1])
                logger.info(f"Alignment fitting using a polynomial of order {k}")
                for iz in range(nz):
                    if iz != izref:
                        polx = np.polynomial.polynomial.polyfit(proj_idx_interval, dx[:, iz], k)
                        poly = np.polynomial.polynomial.polyfit(proj_idx_interval, dy[:, iz], k)
                        dx[:, iz] = np.polynomial.polynomial.polyval(proj_idx_interval, polx)
                        dy[:, iz] = np.polynomial.polynomial.polyval(proj_idx_interval, poly)
            elif 'spline' in align_method:
                # Use a 3rd-order Univariate spline to smooth shifts
                # First use a median filter to remove outliers
                dx = median_filter(dx, (3, 1), mode='nearest')
                dy = median_filter(dy, (3, 1), mode='nearest')
                k = int(align_method[-1])
                logger.info(f"Alignment fitting using an univariate spline of order 3 and {k} knots")
                for iz in range(nz):
                    if iz != izref:
                        # We use np.linspace(i0,i1, N)[1:-1] because an LSQUnivariateSpline always adds
                        # a knot at the beginning and the end. That way there are effectively N
                        # knots, evenly spaced.
                        csx = LSQUnivariateSpline(proj_idx_interval, dx[:, iz],
                                                  np.linspace(proj_idx_interval[0],
                                                              proj_idx_interval[-1], k)[1:-1],
                                                  k=3)
                        csy = LSQUnivariateSpline(proj_idx_interval, dy[:, iz],
                                                  np.linspace(proj_idx_interval[0],
                                                              proj_idx_interval[-1], k)[1:-1],
                                                  k=3)

                        dx[:, iz] = csx(proj_idx_interval)
                        dy[:, iz] = csy(proj_idx_interval)
            elif 'median' in align_method:
                k = int(align_method[-1])
                logger.info(f"Alignment filtering using a median filter of size {k}")
                dx = median_filter(dx, (k, 1), mode='nearest')
                dy = median_filter(dy, (k, 1), mode='nearest')
            else:
                logger.info(f"Alignment without fitting or filtering")

            logger.info(f"Alignment standard deviation: dx[iz]="
                        f"{' '.join([f'{np.std(dx[:, iz] - dxraw[:, iz]):5.2f}' for iz in range(nz)])}")
            logger.info(f"Alignment standard deviation: dy[iz]="
                        f"{' '.join([f'{np.std(dy[:, iz] - dyraw[:, iz]):5.2f}' for iz in range(nz)])}")

            # interpolate missing points
            if align_interval > 1:
                dx0, dy0 = dx, dy
                dx = np.empty((nb_proj, nz), dtype=np.float32)
                dy = np.empty((nb_proj, nz), dtype=np.float32)
                for iz in range(nz):
                    dx[:, iz] = interp1d(proj_idx_interval, dx0[:, iz], kind='quadratic')(proj_idx)
                    dy[:, iz] = interp1d(proj_idx_interval, dy0[:, iz], kind='quadratic')(proj_idx)

            # Save alignments to an hdf5 file
            # Assume this is sequential, so no access conflict
            # TODO: use a separate worker to avoid any conflict
            with h5py.File(prefix_output + '_shifts.h5', 'a') as h:
                # Save parameters used for the script.
                if 'parameters' not in h and params is not None:
                    config = h.create_group("parameters")
                    config.attrs['NX_class'] = 'NXcollection'
                    for k, v in vars(params).items():
                        if v is not None:
                            if type(v) is dict:
                                # This can happen if complex configuration is passed on
                                if len(v):
                                    kd = config.create_group(k)
                                    kd.attrs['NX_class'] = 'NXcollection'
                                    for kk, vv in v.items():
                                        kd.create_dataset(kk, data=vv)
                            else:
                                config.create_dataset(k, data=v)
                    # Also save magnification
                    config.create_dataset("magnification", data=magnification)

                if 'shifts' not in h:
                    entry = h.create_group("shifts")
                    entry.attrs['NX_class'] = 'NXdata'
                    entry.attrs['signal'] = 'dx'
                    entry.attrs['interpretation'] = 'spectrum'
                    h.attrs['default'] = 'shifts'
                else:
                    entry = h['shifts']

                proj_idx1 = np.array(proj_idx, dtype=int)
                dx1 = dx.copy()
                dy1 = dy.copy()
                for iz in range(1, nz):
                    # Always save the displacements relative to the first distance,
                    # so we can see if the curves looks nice (after motion correction)
                    dx1[:, iz] -= dx1[:, 0]
                    dy1[:, iz] -= dy1[:, 0]
                    dxraw[:, iz] -= dxraw[:, 0]
                    dyraw[:, iz] -= dyraw[:, 0]
                if motion:
                    mx1 = mx.copy()
                    my1 = my.copy()
                if 'dx' in entry:
                    dx1 = np.append(dx1, entry['dx'][()].transpose(), axis=0)
                    dy1 = np.append(dy1, entry['dy'][()].transpose(), axis=0)
                    proj_idx1 = np.append(proj_idx1, entry['idx'][()])
                    idx = np.argsort(proj_idx1)
                    dx1 = np.take(dx1, idx, axis=0)
                    dy1 = np.take(dy1, idx, axis=0)
                    proj_idx1 = proj_idx1[idx]
                    del entry['dx'], entry['dy'], entry['idx']

                    if motion:
                        mx1 = np.append(mx1, entry['mx'][()].transpose(), axis=0)
                        my1 = np.append(my1, entry['my'][()].transpose(), axis=0)
                        mx1 = np.take(mx1, idx, axis=0)
                        my1 = np.take(my1, idx, axis=0)
                        del entry['mx'], entry['my']

                    dxraw = np.append(dxraw, entry['dx_raw'][()].transpose(), axis=0)
                    dyraw = np.append(dyraw, entry['dy_raw'][()].transpose(), axis=0)
                    proj_idxraw = np.append(proj_idx_interval, entry['idx_raw'][()])
                    idx = np.argsort(proj_idxraw)
                    dxraw = np.take(dxraw, idx, axis=0)
                    dyraw = np.take(dyraw, idx, axis=0)
                    proj_idxraw = proj_idxraw[idx]

                    # dx_debug = np.append(dx_debug, entry['dx_debug'][()].transpose(), axis=0)
                    # dy_debug = np.append(dy_debug, entry['dy_debug'][()].transpose(), axis=0)
                    # dx_debug = np.take(dx_debug, idx, axis=0)
                    # dy_debug = np.take(dy_debug, idx, axis=0)

                    del entry['dx_raw'], entry['dy_raw'], entry['idx_raw']
                    # del entry['dx_debug'], entry['dy_debug']
                else:
                    proj_idxraw = proj_idx_interval
                entry.create_dataset('dx', data=dx1.transpose())
                entry.create_dataset('dy', data=dy1.transpose())
                entry.create_dataset('idx', data=proj_idx1)
                if motion:
                    entry.create_dataset('mx', data=mx1.transpose())
                    entry.create_dataset('my', data=my1.transpose())
                entry.create_dataset('dx_raw', data=dxraw.transpose())
                entry.create_dataset('dy_raw', data=dyraw.transpose())
                entry.create_dataset('idx_raw', data=proj_idxraw)
                # # For debugging: save the measured displacement, without correcting
                # # for motion or changing the reference plane
                # entry.create_dataset('dx_debug', data=dx_debug.transpose())
                # entry.create_dataset('dy_debug', data=dy_debug.transpose())
        else:
            logger.info("Aligning images: using shifts imported from holoCT (rhapp)")
            if dx_holoct is None:
                raise RuntimeError("Image alignment: method is rhapp but align_rhapp was not supplied")
            dx = dx_holoct
            dy = dy_holoct

        # Now the fit/filtering & export is done, add back the random motion
        # (Also when using the rhapp files, which have the smooth displacements)
        dx -= mx
        dy -= my

        if normalise:
            logger.info("Alignment with --normalise: pre-shifting all images (disables probe optimisation)")
            vkw = [{'i': i, 'nproj': nb_proj_sharray, 'nz': nz, 'ny': ny, 'nx': nx,
                    'padding': padding, 'dx': dx[i], 'dy': dy[i]} for i in range(nb_proj)]
            with Pool(nproc, initializer=align_images_pool_init, initargs=(iobs_array,)) as pool:
                pool.map(shift_images_kw, vkw, chunksize=1)
            dx[:] = 0
            dy[:] = 0

        logger.info("Align images: dt = %6.2fs" % (timeit.default_timer() - t0))
    else:
        if padding:
            # Move data to centre of frame & pad
            for i in range(len(iobs)):
                iobs[i, 0] = pad2(iobs[i, 0, :ny, :nx], pad_width=padding, mode=pad_method, mask2neg=True)
            ref_zoom = np.empty((nz, ny + 2 * pady, nx + 2 * padx), dtype=np.float32)
            ref_zoom[0] = pad2(ref[0], pad_width=padding, mode=pad_method, mask2neg=True)
        else:
            ref_zoom = ref
        dx, dy = None, None
    logger.info("Returning Iobs and reference frames")
    # # DEBUG
    # with h5py.File(f'iobs_align_vsnr_{proj_idx[0]:02d}.h5', 'w') as h:
    #     h.create_dataset("data", data=iobs)
    #     h.create_dataset("ref", data=ref)
    #     h.create_dataset("ref_zoom", data=ref_zoom)

    return ref_zoom.copy(), dx, dy


def find_cor_iobs_window(data_src, reference_plane, motion, side, logger,
                         q: Queue, reference_plane_magnification, tomo_rot_center_near):
    # from tomoscan.esrf.scan.edfscan import EDFTomoScan
    # from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
    from nabu.resources.dataset_analyzer import analyze_dataset
    from nabu.pipeline.estimators import CORFinder
    if data_src.endswith('.nx'):
        di = analyze_dataset(data_src % (reference_plane - 1))
    else:
        di = analyze_dataset(data_src % reference_plane)
    # TODO: why did we change the rotation angles ?
    # angular_step = EDFTomoScan.get_scan_range(data_src % reference_plane) / (di.n_angles - 1)
    # di.rotation_angles = np.deg2rad(np.linspace(0, angular_step * di.n_angles, di.n_angles, True))

    # Nabu = 'Flat-fielding with more than 1 darks is not supported'
    if len(di.darks) > 1:
        try:
            k = next(iter(di.darks))
            di.darks = {k: di.darks[k]}
        except Exception as ex:
            logger.error(traceback.format_exc())

    logger.info("Launching Nabu's CORFinder/sliding/growing-window (motion=%s)" % str(motion))
    cor_options = {"side": side} if side is not None else {}
    if tomo_rot_center_near is not None:
        cor_options["tomo_rot_center_near"] = tomo_rot_center_near
    try:
        logger.info(f"CORFinder(sliding-window, cor_options={str(cor_options)})")
        cor_sliding = CORFinder(method="sliding-window", dataset_info=di, do_flatfield=True,
                                cor_options=cor_options, logger=logger).find_cor()
    except Exception as ex:
        logger.warn(f"Failed finding CoR using Nabu (method=sliding-window): {str(ex)}")
        cor_sliding = None
    try:
        cf = CORFinder(method="growing-window", dataset_info=di, do_flatfield=True,
                       cor_options=cor_options, logger=logger)
        cor_growing = cf.find_cor()
    except Exception as ex:
        logger.warn(f"Failed finding CoR using Nabu (method=growing-window): {str(ex)}")
        cor_growing = None
    try:
        cor_accurate = CORFinder(method="octave-accurate", dataset_info=di, do_flatfield=True,
                                 cor_options=cor_options, logger=logger).find_cor()
    except Exception as ex:
        logger.warn(f"Failed finding CoR using Nabu (method=octave-accurate): {str(ex)}")
        cor_accurate = None

    if motion:
        motion_txt = '/correct_motion.txt'
        if os.path.isfile(f"{data_src % reference_plane}/correct_motion.txt"):
            logger.info(f"Taking int account motion correction for CoR: "
                        f"{data_src % reference_plane}/correct_motion.txt")
        elif os.path.isfile(f"{data_src % reference_plane}/correct.txt"):
            logger.info(f"Taking int account motion correction for CoR: "
                        f"{data_src % reference_plane}/correct.txt")
            motion_txt = '/correct.txt'
        else:
            raise RuntimeError("--motion is set, but could not find correct_motion.txt "
                               "or correct.txt file for the reference plane")

        m = np.loadtxt(data_src % reference_plane + motion_txt, dtype=np.float32)
        mx = m[:, 0].take(cf._radios_indices) * reference_plane_magnification
        if cor_sliding is not None:
            logger.info("Sliding window: CoR: %6.2f , motion correction: %5.2f  %5.2f => %6.2f" %
                        (cor_sliding, mx[0], mx[1], cor_sliding - 0.5 * (mx[1] + mx[0])))
            cor_sliding -= 0.5 * (mx[1] + mx[0])
        if cor_growing is not None:
            logger.info("Growing window: CoR: %6.2f , motion correction: %5.2f  %5.2f => %6.2f" %
                        (cor_growing, mx[0], mx[1], cor_growing - 0.5 * (mx[1] + mx[0])))
            cor_growing -= 0.5 * (mx[1] + mx[0])
        if cor_accurate is not None:
            logger.info("Ocatve-accurate: CoR: %6.2f , motion correction: %5.2f  %5.2f => %6.2f" %
                        (cor_accurate, mx[0], mx[1], cor_accurate - 0.5 * (mx[1] + mx[0])))
            cor_accurate -= 0.5 * (mx[1] + mx[0])

    q.put((cor_sliding, cor_growing, cor_accurate))


def find_cor_phase_registration(d0, d180, i0, i180, low_cutoff=None, low_width=0.03,
                                high_cutoff=None, high_width=0.03):
    from ...utils.registration import phase_cross_correlation
    from scipy.fft import fft2, ifft2
    from scipy.ndimage import fourier_shift
    d0 -= d0.mean()
    d180 -= d180.mean()
    dxy, err, dphi = phase_cross_correlation(d0, d180, upsample_factor=10, low_cutoff=low_cutoff, low_width=low_width,
                                             high_cutoff=high_cutoff, high_width=high_width)
    dy, dx = dxy
    # Shift image and re-correlate after cropping wrapped borders
    d0 = ifft2(fourier_shift(fft2(d0), [-dy / 2, -dx / 2])).real
    d180 = ifft2(fourier_shift(fft2(d180), [dy / 2, dx / 2])).real
    adx = int(abs(np.round(dx / 2)))
    if adx >= 1:
        d0 = d0[:, adx:-adx]
        d180 = d180[:, adx:-adx]
    ady = int(abs(np.round(dy / 2)))
    if ady >= 1:
        d0 = d0[ady:-ady]
        d180 = d180[ady:-ady]
    dxy, err, dphi = phase_cross_correlation(d0, d180, upsample_factor=20, low_cutoff=low_cutoff, low_width=low_width,
                                             high_cutoff=high_cutoff, high_width=high_width)
    dy1, dx1 = dxy
    return dx + dx1, dy + dy1, err


def find_cor_phase_registration_kw(kwargs):
    return find_cor_phase_registration(**kwargs)


def find_cor(data_src, proj_idx, binning, reference_plane, method,
             motion, side, ph_shape, ph_array_dtype, logger, ph_array,
             reference_plane_magnification, tomo_rot_center_near):
    """
    Find the Center of Rotation using tomotools & Nabu, with the original iobs or the
    phased projections or sinograms
    """
    from tomoscan.esrf.scan.edfscan import EDFTomoScan
    from nabu.resources.dataset_analyzer import analyze_dataset
    # from nabu.estimation.cor import CenterOfRotationSlidingWindow
    logger.info("Searching for Centre of rotation [NB: reported values are not binned]")
    cor = None
    method0 = method

    # Default CoR using Nabu's CORFinder ##################################
    logger.info("Launching Nabu's CORFinder: sliding/growing-window & octave-accurate methods")
    q_iobs_window = Queue()
    p_iobs_window = Process(target=find_cor_iobs_window,
                            args=(data_src, reference_plane, motion, side, logger,
                                  q_iobs_window, reference_plane_magnification,
                                  tomo_rot_center_near))
    p_iobs_window.start()

    # CoR using +180 phased projections (even ones) ##################################
    logger.info("CoR estimation using +180 phased projections")
    # Try to get ~20 couple of projections if we have more than 180°
    if data_src.endswith('.nx'):
        di = analyze_dataset(data_src % (reference_plane - 1))
        tomo_n = di.n_angles
        scan_range = di.rotation_angles[-1] - di.rotation_angles[0]
    else:
        scan_range = EDFTomoScan.get_scan_range(data_src % reference_plane)
        tomo_n = EDFTomoScan.get_tomo_n(data_src % reference_plane)
    tomo_angle_step_orig = scan_range / (tomo_n - 1)  # Keep in degrees to match user input
    tomo_angle_step = tomo_angle_step_orig * (proj_idx.max() - proj_idx.min()) / (len(proj_idx) - 1)
    nproj, ny, nx = ph_shape
    angles = np.deg2rad(np.arange(nproj) * tomo_angle_step)
    ph = np.frombuffer(ph_array.get_obj(), dtype=ph_array_dtype).reshape(ph_shape)

    angles -= angles.min()
    nb_above_180 = (angles > np.pi).sum()
    step = 2  # We only want even frames
    if nb_above_180 > 20:
        step = nb_above_180 // 20
        step += step % 2
    idx0 = np.arange(0, nb_above_180 + 1, step)
    vcor = []
    vkw = []
    for i0 in idx0:
        da = (angles[::2] - angles[i0] - np.pi) % (2 * np.pi)
        i180 = np.argmin(np.minimum(da, 2 * np.pi - da)) * 2
        # Projections, without borders
        d0 = ph[i0, 100:-100, 100:-100].astype(np.float32)
        d180 = np.fliplr(ph[i180, 100:-100, 100:-100]).astype(np.float32)
        vkw.append({'d0': d0, 'd180': d180, 'i0': i0, 'i180': i180,
                    'low_cutoff': None, 'low_width': None,
                    'high_cutoff': None, 'high_width': None})
    nprocmax = max(1, int(np.ceil(len(vkw) / 2)))
    with Pool(nproc if nproc < nprocmax else nprocmax) as pool:
        results = pool.imap(find_cor_phase_registration_kw, vkw)
        for i in range(len(vkw)):
            dx, dy, err = results.next()
            logger.info("Cross-correlation between #%4d and #%4d: dx=%5.2f dy=%5.2f err=%6e => cor=%6.2f/binning" %
                        (vkw[i]['i0'], vkw[i]['i180'], dx, dy, err, (dx + nx) / 2 * binning))
            vcor.append((dx + nx) / 2 * binning)

    vcor = np.array(vcor)
    with np.printoptions(precision=2, suppress=False, floatmode='fixed', linewidth=200):
        str_vcor = str(vcor)
    logger.info("CoR using registration of phased projections at +180°: " + str_vcor)
    cor_ph_reg = np.median(vcor)
    corstd = np.std(vcor)
    logger.info("CoR using registration of phased projections at +180°: <%6.2f> +/-<%6.3f>" % (cor_ph_reg, corstd))

    if method == "phase_registration":
        cor = cor_ph_reg

    # CoR using +180 phased projections (even ones) with a high-pass filter
    # Same method as previous but with the high-pass filter
    logger.info("CoR estimation using +180 phased projections with a high-pass filter")
    vcorhigh = []
    vkw = []
    for i0 in idx0:
        da = (angles[::2] - angles[i0] - np.pi) % (2 * np.pi)
        i180 = np.argmin(np.minimum(da, 2 * np.pi - da)) * 2
        # Projections, without borders
        d0 = ph[i0, 100:-100, 100:-100].astype(np.float32)
        d180 = np.fliplr(ph[i180, 100:-100, 100:-100]).astype(np.float32)
        vkw.append({'d0': d0, 'd180': d180, 'i0': i0, 'i180': i180,
                    'low_cutoff': 0.03, 'low_width': 0.03,
                    'high_cutoff': None, 'high_width': None})
    nprocmax = max(1, int(np.ceil(len(vkw) / 2)))
    with Pool(nproc if nproc < nprocmax else nprocmax) as pool:
        results = pool.imap(find_cor_phase_registration_kw, vkw)
        for i in range(len(vkw)):
            dx, dy, err = results.next()
            logger.info("Cross-correlation between #%4d and #%4d [high-pass]: dx=%5.2f dy=%5.2f err=%6e "
                        "=> cor=%6.2f/binning" %
                        (vkw[i]['i0'], vkw[i]['i180'], dx, dy, err, (dx + nx) / 2 * binning))
            vcorhigh.append((dx + nx) / 2 * binning)

    vcorhigh = np.array(vcorhigh)
    with np.printoptions(precision=2, suppress=False, floatmode='fixed', linewidth=200):
        str_vcorhigh = str(vcorhigh)
    logger.info("CoR using registration of phased projections (high-pass) at +180°: " + str_vcorhigh)
    cor_ph_reg_high = np.median(vcorhigh)
    corhighstd = np.std(vcorhigh)
    logger.info("CoR using registration of phased projections (high-pass) at +180°: <%6.2f> +/-<%6.3f>"
                % (cor_ph_reg_high, corhighstd))

    if method == "phase_registration_high":
        cor = cor_ph_reg_high

    # CoR using +180 phased projections (even ones) difference
    # Same method as previous but cross-correlate the difference between successive images
    logger.info("CoR estimation using +180 phased - differential approach")
    vcordiff = []
    vkw = []
    for i0 in idx0:
        da = (angles[::2] - angles[i0] - np.pi) % (2 * np.pi)
        i180 = np.argmin(np.minimum(da, 2 * np.pi - da)) * 2
        if (i0 + 2) >= nproj or (i180 + 2) >= nproj:
            continue
        # Projections, without borders
        d0 = (ph[i0, 100:-100, 100:-100] - ph[i0 + 2, 100:-100, 100:-100]).astype(np.float32)
        d180 = np.fliplr(ph[i180, 100:-100, 100:-100] - ph[i180 + 2, 100:-100, 100:-100]).astype(np.float32)
        vkw.append({'d0': d0, 'd180': d180, 'i0': i0, 'i180': i180,
                    'low_cutoff': None, 'low_width': None,
                    'high_cutoff': None, 'high_width': None})
    if len(vkw):
        nprocmax = max(1, int(np.ceil(len(vkw) / 2)))
        with Pool(nproc if nproc < nprocmax else nprocmax) as pool:
            results = pool.imap(find_cor_phase_registration_kw, vkw)
            for i in range(len(vkw)):
                dx, dy, err = results.next()
                logger.info("Cross-correlation between #%4d and #%4d [diff]: dx=%5.2f dy=%5.2f err=%6e "
                            "=> cor=%6.2f/binning" %
                            (vkw[i]['i0'], vkw[i]['i180'], dx, dy, err, (dx + nx) / 2 * binning))
                vcordiff.append((dx + nx) / 2 * binning)

        vcordiff = np.array(vcordiff)
        with np.printoptions(precision=2, suppress=False, floatmode='fixed', linewidth=200):
            str_vcordiff = str(vcordiff)
        logger.info("CoR using registration of phased projections diff at +180°: " + str_vcordiff)
        cor_ph_reg_diff = np.median(vcordiff)
        cordiffstd = np.std(vcordiff)
        logger.info("CoR using registration of phased projections diff at +180°: <%6.2f> +/-<%6.3f>"
                    % (cor_ph_reg_diff, cordiffstd))

        if method == "phase_registration_diff":
            cor = cor_ph_reg_diff

    # CoR using +180 phased projections (even ones) difference and a high-pass filter
    # Same method as previous but with the high-pass filter
    logger.info("CoR estimation using +180 phased projections- differential & high-pass filter")
    vcordiffhigh = []
    vkw = []
    for i0 in idx0:
        da = (angles[::2] - angles[i0] - np.pi) % (2 * np.pi)
        i180 = np.argmin(np.minimum(da, 2 * np.pi - da)) * 2
        if (i0 + 2) >= nproj or (i180 + 2) >= nproj:
            continue
        # Projections, without borders
        d0 = (ph[i0, 100:-100, 100:-100] - ph[i0 + 2, 100:-100, 100:-100]).astype(np.float32)
        d180 = np.fliplr(ph[i180, 100:-100, 100:-100] - ph[i180 + 2, 100:-100, 100:-100]).astype(np.float32)
        vkw.append({'d0': d0, 'd180': d180, 'i0': i0, 'i180': i180,
                    'low_cutoff': 0.03, 'low_width': 0.03,
                    'high_cutoff': None, 'high_width': None})
    if len(vkw):
        nprocmax = max(1, int(np.ceil(len(vkw) / 2)))
        with Pool(nproc if nproc < nprocmax else nprocmax) as pool:
            results = pool.imap(find_cor_phase_registration_kw, vkw)
            for i in range(len(vkw)):
                dx, dy, err = results.next()
                logger.info("Cross-correlation between #%4d and #%4d [diff/high-pass]: dx=%5.2f dy=%5.2f err=%6e "
                            "=> cor=%6.2f/binning" %
                            (vkw[i]['i0'], vkw[i]['i180'], dx, dy, err, (dx + nx) / 2 * binning))
                vcordiffhigh.append((dx + nx) / 2 * binning)

        vcordiffhigh = np.array(vcordiffhigh)
        with np.printoptions(precision=2, suppress=False, floatmode='fixed', linewidth=200):
            str_vcordiffhigh = str(vcordiffhigh)
        logger.info("CoR using registration of phased projections diff/high-pass at +180°: " + str_vcordiffhigh)
        cor_ph_reg_diff_high = np.median(vcordiffhigh)
        cordiffhighstd = np.std(vcordiffhigh)
        logger.info("CoR using registration of phased projections diff/high-pass at +180°: <%6.2f> +/-<%6.3f>"
                    % (cor_ph_reg_diff_high, cordiffhighstd))

        if method == "phase_registration_diff_high":
            cor = cor_ph_reg_diff_high

    # Get back result from iobs_sliding_window process
    cor_sliding, cor_growing, cor_accurate = q_iobs_window.get()

    # Check if chosen method (from Nabu) is available
    msg_allfail = "All Nabu center-of-rotation methods have failed. Check that all " \
                  "data (notably averaged dark and reference files) are present ? " \
                  "You can also give directly the CoR using e.g. " \
                  "'--tomo_rot_center 1279.5', or choose another CoR method using the " \
                  "phased projections."

    if method == "radios_sliding_window" and cor_sliding is None:
        if cor_growing is not None:
            method = "radios_growing_window"
        elif cor_accurate is not None:
            method = "radios_accurate"
        else:
            logger.error(msg_allfail)
            raise RuntimeError(msg_allfail)
        logger.error(f"CoR: method 'accurate' not available, switching to {method}")
    if method == "radios_growing_window" and cor_growing is None:
        if cor_sliding is not None:
            method = "radios_sliding_window"
        elif cor_accurate is not None:
            method = "radios_accurate"
        else:
            logger.error(msg_allfail)
            raise RuntimeError(msg_allfail)
        logger.error(f"CoR: method growing_window not available, switching to {method}")
    if method == "radios_accurate" and cor_accurate is None:
        if cor_sliding is not None:
            method = "radios_sliding_window"
        elif cor_accurate is not None:
            method = "radios_growing_window"
        else:
            logger.error(msg_allfail)
            raise RuntimeError(msg_allfail)
        logger.error(f"CoR: method 'accurate' not available, switching to {method}")

    p_iobs_window.join()
    if method == "radios_sliding_window":
        cor = cor_sliding
    if cor_sliding is not None:
        logger.info("CORFinder/sliding-window result: cor=%6.2f/binning" % cor_sliding)

    if method == "radios_growing_window":
        cor = cor_growing
    if cor_growing is not None:
        logger.info("CORFinder/growing-window result: cor=%6.2f/binning" % cor_growing)

    if method == "radios_accurate":
        cor = cor_accurate
    if cor_accurate is not None:
        logger.info("CORFinder/octave-accurate result: cor=%6.2f/binning" % cor_accurate)

    if cor is None:
        if cor_sliding is not None:
            logger.warn(f"CORFinder - could not apply selected method ({str(method0)}): "
                        f"reverting to radios_sliding_window: {cor:6.2f}")
            cor = cor_sliding
        else:
            logger.error(f"CORFinder - could not apply selected method ({str(method0)}): "
                         f"and radios_sliding_window is not available as a backup method "
                         f"(please check the data files for averaged dark and ref files if "
                         f"Nabu's CoR estimator are failing). "
                         f"You can also give directly the CoR using e.g. "
                         f"'--tomo_rot_center 1279.5'")
    else:
        logger.info("Returning CoR using method [%s]: %6.2f" % (method, cor))
    return cor

    if False:
        # Evaluate CoR using multiple differential projections #####
        # Tomo angular parameters
        scan_range, tomo_n = EDFTomoScan.get_scan_range(data_src), EDFTomoScan.get_tomo_n(data_src)
        tomo_angle_step_orig = scan_range / (tomo_n - 1)  # Keep in degrees to match user input
        nproj = len(proj_idx)
        tomo_angle_step = tomo_angle_step_orig * (proj_idx.max() - proj_idx.min()) / (nproj - 1)
        angles = np.deg2rad(np.arange(nproj) * tomo_angle_step)
        angles -= angles.min()
        nb_above_180 = (angles > np.pi).sum()
        step = 1
        if nb_above_180 > 20:
            step = nb_above_180 // 20
            step += step % 2
        idx0 = np.arange(0, nb_above_180 + 1, step)
        print(idx0)
        vcor = []
        v_projs = []
        for i0 in idx0:
            img_name = data_src[:-2] + "%d_/" + prefix + "%d_%04d.edf"
            i180 = np.argmin(np.abs(angles - angles[i0] - np.pi))
            print(i0, i180)
            # Load projections, difference with neighbouring ones
            d0 = fabio.open(img_name % (reference_plane, reference_plane, i0)).data
            d0m = [d0]
            for i in [i0 - 1, i0 + 1]:
                if 0 <= i < tomo_n:
                    d0m.append(fabio.open(img_name % (reference_plane, reference_plane, i)).data)
            d0 = d0.astype(np.float32) - np.array(d0m).mean(axis=0)

            d180 = fabio.open(img_name % (reference_plane, reference_plane, i180)).data
            d180m = [d180]
            for i in [i180 - 1, i180 + 1]:
                if 0 <= i < tomo_n:
                    d180m.append(fabio.open(img_name % (reference_plane, reference_plane, i)).data)
            d180 = d180.astype(np.float32) - np.array(d180m).mean(axis=0)

            if binning is not None:
                if binning > 1:
                    d0 = rebin(d0, binning)
                    d180 = rebin(d180, binning)

            d180 = np.fliplr(d180)
            cor = CenterOfRotationSlidingWindow(logger=logger).find_shift(d0, d180, 'center') + d0.shape[1] / 2
            if hasattr(cor, "__iter__"):
                cor = cor[0]
            print("CoR on derivative projections / sliding window: i=%4d, %4d  CoR=%6.2f" % (i0, i180, cor))
            vcor.append(cor)
            v_projs.append(d0)
            v_projs.append(d180)
        with h5py.File("cor.h5", "w") as h:
            h.create_dataset('projs', data=np.array(v_projs))
        logger.info("CoR using derivative projections: %s" % str(vcor))
        cor = np.median(np.array(vcor)) * binning
        print("Final CoR based on %d +180° derivative projections: %6.2f" % (len(idx0), cor))
        return cor


def find_cor_half(binning, logger, ph_array, ph_shape, ph_array_dtype, proj_idx, data_src, reference_plane):
    """
    Find the Center of Rotation using Nabu, with half of the projections
    (all the even ones) reconstructed.

    :param logger: logger object
    :return: the center of rotation (multiplied by the binning value)
    """
    logger.info("Finding rot center from phased projections.")
    from tomoscan.esrf.scan.edfscan import EDFTomoScan
    from nabu.estimation.cor_sino import SinoCor
    from nabu.estimation.cor import (CenterOfRotation, CenterOfRotationAdaptiveSearch,
                                     CenterOfRotationSlidingWindow, CenterOfRotationGrowingWindow)

    # Tomo angular parameters
    scan_range, tomo_n = EDFTomoScan.get_scan_range(data_src % reference_plane), \
        EDFTomoScan.get_tomo_n(data_src, reference_plane)
    tomo_angle_step_orig = scan_range / (tomo_n - 1)  # Keep in degrees to match user input
    tomo_angle_step = tomo_angle_step_orig * (proj_idx.max() - proj_idx.min()) / (len(proj_idx) - 1)
    nproj, ny, nx = ph_shape
    angles = np.deg2rad(np.arange(nproj) * tomo_angle_step)
    ph = np.frombuffer(ph_array.get_obj(), dtype=ph_array_dtype).reshape(ph_shape)
    # Evaluate CoR from half-sinogram (even projection)
    side = "right"
    window_width = None
    neighborhood = 7
    shift_value = 0.1

    # Extract sinograms with even projections #################
    sino = ph[::2, ny // 4: 3 * ny // 4 + 1:ny // 20]
    # We need an even shape
    if len(sino) % 2:
        sino = sino[:-1]
    # We use N layers and use the median value
    vcor_sino = []
    v_sino = []
    for i in range(sino.shape[1]):
        s = sino[:, i].astype(np.float32)
        # Perform some filtering
        s = median_filter(s, (3, 1))
        s -= s.mean(axis=0)
        cor_finder = SinoCor(sinogram=s, logger=logger)
        cor_finder.estimate_cor_coarse(side=side, window_width=window_width)
        vcor_sino.append(cor_finder.estimate_cor_fine(neighborhood=neighborhood, shift_value=shift_value))
        v_sino.append(s)
    vcor_sino = np.array(vcor_sino)
    cor = np.median(vcor_sino)
    with np.printoptions(precision=2, suppress=False, floatmode='fixed'):
        str_vcor = str(vcor_sino)
    logger.info("CoR for different layers using sinograms: " + str_vcor)

    if False:
        # CoR using +180 projections # Extract sinograms with even projections #################
        # Try to get ~20 couple of projections if we have more than 180°
        angles -= angles.min()
        nb_above_180 = (angles > np.pi).sum()
        step = 2  # We only want even frames
        if nb_above_180 > 20:
            step = nb_above_180 // 20
            step += step % 2
        idx0 = np.arange(0, nb_above_180 + 1, step)
        print(idx0)
        vcor = []
        v_projs = []
        for i0 in idx0:
            i180 = np.argmin(np.abs(angles - angles[i0] - np.pi))
            print(i0, i180)
            # Projections, without borders
            d0 = ph[i0, 100:-100, 100:-100].copy()
            d180 = np.fliplr(ph[i180, 100:-100, 100:-100]).copy()
            d0 -= d0.mean()
            d180 -= d180.mean()
            cor = CenterOfRotationSlidingWindow(logger=logger).find_shift(d0, d180, 'center') + nx / 2
            if hasattr(cor, "__iter__"):
                cor = cor[0]
            vcor.append(cor)
            # cor2 = CenterOfRotationGrowingWindow(logger=logger).find_shift(d0, d180, 'center') + nx / 2
            # if hasattr(cor2, "__iter__"):
            #     cor2 = cor2[0]
            # vcor.append(cor2)
            v_projs += [d0, d180]
        logger.info("CoR for couple of projections at +180° : " + str(vcor))
        with h5py.File("cor.h5", "w") as h:
            h.create_dataset('sino', data=np.array(v_sino))
            h.create_dataset('projs', data=np.array(v_projs))

        cor = np.median(vcor)
    logger.info("Final rotation center using median value: %8.2f" % (cor * binning))
    return cor * binning


def prep_holotomodata(iobs_array, iobs_shape, ref, pixel_size, wavelength, detector_distance,
                      dx, dy, proj_idx, padding, algorithm, gpu_mem_f, pu, logger, data_old=None,
                      psf=None, stack_size=None):
    from pynx.holotomo import HoloTomoData
    from pyvkfft.cuda import VkFFTApp
    # Iobs from shared memory buffer
    nproj, nz, ny, nx = iobs_shape
    nproj_sharray = len(iobs_array) // (nz * ny * nx)  # Actual len of shared array
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape((nproj_sharray, nz, ny, nx))
    iobs = iobs[:nproj]
    if stack_size is None:
        # Estimate how much GPU memory will be used
        gpu_mem_f0 = 1.4  # Add +40% to the estimation, as a margin for errors
        mem = pu.cu_device.total_memory()
        logger.info("Available GPU memory: %6.2fGB  data size: %6.2fGB" % (mem / 1024 ** 3, iobs.nbytes / 1024 ** 3))
        # Required memory: data, iobs, object (complex + phase), 1 copy of psi per projection (DRAP/DM/RAAR)
        mem_stack = 4 * nz * nx * ny  # Iobs
        mem_stack += 8 * nz * nx * ny  # Psi
        mem_stack += 12 * nx * ny  # object projections & phase
        if 'RAAR' in algorithm or 'DM' in algorithm or 'DRAP' in algorithm:
            mem_stack += 8 * nz * nx * ny
        # pyvkfft requires an extra buffer for inplace transforms, for large arrays
        mem_fft = nz * VkFFTApp((ny, nx), dtype=np.complex64).get_tmp_buffer_nbytes()
        # Memory required by the GPU context (estimated)
        mem_ctx = 0.5 * 1024 ** 3
        if psf is not None:
            if psf > 0:
                mem_stack += 4 * nz * (nx + 2) * ny
                mem_fft *= 1.5
        logger.info("(Over-)estimated GPU memory requirement: %8.4fGB/projection"
                    % (mem_stack * gpu_mem_f0 / 1024 ** 3))
        # Can we fit this in memory with 1 stack ?
        stack_size = int(np.ceil((mem - mem_ctx) / ((mem_stack * gpu_mem_f0 + mem_fft) * gpu_mem_f)))
        logger.info(f"Maximum single stack size:  {stack_size} <? nproj={nproj}")
        if False:  # TODO stack_size > nproj:
            stack_size = nproj
        else:
            # We need two stacks in memory (one for computing, the other to swap)
            mem_stack *= 2
            stack_size = int(np.ceil((mem - mem_ctx) / ((mem_stack * gpu_mem_f0 + mem_fft) * gpu_mem_f)))
            if nproj // stack_size == 2:
                # We need either n=1 or n>=3 stacks for swapping
                stack_size = int(np.ceil(nproj // 3))
        m = stack_size * mem_stack * gpu_mem_f0 * (1 + nproj > stack_size) + mem_fft
        logger.info(f"Using {nproj // stack_size} stacks with size = {stack_size}, "
                    f"over-estimated GPU mem = {m / 1024 ** 3:.1f} GB")
    ################################################################
    # Create HoloTomoData and HoloTomo objects
    ################################################################
    if data_old is not None:
        logger.info("Re-using existing HolotomoData object with pinned memory")
        data = data_old
        data.replace_data(iobs, ref, dx=dx, dy=dy, idx=proj_idx, pu=pu)
    else:
        logger.info("Creating HolotomoData object and allocating projections in pinned memory")
        data = HoloTomoData(iobs, ref, pixel_size_detector=pixel_size, wavelength=wavelength,
                            detector_distance=detector_distance, stack_size=stack_size,
                            dx=dx, dy=dy, idx=proj_idx, padding=padding, pu=pu)
    return data


def run_algorithms(data, iobs_shape, padding, delta_beta, algorithm, ph_array, ph_shape,
                   i0, nb_chunk, obj_smooth, obj_inertia, obj_min, obj_max, nb_probe,
                   sino_filter, pu, logger, old_ht=None, return_probe=False,
                   normalise=False, probe0=None, magnification=None,
                   constrain_probe_direct_beam=False, psf=None, binning=None):
    from pynx.holotomo import HoloTomo
    from pynx.holotomo.operator import AP, DRAP, DM, RAAR, ScaleObjProbe, InitPSF, \
        BackPropagatePaganin, BackPropagateCTF, FreePU, MemUsage, SinoFilter
    from ..utils import zoom_pad_images

    # Iobs from shared memory buffer
    nproj, nz, ny, nx = iobs_shape

    ################################################################
    # Create HoloTomo object
    ################################################################
    if old_ht is None:
        ################################################################
        # Probe init
        ################################################################
        if probe0 is not None:
            # Load probe from the result of an NFP analysis
            probe = []
            for iz in range(nz):
                if len(probe0) == 1:
                    s = probe0[0] % (iz + 1)
                else:
                    s = probe0[iz]
                with h5py.File(s, 'r') as h:
                    logger.info(f"Loading initial probe (iz={iz}) from NFP optimisation: {s}")
                    # We use only the first mode
                    if len(h['entry_last/probe/data']) > 1:
                        logger.warn(f"Loaded probe had {len(h['entry_last/probe/data'])} modes ! "
                                    f"Only using the first one")
                    probe.append(h['entry_last/probe/data'][0])
            logger.info(f"Zoom & pad probe arrays")
            # TODO: should maybe zoom amplitude & phase instead ?
            probe = zoom_pad_images(probe, magnification, padding, nz, mask2neg=False)
            probe = np.reshape(probe, (nz, 1, ny, nx))
            nb_probe = 1
            coherent_probe_modes = False
        else:
            probe = np.ones((nz, nb_probe, ny, nx), dtype=np.complex64)
            if nb_probe > 1:
                # Use coherent probe modes ?
                logger.info("Using %d coherent probe modes with linear ramp coefficients" % nb_probe)
                # Use linear ramps for the probe mode coefficients
                coherent_probe_modes = np.zeros((nproj, nz, nb_probe))
                dn = nproj // (nb_probe - 1)
                for iz in range(nz):
                    for i in range(nb_probe - 1):
                        if i < (nb_probe - 2):
                            coherent_probe_modes[i * dn:(i + 1) * dn, iz, i] = np.linspace(1, 0, dn)
                            coherent_probe_modes[i * dn:(i + 1) * dn, iz, i + 1] = np.linspace(0, 1, dn)
                        else:
                            n = nproj - i * dn
                            coherent_probe_modes[i * dn:, iz, i] = np.linspace(1, 0, n)
                            coherent_probe_modes[i * dn:, iz, i + 1] = np.linspace(0, 1, n)

                # Use only 1 mode for each projection
                # coherent_probe_modes[:] = coherent_probe_modes >= 0.5

                # All in first mode - should yield the same result as without modes
                # coherent_probe_modes[:] = 0
                # coherent_probe_modes[:, :, 0] = 1
            else:
                coherent_probe_modes = False

        logger.info("Creating Holotomo object")
        p = HoloTomo(data=data, obj=None, probe=probe, coherent_probe_modes=coherent_probe_modes, pu=pu)
    else:
        logger.info("Re-using Holotomo object")
        p = old_ht
        p._cu_timestamp_counter, p._cl_timestamp_counter = -1, -1
        p.cycle = 0

    ################################################################
    # Algorithms
    ################################################################
    logger.info("Starting algorithms")
    if constrain_probe_direct_beam:
        logger.info("Adding a direct beam constraint for the probe")
        p.constrain_probe_direct_beam = True
    t0 = timeit.default_timer()
    # print(nz, 1, ny, nx, p.data.nz, p.nb_probe, p.data.ny, p.data.nx)
    p = ScaleObjProbe() * p
    beta = 0.75
    dm_alpha = 0.02
    show_obj_probe = False
    probe_inertia = 0.01

    update_obj = True

    update_probe = not normalise and probe0 is None
    verbose = 10
    if psf is not None:
        logger.info(f"Initialising psf with gaussian fwhm={psf:.2f}/binning pixels")
        p = InitPSF(psf / binning) * p
    for algo in algorithm.split(",")[::-1]:
        if "=" in algo:
            logger.info("Changing parameter? %s" % algo)
            k, v = algo.split("=")
            if k == "delta_beta":
                delta_beta = eval(v)
            elif k == "psf":
                p = InitPSF(float(v) / binning) * p
            elif k == "verbose":
                verbose = int(eval(v))
            elif k == "beta":
                beta = eval(v)
            elif k == "dm_alpha":
                dm_alpha = eval(v)
            elif k == "obj_smooth":
                obj_smooth = eval(v)
            elif k == "obj_inertia":
                obj_inertia = eval(v)
            elif k == "obj_min":
                obj_min = eval(v)
            elif k == "obj_max":
                obj_max = eval(v)
            elif k == "probe_inertia":
                probe_inertia = eval(v)
            elif k == "probe":
                update_probe = eval(v)
            elif k == "obj":
                update_obj = eval(v)
            elif k == "constrain_probe_direct_beam":
                p.constrain_probe_direct_beam = bool(eval(v))
                logger.info(f"Direct beam constraint for the probe set to:{p.constrain_probe_direct_beam}")
            else:
                logger.info("WARNING: did not understand algorithm step: %s" % algo)
        elif "paganin" in algo.lower():
            logger.info(f"Paganin back-projection (δ/ß={delta_beta:.1f}) + 2 AP/p cycles")
            if not update_probe:
                p = BackPropagatePaganin(delta_beta=delta_beta, keep_probe=probe0 is not None) * p
            else:
                ap = AP(update_object=False, update_probe=True,
                        calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                        obj_smooth=obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                        show_obj_probe=show_obj_probe, probe_inertia=probe_inertia)
                p = ap ** 2 * BackPropagatePaganin(delta_beta=delta_beta) * p
        elif "ctf" in algo.lower():
            logger.info(f"CTF back-projection (δ/ß={delta_beta:.1f}) + 2 AP/p cycles")
            if not update_probe:
                p = BackPropagateCTF(delta_beta=delta_beta, keep_probe=probe0 is not None) * p
            else:
                ap = AP(update_object=False, update_probe=True,
                        calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                        obj_smooth=obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                        show_obj_probe=show_obj_probe, probe_inertia=probe_inertia)
                p = ap ** 2 * BackPropagateCTF(delta_beta=delta_beta) * p
        else:
            logger.info("Algorithm step: %s" % algo)
            if ('dm' in algo.lower() or 'raar' in algo.lower() or 'drap' in algo.lower()) \
                    and not update_probe and delta_beta > 0:
                logger.warn(f"Using DM/RAAR/DRAP without updating the probe combined with "
                            f"delta_beta={delta_beta}>0 will likely be unstable")
            dm = DM(update_object=update_obj, update_probe=update_probe,
                    calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                    obj_smooth=obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                    show_obj_probe=show_obj_probe, probe_inertia=probe_inertia,
                    alpha=dm_alpha)
            ap = AP(update_object=update_obj, update_probe=update_probe,
                    calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                    obj_smooth=obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                    show_obj_probe=show_obj_probe, probe_inertia=probe_inertia)
            apn = AP(update_object=update_obj, update_probe=update_probe,
                     calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                     obj_smooth=obj_smooth, obj_inertia=obj_inertia,
                     weight_empty=1, show_obj_probe=show_obj_probe,
                     iobs_normalise=True, probe_inertia=probe_inertia)
            raar = RAAR(update_object=update_obj, update_probe=update_probe, beta=beta,
                        calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                        obj_smooth=obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                        show_obj_probe=show_obj_probe, probe_inertia=probe_inertia)
            drap = DRAP(update_object=update_obj, update_probe=update_probe, beta=beta,
                        calc_llk=verbose, delta_beta=delta_beta, obj_min=obj_min, obj_max=obj_max,
                        obj_smooth=obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                        show_obj_probe=show_obj_probe, probe_inertia=probe_inertia)
            p = eval(algo.lower()) * p

    if sino_filter is not None:
        p = SinoFilter(sino_filter) * p

    # p = FreePU() * MemUsage() * p
    logger.info("Algorithms: dt = %6.2fs" % (timeit.default_timer() - t0))
    # We use float32 because float16 operations are extremely slow
    idx, ph0 = p.get_obj_phase_unwrapped(crop_padding=True, dtype=np.float32, sino_filter=None)
    logger.info("Got unwrapped phases (sino_filter:%s): dt = %6.2fs" % (str(sino_filter), timeit.default_timer() - t0))
    # Move phases to shared memory array
    ph = np.frombuffer(ph_array.get_obj(), dtype=np.float16).reshape(ph_shape)
    # Move ph to shared memory buffer
    if nb_chunk == 1:
        ph[:] = ph0
    else:
        # print(i0 + i * nb_chunk, i, ph0.shape, ph.shape)
        ph[i0::nb_chunk] = ph0
    logger.info("Copied back phases for the main process")

    if not return_probe:
        return idx, None, p

    pady, padx = p.data.padding
    if pady and padx:
        probe = p.get_probe()[..., pady:-pady, padx:-padx]
    elif pady:
        probe = p.get_probe()[..., pady:-pady]
    elif padx:
        probe = p.get_probe()[..., :, padx:-padx]
    else:
        probe = p.get_probe()

    # only for debugging - save object and probe
    # p.save_obj_probe_chunk("obj_probe_%04d", save_obj_complex=True, save_probe=True, dtype=np.float16)
    return idx, probe, p


# Shared arrays dictionary for the FBP pool
shared_arrays_fbp = {}


def fbp(vy, nproj, ny, nx, tomo_rot_center, tomo_angle_step, tomo_angle_start,
        profile, igpu, sino_filter="ram-lak", halftomo=None):
    """

    :param vy:
    :param nproj:
    :param ny:
    :param nx:
    :param tomo_rot_center:
    :param tomo_angle_step:
    :param profile:
    :param igpu:
    :param sino_filter: if None, assume that the sinogram filtering before the FBP has already been
        done.
    :param halftomo: either None, 'right' or 'left'
    :return:
    """
    # os.environ['PYNX_PU'] = 'cuda.%d' % igpu
    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
        os.environ['CUDA_VISIBLE_DEVICES'] = "%d" % igpu
    else:
        vigpu = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
        os.environ['CUDA_VISIBLE_DEVICES'] = vigpu[igpu]
    from nabu.reconstruction.fbp import Backprojector
    from nabu.reconstruction.sinogram import get_extended_sinogram_width
    from ...processing_unit.cu_resources import cu_resources
    from ...processing_unit import default_processing_unit as pu
    import pycuda.gpuarray as cua
    pu.select_gpu(language='cuda')
    pu.enable_profiling(profile)

    sino_width = get_extended_sinogram_width(nx, tomo_rot_center) if halftomo else nx
    cu_ctx = cu_resources.get_context(pu.cu_device)

    ph_array = shared_arrays_fbp['ph_array']
    vol_array = shared_arrays_fbp['vol_array']
    ph = np.frombuffer(ph_array.get_obj(), dtype=np.float16).reshape((nproj, ny, nx))
    vol_shape = (ny, sino_width, sino_width)
    vol = np.frombuffer(vol_array.get_obj(), dtype=np.float16).reshape(vol_shape)

    idx_out = shared_arrays_fbp['idx_out']
    idx_in = shared_arrays_fbp['idx_in']

    extra_options = {"centered_axis": True} if halftomo is not None else None
    B = Backprojector((nproj, nx), rot_center=tomo_rot_center,
                      angles=np.deg2rad(np.arange(nproj) * tomo_angle_step + tomo_angle_start),
                      filter_name=sino_filter, backend_options={'ctx': cu_ctx},
                      halftomo=halftomo is not None, extra_options=extra_options)

    # sys.stdout.write('Reconstructing %d slices (%d x %d): ' % (ny, nproj, nx))
    volmin, volmax = [], []
    for i in vy:
        sino = ph[:, i, :].astype(np.float32)
        # if i % 50 == 0:
        #     sys.stdout.write('%d..' % (ny - i))
        #     sys.stdout.flush()
        if sino_filter is None:
            # Filtering has already been done when getting the unwrapped phases,
            # using the padded areas.
            # Need to normalise to keep the same scale as through nabu filtering
            sino *= np.pi / len(sino)
            if halftomo is not None:
                # See https://gitlab.esrf.fr/tomotools/nabu/-/issues/570
                # and https://gitlab.esrf.fr/tomotools/nabu/-/issues/483
                vol[i] = B.backproj(cua.to_gpu(sino))
            else:
                vol[i] = B.backproj(sino)
        else:
            vol[i] = B.fbp(sino)
        vol[i][idx_out] = np.nan
        # Get min/max, avoid slow nanmin and nanmax later
        volmin.append(vol[i][idx_in].min())
        volmax.append(vol[i][idx_in].max())
    # print("\n")
    if profile:
        pu.enable_profiling(False)
    return volmin, volmax


def fbp_pool_init(ph_array, vol_array, vnx):
    shared_arrays_fbp['ph_array'] = ph_array
    shared_arrays_fbp['vol_array'] = vol_array
    # Mask outside of reconstruction
    ix, iy = np.meshgrid(np.arange(vnx) - vnx / 2, np.arange(vnx) - vnx / 2)
    r2 = ix ** 2 + iy ** 2
    shared_arrays_fbp['idx_out'] = r2 >= (vnx * vnx / 4)
    shared_arrays_fbp['idx_in'] = r2 < (vnx * vnx / 4)


def fbp_kw(kwargs):
    return fbp(**kwargs)


def vol2tiff_slices(filename, volslice, pxum, vmin, vmax):
    """
    Save a volume slice to a tiff file.
    :param filename: the filename e.g. sample_vol/sample_vol00053.tiff
    :param volslice: the volume as a float array
    :param pxum: pixel/voxel spacing in micron
    :param vmin: min value to scale the slice to uint16
    :param vmin: max value to scale the slice to uint16
    :return: nothing
    """
    volint = ((volslice - vmin) / (vmax - vmin)) * (2 ** 16 - 1)
    volint[np.isnan(volint)] = 2 ** 15
    imwrite(filename, np.round(volint).astype(np.uint16), imagej=True, resolution=(1 / pxum, 1 / pxum),
            metadata={'spacing': pxum, 'unit': 'um', 'hyperstack': False})


def run_fbp(data_src, reference_plane, prefix_output, ph_array, ph_shape, idx, tomo_rot_center, pixel_size,
            tomo_angle_step, tomo_angle_start,
            save_fbp_vol='tiff', save_3sino=True, logger=None, profile=False, ngpu=1,
            sino_filter='ram-lak', halftomo=None, params=None):
    if logger is None:  # Should not happen
        logger = logging.getLogger()
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_agg import FigureCanvasAgg
    from nabu.reconstruction.sinogram import get_extended_sinogram_width
    os.environ['NABU_FFT_BACKEND'] = 'vkfft'  # https://gitlab.esrf.fr/tomotools/nabu/-/issues/483#note_391626
    logger.info("#" * 48)
    logger.info("# 3D Tomography reconstruction (FBP) with Nabu using %d GPU(s)" % ngpu)
    logger.info("#" * 48)
    t0 = timeit.default_timer()

    ph = np.frombuffer(ph_array.get_obj(), dtype=np.float16).reshape(ph_shape)

    if save_3sino:
        filename = prefix_output + "_3sino.npz"
        logger.info("#" * 48)
        logger.info(" Saving 3 sinograms to: %s" % filename)
        logger.info("#" * 48)
        ny = ph.shape[1]
        np.savez_compressed(filename, sino=np.take(ph, (ny // 4, ny // 2, 3 * ny // 4), axis=1).astype(np.float32))

    # Tomo angular step
    logger.info("Tomography angular step:  %6.4f°" % tomo_angle_step)

    nproj, ny, nx = ph.shape
    if halftomo is not None:
        vnx = get_extended_sinogram_width(nx, tomo_rot_center)
        logger.info(f"Half-tomo: nx={nx}, CoR={tomo_rot_center:6.2f} => extended sino width={vnx}")
    else:
        vnx = nx
    vol_shape = (ny, vnx, vnx)

    logger.info("Allocating volume array: %s [%6.2fGB]" % (str(vol_shape), np.prod(vol_shape) * 2 / 1024 ** 3))
    vol_array = Array(ctypes.c_int16, int(np.prod(vol_shape)))
    vol = np.frombuffer(vol_array.get_obj(), dtype=np.float16).reshape(vol_shape)

    logger.info('Starting Filtered Back-Projection using Nabu - %d slices (%d x %d)' % (ny, nproj, nx))
    if sino_filter is None:
        logger.info('sino_filter=None - assuming FBP filtering has already been performed')
    else:
        logger.info('sino_filter=%s' % sino_filter)
    volmin, volmax = [], []
    n_fbp_proc = 4 * ngpu
    vkw = [{'vy': range(i, ny, n_fbp_proc), 'nproj': nproj, 'ny': ny, 'nx': nx,
            'tomo_rot_center': tomo_rot_center,
            'tomo_angle_step': tomo_angle_step, 'tomo_angle_start': tomo_angle_start,
            'profile': profile,
            'igpu': i % ngpu, 'sino_filter': sino_filter, 'halftomo': halftomo} for i in range(n_fbp_proc)]
    with Pool(n_fbp_proc, initializer=fbp_pool_init, initargs=(ph_array, vol_array, vnx)) as pool:
        results = pool.imap(fbp_kw, vkw)
        for i in range(len(vkw)):
            r = results.next()
            volmin += r[0]
            volmax += r[1]
    volmin = np.array(volmin)
    volmax = np.array(volmax)
    if np.any(np.isnan(volmin)) or np.any(np.isnan(volmax)):
        logger.warning("Some slices min/max are NaNs")
    if np.any(np.isinf(volmin)):
        logger.warning("Some slices min are -Inf")
        volmin[np.isinf(volmin)] = np.nan  # So nanmin/nanmax works later on
    if np.any(np.isinf(volmax)):
        logger.warning("Some slices max are +Inf")
        volmax[np.isinf(volmax)] = np.nan  # So nanmin/nanmax works later on

    dt = timeit.default_timer() - t0
    logger.info("Time to perform FBP reconstruction:  %4.1fs" % dt)

    logger.info("#" * 48)
    logger.info(" Saving central XYZ cuts of volume to: %s_volXYZ.png" % prefix_output)
    logger.info("#" * 48)
    fig = Figure(figsize=(16, 16))
    ax1, ax2, ax3, ax4 = fig.subplots(2, 2).flat
    ax4.set_visible(False)
    c = vol[ny // 2].astype(np.float32)
    vmin, vmax = np.nanpercentile(c, (0.1, 99.9))
    ax = ax1.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')
    # ax1.set_title(os.path.split(prefix_output)[-1] + "[XY]")
    # ax.colorbar()

    # We use a layout similar to what is used in imagej
    c = vol[..., vnx // 2].astype(np.float32).transpose()
    vmin, vmax = np.nanpercentile(c, (0.1, 99.9))
    ax = ax2.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')
    # ax.colorbar()
    # ax2.set_title(os.path.split(prefix_output)[-1] + "[YZ]")

    c = vol[:, vnx // 2].astype(np.float32)
    vmin, vmax = np.nanpercentile(c, (0.1, 99.9))
    ax = ax3.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')
    # ax3.set_title(os.path.split(prefix_output)[-1] + "[XZ]")
    # ax.colorbar()

    # Step between text lines of size 6
    fontsize = 8
    dy = fontsize * 1.2 / 72 / fig.get_size_inches()[1]
    if params is not None:
        fig.text(0.5, 1, r"%s   binning=%d   nz=%d   %s   $\delta/\beta=%3.0f$" %
                 (os.path.split(prefix_output)[-1], params.binning, params.nz,
                  params.algorithm, params.delta_beta),
                 fontsize=10 if len(params.algorithm) > 60 else 12,
                 horizontalalignment='center',
                 verticalalignment='top', stretch='condensed')
        # Print parameters
        y0 = 0.45 - 1.5 * dy
        n = 1
        vpars = vars(params)
        vpars['host'] = "%s - %s" % (platform.node(), platform.platform())
        vpars['directory'] = os.getcwd()
        try:
            vpars['user'] = os.getlogin()
        except OSError:
            # When using slurm, this is not running in a shell with user info
            if 'USER' in os.environ:
                vpars['user'] = os.environ['USER']
        if params.tomo_rot_center is None:
            vpars['tomo_rot_center_real'] = tomo_rot_center
        vpars['pixel_size_real_nm'] = pixel_size * 1e9
        vk = list(vpars.keys())
        vk.sort()
        for k in vk:
            v = vpars[k]
            if v is not None and k not in ['slurm']:
                fig.text(0.55, y0 - n * dy, "%s = %s" % (k, str(v)), fontsize=fontsize, horizontalalignment='left',
                         stretch='condensed')
                n += 1
    else:
        fig.text(0.5, 1, "%s" % (os.path.split(prefix_output)[-1]), fontsize=12,
                 horizontalalignment='center', verticalalignment='top', stretch='condensed')

    # Elapsed time for figure
    dt = time.time() - DeltaTimeFormatter.start_time
    if dt < 60:
        dt = f"{dt:.1f}s"
    elif dt < 3600:
        dt = time.strftime("%Mm%Ss", time.gmtime(dt))
    else:
        dt = time.strftime("%Hh%Mm%Ss", time.gmtime(dt))
    fig.text(dy, dy, f"PyNX v{_pynx_version}, finished at "
                     f"{time.strftime('%Y/%m/%d %H:%M:%S')}, dt={dt}",
             fontsize=fontsize, horizontalalignment='left', stretch='condensed')

    fig.tight_layout()
    canvas = FigureCanvasAgg(fig)
    canvas.print_figure(prefix_output + "_volXYZ.png", dpi=120)

    if isinstance(save_fbp_vol, bool):
        if save_fbp_vol:
            save_fbp_vol = "hdf5"
        else:
            save_fbp_vol = "False"
    if "tif" in save_fbp_vol:
        logger.info("#" * 48)
        logger.info(" Exporting FBP volume as an uint16 TIFF file")
        logger.info("#" * 48)
        t1 = timeit.default_timer()
        logger.info("Computing min/max of volume")
        # The following operations are dead slow using float16
        vmin, vmax = np.nanmin(volmin), np.nanmax(volmax)
        logger.info("Range for int16 conversion: min=%10.4f  max=%10.4f" % (vmin, vmax))
        pxum = pixel_size * 1e6
        if vol.size * 2 >= 4 * 1024 ** 3:  # 4GB tiff file limit
            logger.info("Volume is %6.2fGB > 4GB => saving volume as separate tiff slices" %
                        (vol.size * 2 / 1024 ** 3))
            head, tail = os.path.split(prefix_output)
            logger.info(f"Saving volume to: {prefix_output}_vol/{tail}_vol%05d.tiff")
            os.makedirs("%s_vol" % prefix_output, exist_ok=True)
            with Pool(nproc) as pool:
                vpar = [(f"{prefix_output}_vol/{tail}_vol{i:05d}.tiff", vol[i], pxum, vmin, vmax)
                        for i in range(ny)]
                pool.starmap(vol2tiff_slices, vpar)
            # for i in range(ny):
            #     filename = "%s_vol/%s_vol%05d.tiff" % (prefix_output, prefix_output, i)
            #     imwrite(filename, volint[i], imagej=True, resolution=(1 / pxum, 1 / pxum),
            #             metadata={'spacing': pxum, 'unit': 'um', 'hyperstack': False}, maxworkers=nproc)
        else:
            logger.info("Converting volume to uint16")
            volint = ((vol - vmin) * ((2 ** 16 - 1) / (vmax - vmin))).astype(np.uint16)
            filename = prefix_output + "_vol.tiff"
            logger.info("Writing tiff file: %s" % filename)
            imwrite(filename, volint, imagej=True, resolution=(1 / pxum, 1 / pxum),  # bigtiff=True,
                    metadata={'spacing': pxum, 'unit': 'um', 'hyperstack': False}, maxworkers=nproc)
            logger.info("Finished saving volume to: %s" % filename)
        dt = timeit.default_timer() - t1
        logger.info("Time to export volume as tiff:  %4.1fs" % dt)
    if "hdf5" in save_fbp_vol:
        logger.info("#" * 48)
        logger.info(" Saving FBP volume as an hdf5/float16")
        logger.info("#" * 48)
        t1 = timeit.default_timer()
        import hdf5plugin
        import h5py as h5
        filename = prefix_output + "pynx_vol.h5"
        f = h5.File(filename, "w")
        f.attrs['creator'] = 'PyNX'
        # f.attrs['NeXus_version'] = '2018.5'  # Should only be used when the NeXus API has written the file
        f.attrs['HDF5_Version'] = h5.version.hdf5_version
        f.attrs['h5py_version'] = h5.version.version
        f.attrs['default'] = 'entry_1'

        entry_1 = f.create_group("entry_1")
        entry_1.attrs['NX_class'] = 'NXentry'
        entry_1.attrs['default'] = 'data_1'
        data_1 = entry_1.create_group("data_1")
        data_1.attrs['NX_class'] = 'NXdata'
        data_1.attrs['signal'] = 'data'
        data_1.attrs['interpretation'] = 'image'
        data_1['title'] = 'PyNX (FBP:Nabu)'
        nz, ny, nx = vol.shape
        data_1.create_dataset("data", data=vol.astype(np.float16), chunks=(1, ny, nx), shuffle=True,
                              **hdf5plugin.LZ4())
        f.close()
        logger.info("Finished saving %s" % filename)

        dt = timeit.default_timer() - t1
        logger.info("Time to save hdf5 volume:  %4.1fs" % dt)
    if "cuts" in save_fbp_vol:
        logger.info(f" Saving 9x9x9 XYZ cuts of volume to: {prefix_output}_volXYZ.h5")
        with h5py.File(f"{prefix_output}_volXYZ.h5", 'w') as h:
            h.attrs['creator'] = 'PyNX'
            # f.attrs['NeXus_version'] = '2018.5'  # Should only be used when the NeXus API has written the file
            h.attrs['HDF5_Version'] = h5py.version.hdf5_version
            h.attrs['h5py_version'] = h5py.version.version
            h.attrs['default'] = 'XY'
            pxum = pixel_size * 1e6

            data = h.create_group("XY")
            data.attrs['NX_class'] = 'NXdata'
            data.attrs['signal'] = 'data'
            data.attrs['interpretation'] = 'image'
            data['title'] = 'XY cuts (from nz/10 to 9*nz/10'
            nz, ny, nx = vol.shape
            data.create_dataset("data", data=vol.take(range(nz // 10, 9 * nz // 10 + 1, nz // 10),
                                                      axis=0).astype(np.float32),
                                chunks=True, compression="gzip")
            data.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
            # Flip: to have origin at top, left corner
            yc = np.arange(ny) * pxum - ny * pxum / 2
            data.create_dataset('row_coords', data=yc)
            data['row_coords'].attrs['units'] = "µm"
            data['row_coords'].attrs['long_name'] = 'Y (µm)'
            xc = np.arange(nx) * pxum - nx * pxum / 2
            data.create_dataset('col_coords', data=xc)
            data['col_coords'].attrs['units'] = "µm"
            data['col_coords'].attrs['long_name'] = 'X (µm)'

            data = h.create_group("YZ")
            data.attrs['NX_class'] = 'NXdata'
            data.attrs['signal'] = 'data'
            data.attrs['interpretation'] = 'image'
            data['title'] = 'YZ cuts (from nx/10 to 9*nx/10'
            nz, ny, nx = vol.shape
            tmp = vol.take(range(nx // 10, 9 * nx // 10 + 1, nx // 10), axis=2).astype(np.float32)
            data.create_dataset("data", data=tmp.swapaxes(0, 2), chunks=True, compression="gzip")
            data.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
            yc = np.arange(nz) * pxum - nz * pxum / 2
            data.create_dataset('row_coords', data=yc)
            data['row_coords'].attrs['units'] = "µm"
            data['row_coords'].attrs['long_name'] = 'Z (µm)'
            xc = np.arange(ny) * pxum - ny * pxum / 2
            data.create_dataset('col_coords', data=xc)
            data['col_coords'].attrs['units'] = "µm"
            data['col_coords'].attrs['long_name'] = 'Y (µm)'

            data = h.create_group("XZ")
            data.attrs['NX_class'] = 'NXdata'
            data.attrs['signal'] = 'data'
            data.attrs['interpretation'] = 'image'
            data['title'] = 'XZ cuts (from ny/10 to 9*ny/10'
            nz, ny, nx = vol.shape
            tmp = vol.take(range(ny // 10, 9 * ny // 10 + 1, ny // 10), axis=1).astype(np.float32)
            data.create_dataset("data", data=tmp.swapaxes(0, 1), chunks=True, compression="gzip")
            data.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
            yc = np.arange(nz) * pxum - nz * pxum / 2
            data.create_dataset('row_coords', data=yc)
            data['row_coords'].attrs['units'] = "µm"
            data['row_coords'].attrs['long_name'] = 'Z (µm)'
            xc = np.arange(nx) * pxum - nx * pxum / 2
            data.create_dataset('col_coords', data=xc)
            data['col_coords'].attrs['units'] = "µm"
            data['col_coords'].attrs['long_name'] = 'X (µm)'

    dt = timeit.default_timer() - t0
    logger.info("Time for 3D FBP & save volume:  %4.1fs" % dt)


def worker_params_id16(qin: Queue, qout: Queue):
    # only to avoid imports triggering cuda init
    if 'id16a' in sys.argv[0]:
        from pynx.holotomo.utils import get_params_id16a as get_params_id16
    else:
        from pynx.holotomo.utils import get_params_id16b as get_params_id16
    for kwargs in iter(qin.get, 'STOP'):
        logger, stream = new_logger_stream("ID16Params")
        try:
            result = get_params_id16(**kwargs, logger=logger)
        except Exception as ex:
            logger.error("Error loading ID16 parameters (most likely: path error or "
                         "missing files)\n" + traceback.format_exc())
            qout.put((None, stream.getvalue()))
            return
        qout.put((result, stream.getvalue()))


def worker_load_align_zoom(qin, qout, iobs_array):
    # Load & align images
    for kwargs in iter(qin.get, 'STOP'):
        logger, stream = new_logger_stream("Load-Align-Zoom")
        try:
            ref, dx, dy = load_align_zoom(**kwargs, iobs_array=iobs_array, logger=logger)
        except Exception as ex:
            logger.error("Error loading data:\n" + traceback.format_exc())
            qout.put((None, None, None, stream.getvalue()))
            return
        qout.put((ref, dx, dy, stream.getvalue()))


def worker_find_cor(qin, qout, ph_array):
    # Center of rotation
    for kwargs in iter(qin.get, 'STOP'):
        logger, stream = new_logger_stream("CORFinder")
        logger.info("Searching for centre of rotation")
        try:
            result = find_cor(**kwargs, logger=logger, ph_array=ph_array)
            logger.info("Returning CoR = %8.2f" % result)
        except Exception:
            logger.error("Failed obtaining center of rotation:\n" + traceback.format_exc())
            result = None
        qout.put((result, stream.getvalue()))


def worker_algorithms(qin, qout, iobs_array, ph_array, profile, igpu):
    # os.environ['PYNX_PU'] = 'cuda.%d' % igpu
    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
        os.environ['CUDA_VISIBLE_DEVICES'] = "%d" % igpu
    else:
        vigpu = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
        os.environ['CUDA_VISIBLE_DEVICES'] = vigpu[igpu]
    from pynx.holotomo.operator import ScaleObjProbe
    from pycuda.driver import MemoryError as cuMemoryError
    pu = None
    gpu_mem_f = 1  # use this to lower GPU memory requirements if we hit a MemoryError
    data, ht = None, None
    for kwargs in iter(qin.get, 'STOP'):
        logger, stream = new_logger_stream("Algorithms%d" % igpu)
        logger.info("#" * 48)
        logger.info("# Algorithms for chunk %d/%d:  [GPU #%d]" % (kwargs["ichunk"] + 1, kwargs['nb_chunk'], igpu))
        logger.info("#" * 48)
        todo = True
        algo = kwargs['algorithm'] if kwargs['algorithm_real'] is None else kwargs['algorithm_real']
        try:
            while todo:
                try:
                    logger.info("Preparing HolotomoData")
                    try:
                        if pu is None:
                            pu = ScaleObjProbe().processing_unit
                            pu.enable_profiling(profiling=profile)
                            pu.set_logger(logger)
                        data = prep_holotomodata(iobs_array=iobs_array, iobs_shape=kwargs['iobs_shape'],
                                                 ref=kwargs['ref'],
                                                 pixel_size=kwargs['pixel_size'], wavelength=kwargs['wavelength'],
                                                 detector_distance=kwargs['detector_distance'],
                                                 dx=kwargs['dx'], dy=kwargs['dy'], proj_idx=kwargs['proj_idx'],
                                                 padding=kwargs['padding'], algorithm=algo,
                                                 gpu_mem_f=gpu_mem_f, pu=pu, logger=logger, data_old=data,
                                                 psf=kwargs['psf'])
                    except Exception as ex:
                        logger.error("Error preparing HolotomoData: " + traceback.format_exc())
                        qout.put("Error preparing HolotomoData")
                        return
                    qout.put("Finished preparing HolotomoData")  # This frees iobs to load the next dataset.
                    logger.info("Beginning algorithms")
                    idx, probe, ht = run_algorithms(data, iobs_shape=kwargs['iobs_shape'], padding=kwargs['padding'],
                                                    delta_beta=kwargs['delta_beta'], algorithm=algo,
                                                    ph_shape=kwargs['ph_shape'], i0=kwargs['i0'],
                                                    nb_chunk=kwargs['nb_chunk'],
                                                    obj_smooth=kwargs['obj_smooth'], obj_inertia=kwargs['obj_inertia'],
                                                    obj_min=kwargs['obj_min'], obj_max=kwargs['obj_max'],
                                                    nb_probe=kwargs['nb_probe'], ph_array=ph_array,
                                                    sino_filter=kwargs['sino_filter'], pu=pu, logger=logger,
                                                    old_ht=ht, return_probe=kwargs['return_probe'],
                                                    normalise=kwargs['normalise'], probe0=kwargs['probe0'],
                                                    magnification=kwargs['magnification'],
                                                    constrain_probe_direct_beam=kwargs['constrain_probe_direct_beam'],
                                                    psf=kwargs['psf'], binning=kwargs['binning'])
                    todo = False
                except cuMemoryError as ex:
                    try:
                        del data
                        data = None
                        gc.collect()
                        gc.collect()
                    except NameError:  # Is 'data' defined here ?
                        pass
                    if gpu_mem_f > 1:
                        logger.warning("cuda MemoryError - again - giving up [GPU#%d]" % igpu)
                        qout.put((None, None, stream.getvalue()))
                        raise ex
                    else:
                        logger.warning("cuda MemoryError - try again, lowering GPU memory usage [GPU#%d]" % igpu)
                        os.system('nvidia-smi')
                        gpu_mem_f = 2
            logger.info("Finished algorithms for chunk %d/%d: sending back unwrapped phases and probe" %
                        (kwargs['ichunk'] + 1, kwargs['nb_chunk']))
            qout.put((idx, probe, stream.getvalue()))
        except Exception as ex:
            s = "Error running algorithms for chunk %d/%d:\n" % (kwargs['ichunk'] + 1, kwargs['nb_chunk'])
            logger.error(s + traceback.format_exc())
            qout.put((None, None, stream.getvalue()))
            return
    if profile:
        pu.enable_profiling(False)
    del data, ht
    gc.collect()
    gc.collect()


def worker_save_projections_hdf5(qin: JoinableQueue, ph_array):
    from pynx.holotomo import save_obj_probe_chunk
    for kwargs in iter(qin.get, 'STOP'):
        result = save_obj_probe_chunk(**kwargs, obj_phase=ph_array)
        qin.task_done()
    qin.task_done()


def worker_save_projections_nxtomo(qin: JoinableQueue, ph_array):
    from pynx.holotomo import save_nxtomo
    for kwargs in iter(qin.get, 'STOP'):
        save_nxtomo(**kwargs, obj_phase=ph_array)
        qin.task_done()
    qin.task_done()


def worker_save_projections_edf(qin: JoinableQueue, ph_array):
    from pynx.holotomo.utils import save_phase_edf_kw
    for kwargs in iter(qin.get, 'STOP'):
        idx, prefix = kwargs['idx'], kwargs['prefix']
        head, tail = os.path.split(prefix)
        os.makedirs(prefix + '_edf', exist_ok=True)
        prefix = f"{prefix}_edf/{tail}"
        ph = np.frombuffer(ph_array.get_obj(), dtype=np.float16).reshape(kwargs['ph_shape'])
        vkw = [{'idx': idx[i], 'ph': ph[i], 'prefix_output': prefix} for i in range(len(idx))]
        with Pool(nproc) as pool:
            pool.map(save_phase_edf_kw, vkw)
        qin.task_done()
    qin.task_done()


def worker_fbp(qin, qout, ph_array, profile):
    # Filtered back-projection CT with Nabu
    for kwargs in iter(qin.get, 'STOP'):
        logger, stream = new_logger_stream("Tomo (Nabu FBP)")
        try:
            run_fbp(**kwargs, ph_array=ph_array, logger=logger, profile=profile)
        except Exception:
            s = "Error running FBP and exporting volume or slices:"
            logger.error(s + traceback.format_exc())

        qout.put(stream.getvalue())


def stop_all(vp, vq, logfile):
    print("Stopping execution following an I/O or processing error")
    for q in vq:
        q.put('STOP')
    for p in vp:
        print("Stopping process:", p)
        p.terminate()
        print("Stopped process:", p)
    new_logger_stream("Master", flush_to=logfile)
    sys.exit(1)


def master(params):
    from tomoscan.esrf.scan.edfscan import EDFTomoScan
    from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
    # Setup logging
    DeltaTimeFormatter.start_time = time.time()
    logger, stream = new_logger_stream("Master")
    logger.info("Host: %s - %s" % (platform.node(), platform.platform()))
    logger.info("Working directory: %s" % os.getcwd())
    try:
        logger.info("User: %s " % os.getlogin())
    except OSError:
        # When using slurm, this is not running in a shell with user info
        if 'USER' in os.environ:
            logger.info("User: %s " % os.environ['USER'])
    for k in ['SLURM_GPUS_ON_NODE', 'SLURM_JOB_ACCOUNT', 'SLURM_JOB_GPUS', 'SLURM_JOB_ID',
              'SLURM_JOB_NAME', 'SLURM_JOB_PARTITION', 'SLURM_CPUS_ON_NODE', 'SLURM_NTASKS',
              'SLURM_MEM_PER_NODE']:
        if k in os.environ:
            logger.info("%s: %s " % (k, os.environ[k]))
    logger.info("#" * 48)
    logger.info("# Parsing parameters")
    logger.info("#" * 48)

    c = sys.argv[0]
    for a in sys.argv[1:]:
        c += " %s" % a
    logger.info("Command: %s" % c)
    logger.info("Parameters:")
    vp = vars(params)
    vk = list(vp.keys())
    vk.sort()
    params_dic = {}  # Later, for export
    for k in vk:
        v = vp[k]
        if k == 'data':
            if re.search('%[0-9]*?d', params.data) is None:  # Matches %d, %2d, %02d
                data0 = params.data
                params.data = params.data[:-2] + '%d_'
                v = params.data
                logger.info("     %s: %s -> %s" % (k, data0, str(v)))
            else:
                logger.info("     %s: %s" % (k, str(v)))
        else:
            logger.info("     %s: %s" % (k, str(v)))
        params_dic[k] = v

    if params.dry_run:
        print("\n DRY RUN - exiting")
        sys.exit(0)

    if 'CUDA_VISIBLE_DEVICES' in os.environ:
        ngpu_env = len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))
        if ngpu_env < params.ngpu:
            raise RuntimeError("ngpu=%d have been required but CUDA_VISIBLE_DEVICES=%s" %
                               (params.ngpu, os.environ['CUDA_VISIBLE_DEVICES']))

    # Analysing data name
    data_dir, prefix = os.path.split(params.data)
    prefix = prefix[:-2]  # Remove the trailing '1_'

    # Create logfile
    path = os.path.split(params.prefix_output)[0]
    if len(path):
        os.makedirs(path, exist_ok=True)
    logfile = open("%s.log" % params.prefix_output, 'w')

    # Get projection urls for all planes
    data_is_nx = True if params.data.endswith('.nx') else False
    proj_urls = {}
    if data_is_nx:
        for p in params.planes:
            nx_name = params.data % (p - 1)
            s = NXtomoScan(nx_name)
            # Pass a string which can be serialised
            # Start the number of projection frames at 0
            proj_urls[p] = {k: f"{v.file_path()}:{v.data_path()}:{v.data_slice()}"
                            for k, v in enumerate(s.projections.values())}
    else:
        for p in params.planes:
            s = EDFTomoScan.get_proj_urls(params.data % p, n_frames=1)
            proj_urls[p] = {k: f"{v.file_path()}" for k, v in s.items()}
        # proj_urls = {p: EDFTomoScan.get_proj_urls(params.data % p, n_frames=1) for p in params.planes}

    # Correct range for projection
    if params.last is None:
        if data_is_nx:
            nb_proj = len(proj_urls[params.planes[0]])
        else:
            # TODO: why does nb_proj from get_tomo_n differ from the number of keys ?
            nb_proj = EDFTomoScan.get_tomo_n(params.data % params.planes[0])
        logger.info("No projection_range given - using all %d projections" % nb_proj)
        # Assume all distances have the same number of projections
        proj_idx = list(proj_urls[params.planes[0]].keys())
        proj_idx.sort()
        proj_idx = proj_idx[:nb_proj]  # Remove extra projections for alignment.
        proj_idx = np.array(proj_idx[params.first::params.step])
        logger.info("Projections: [%d, %d ...%d]" % (proj_idx[0], proj_idx[1], proj_idx[-1]))
    else:
        logger.info("projection_range: [%d-%d::%d]" % (params.first, params.last, params.step))
        proj_idx = np.arange(params.first, params.last + 1, params.step)

    # Update the projection URLs - now should be ordered
    proj_urls = {p: {k: proj_urls[p][k] for k in proj_idx} for p in params.planes}

    nb_proj = len(proj_idx)  # number of images loaded (excluding dark and empty_beam images)
    logger.info("nb_proj=%d" % nb_proj)

    logger, stream = new_logger_stream("Master", flush_to=logfile)

    # Gather magnified pixel sizes and propagation distances
    # This is done in a separate process to avoid imports triggering cuda init
    # ...though this should not be necessary !
    kwargs = {"scan_path": params.data, "planes": params.planes, "sx0": params.sx0, "cx": params.cx, "sx": params.sx,
              "pixel_size_detector": params.pixel_size_detector}

    queue_params_id16_in = Queue()
    queue_params_id16_out = Queue()
    p_params = Process(target=worker_params_id16, args=(queue_params_id16_in, queue_params_id16_out))
    p_params.start()
    logger.info("ID16 parameters worker: " + str(p_params))
    queue_params_id16_in.put(kwargs)
    params_id16, log = queue_params_id16_out.get()
    if params_id16 is None:
        # Something went wrong retrieving parameters
        raise RuntimeError("Error retrieving experiment parameters from data files - aborting")
    queue_params_id16_in.put('STOP')

    logfile.write(log)
    logfile.flush()
    ny, nx = params_id16['ny'] // params.binning, params_id16['nx'] // params.binning
    nz = params.nz

    if params.nrj is None:
        wavelength = params_id16['wavelength']
    else:
        wavelength = 12.3984 / params.nrj
    logger.info("Energy: %6.2f keV" % (12.3984e-10 / wavelength))
    logger.info("Wavelength: %8.6e m" % wavelength)

    if params.distance is None:
        detector_distance = params_id16['distance']
    else:
        detector_distance = np.array(params.distance)

    if params.pixel_size is None:
        pixel_size_z = params_id16['pixel_size']
    else:
        pixel_size_z = np.array(params.pixel_size)
    pixel_size_z *= params.binning
    pixel_size = pixel_size_z[0]

    magnification = pixel_size_z[0] / pixel_size_z

    logger.info("Detector distances [m]: %s" % str(detector_distance))
    logger.info("Pixel sizes [nm]: %s" % str(pixel_size_z * 1e9))

    # Padding:
    if params.padding == 2:
        pady, padx = ny // 2, nx // 2
        logger.info("Padding==2 => using pady=%d  padx=%d" % (pady, padx))
    else:
        pady = padx = params.padding

    # Tomography angle parameters
    tomo_angle_step_orig = params_id16['scan_range'] / (params_id16['tomo_n'] - 1)
    # Shift the angle so the reconstructions are aligned, regardless of the
    # first projection used
    tomo_angle_start = tomo_angle_step_orig * params.first
    # This will correct the step due to the binning and range used
    tomo_angle_step = tomo_angle_step_orig * (proj_idx.max() - proj_idx.min()) / (len(proj_idx) - 1)

    ################################################################
    # Test for radix transforms
    ################################################################
    # Test if we have a radix transform
    logger.info("Testing dimensions prime decomposition: ny = %d = %s  nx= %d = %s" %
                (ny + 2 * pady, str(primes(ny)),
                 nx + 2 * padx, str(primes(nx))))
    primesy, primesx = primes(ny + 2 * pady), primes(nx + 2 * padx)
    if max(primesy) > 13:
        padup = pady
        while not test_smaller_primes(ny + 2 * padup, required_dividers=[2]):
            padup += 1
        paddown = pady
        while not test_smaller_primes(ny + 2 * paddown, required_dividers=[2]):
            paddown -= 1
        s = "The Y dimension (with padding=%d,%d) is incompatible with a radix FFT:\n" \
            "  ny=%d primes=%s  (should be <=13)\n" \
            "  Closest acceptable padding values: %d or %d" % \
            (pady, padx, ny + 2 * pady, str(primesy), paddown, padup)
        logger.warning(s)
        # Use closest padding option, unless we pad to 2x the array size, in which case
        # we can only go down
        if (padup - pady < pady - paddown) and (params.padding != 2):
            pady = padup
        else:
            pady = paddown
        logger.info("Changing Y padding to closest radix-compatible value: %d" % pady)

    if max(primesx) > 13:
        padup = padx
        while not test_smaller_primes(nx + 2 * padup, required_dividers=[2]):
            padup += 1
        paddown = padx
        while not test_smaller_primes(nx + 2 * paddown, required_dividers=[2]):
            paddown -= 1
        s = "The X dimension (with padding=%d) is incompatible with a radix FFT:\n" \
            "  nx=%d primes=%s  (should be <=13)\n" \
            "  Closest acceptable padding values: %d or %d" % \
            (padx, nx + 2 * padx, str(primesx), paddown, padup)
        logger.warning(s)
        # Use the closest padding option, unless we pad to 2x the array size, in which case
        # we can only go down
        if (padup - padx < padx - paddown) and (params.padding != 2):
            padx = padup
        else:
            padx = paddown
        logger.info("Changing X padding to closest radix-compatible value: %d" % padx)
    # Add real padding to params
    setattr(params, 'padding_real', (pady, padx))
    # logger, stream = new_logger_stream("Master", flush_to=logfile)
    # raise RuntimeError(s)

    ################################################################
    # Estimate memory requirements & number of chunks
    ################################################################
    # Try to:
    # * Use a reasonable amount of memory (don't hit shared memory limits)
    # * at least 2 chunks per GPU (so CoR can be determined from half projections)
    # * if possible, use a single stack in GPU memory (much faster iterations)
    mem_req = nz * (ny + 2 * pady) * (nx + 2 * padx) * 4
    mem_req += 12 * (nx + 2 * pady) * (ny + 2 * padx)  # object projections & phase
    if True:  # 'RAAR' in algorithm or 'DM' in algorithm or 'DRAP' in algorithm:
        mem_req += 8 * nz * (nx + 2 * pady) * (ny + 2 * padx)
    logger.info("Estimated memory requirement: %8.4fGB/projection (total = %8.4fGB)" %
                (mem_req / 1024 ** 3, mem_req * nb_proj / 1024 ** 3))
    # Maximum memory used for a single chunk of projections (not counting final result)
    max_mem = params.max_mem_chunk * 1024 ** 3
    nb_chunk = int(np.ceil(nb_proj * mem_req / max_mem))
    nb_chunk += nb_chunk % params.ngpu

    # See if we can use a single stack in GPU memory
    gpu_mem = get_gpu_total_mem()
    logger.info(f"Available GPU memory: {gpu_mem:.1f}GB")
    mem_stack = 4 * nz * (nx + 2 * padx) * (ny + 2 * pady)  # Iobs
    mem_stack += 8 * nz * (nx + 2 * padx) * (ny + 2 * pady)  # Psi
    mem_stack += 12 * (nx + 2 * padx) * (ny + 2 * pady)  # object projections & phase
    if 'raar' in params.algorithm.lower() or 'dm' in params.algorithm.lower() \
            or 'drap' in params.algorithm.lower():
        mem_stack += 8 * nz * (nx + 2 * padx) * (ny + 2 * pady)  # Psi copy
    if params.psf is not None:
        if params.psf > 0:
            mem_stack += 4 * nz * (nx + 2 * padx + 2) * (ny + 2 * pady)
    # pyvkfft may require an extra buffer for inplace transforms, for large arrays
    mem_fft = nz * get_vkfft_buffer_nbytes(ny + 2 * pady, nx + 2 * padx)
    # Maximum stack size fitting in GPU memory
    gpu_mem_margin_fact = 1.3
    stack_size = int(np.ceil((gpu_mem - 0.5) * 1024 ** 3 / (mem_stack * gpu_mem_margin_fact + mem_fft)))
    logger.info(f"Max stack size per GPU: {stack_size} [{mem_stack * gpu_mem_margin_fact / 1024 ** 3:.3f}GB/proj]")
    if False:  # TODO stack_size >= 100:
        # if we have at least 100 projections per chunk, it should be enough
        # to correctly average the probe
        nb_chunk = max(nb_chunk, int(np.ceil(nb_proj / stack_size)))
        # We need a number of chunk proportional to the number of GPUs
        nb_chunk += nb_chunk % params.ngpu
        single_stack = True
    else:
        single_stack = False

    if nb_chunk == params.ngpu:
        nb_chunk *= 2  # This allows for the CoR to be found while the phasing finishes
    nb_proj_chunk = len(proj_idx[::nb_chunk])
    logger.info(f"Using {nb_chunk} chunk(s) [{nb_proj_chunk} projections/chunk]")
    if single_stack:
        m = (mem_stack * gpu_mem_margin_fact + mem_fft) / 1024 ** 3 * nb_proj_chunk + 0.5
        logger.info(f"Will use a single stack in each GPU. "
                    f"Over-estimated GPU memory usage: {m:.1f} GB")

    ################################################################
    # Create shared memory arrays before forking process
    # NB: we use multiprocessing.Array rather than shared_memory
    # as the latter runs into issues (Bus error) for large sizes (around 32GB)
    # on power9 (bug ?)
    ################################################################
    logger.info("#" * 48)
    logger.info("# Creating shared memory arrays")
    logger.info("#" * 48)
    t0 = timeit.default_timer()

    # Create shared memory array for phases
    ph_shape = (nb_proj, ny, nx)
    logger.info("Phase array shape: %s [%6.2fGB]" % (str(ph_shape), np.prod(ph_shape) * 2 / 1024 ** 3))
    ph_array_dtype = np.float16
    ph_array = Array(ctypes.c_int16, int(np.prod(ph_shape)))  # There is no ctypes.c_float16
    ph = np.frombuffer(ph_array.get_obj(), dtype=ph_array_dtype).reshape(ph_shape)

    ################################################################
    # Setup job queues & worker for center of rotation
    ################################################################
    vp, vq = [], []  # global list of process & input queues
    # Center of rotation
    queue_cor_in = Queue()
    queue_cor_out = Queue()
    p_cor = Process(target=worker_find_cor, args=(queue_cor_in, queue_cor_out, ph_array))
    p_cor.start()
    logger.info("Center of rotation worker: " + str(p_cor))
    vp.append(p_cor)
    vq.append(queue_cor_in)
    tomo_rot_center = params.tomo_rot_center

    ################################################################
    # Optionally - reconstruct volume from holoct projections,
    # and exit
    ################################################################
    if params.holoct is not None:
        # Export same files for holoct files, assuming we can find them
        # FBP
        queue_fbp_in = Queue()
        queue_fbp_out = Queue()
        p_fbp = Process(target=worker_fbp, args=(queue_fbp_in, queue_fbp_out, ph_array, params.profile))
        p_fbp.start()
        logger.info("FBP worker: " + str(p_fbp))
        logger.info("Processing volume from holoct reconstructed projections...")
        from ..utils import load_holoct_data_pool_init, load_holoct_data_kw
        import fabio
        # Check we can find holoct files and load them into the phase array
        if not os.path.exists(params.holoct % proj_idx[0]):
            raise RuntimeError("HoloCT volume calculation: could not find %s" % (params.holoct % proj_idx[0]))
        # Use a fake dark array to re-use loading functions
        vkw = [{'i': i, 'idx': proj_idx[i], 'nproj': nb_proj, 'nx': nx, 'ny': ny,
                'img_name': params.holoct, 'binning': params.binning,
                'ph_array_dtype': ph_array_dtype}
               for i in range(nb_proj)]
        logger.info("Loading holoct reconstructed projections...")
        t0 = timeit.default_timer()
        with Pool(nproc, initializer=load_holoct_data_pool_init, initargs=(ph_array,)) as pool:
            pool.map(load_holoct_data_kw, vkw)
        dt = timeit.default_timer() - t0
        logger.info("Loading holoct reconstructed projections... Finished [%6.1fMB/s]"
                    % (np.prod(ph_shape) * 2 * params.binning ** 2 / 1024 ** 2 / dt))
        # Check if a sino filter needs to be applied
        edf = fabio.open(params.holoct % proj_idx[0])
        sino_filter = 'ram-lak'
        if 'post_filter' in edf.header:
            if 'sino' in edf.header['post_filter']:
                sino_filter = None

        if tomo_rot_center is None:
            izref = list(params.planes).index(params.reference_plane)
            logger.info("Evaluating rotation centre using: %s (method=%s)" % (params.data, params.tomo_cor_method))
            queue_cor_in.put({"data_src": params.data, "binning": params.binning,
                              "reference_plane": params.reference_plane, "method": params.tomo_cor_method,
                              "motion": False, "side": params.halftomo,
                              "ph_shape": ph_shape, "ph_array_dtype": ph_array_dtype, "proj_idx": proj_idx,
                              "reference_plane_magnification": magnification[0] / magnification[izref],
                              "tomo_rot_center_near": params.tomo_rot_center_near})

            tomo_rot_center, log = queue_cor_out.get()
            logger.info("Got back CoR: %6.1f" % tomo_rot_center)
            logger, stream = new_logger_stream("Master", flush_to=logfile)
            logfile.write(log)
            logfile.flush()
        else:
            params.tomo_cor_method = "manual"

        # Perform FBP exactly as before using holoCT data
        logger.info("Launching FBP reconstruction from holoct reconstructed projections...")
        kwargs = {"data_src": params.data, "reference_plane": params.reference_plane,
                  "prefix_output": params.prefix_output + "_holoct",
                  "ph_shape": ph_shape, "idx": proj_idx,
                  "tomo_rot_center": tomo_rot_center, "pixel_size": pixel_size,
                  "tomo_angle_step": tomo_angle_step, "tomo_angle_start": tomo_angle_start,
                  "save_fbp_vol": params.save_fbp_vol, "save_3sino": params.save_3sino,
                  "ngpu": params.ngpu, "params": params, "sino_filter": sino_filter,
                  "halftomo": params.halftomo}
        queue_fbp_in.put(kwargs)
        log = queue_fbp_out.get()
        logger, stream = new_logger_stream("Master", flush_to=logfile)
        logfile.write(log)
        logfile.flush()
        # Flush logger
        logger, stream = new_logger_stream("Master", flush_to=logfile)
        logger.info("Finished FBP reconstruction from holoct reconstructed projections - exiting")
        for p, q in [(p_cor, queue_cor_in), (p_fbp, queue_fbp_in)]:
            q.put('STOP')
            p.join()
        sys.exit(0)

    # Create shared memory array for iobs
    proj_idx_chunk = proj_idx[::nb_chunk]
    iobs_shape = (len(proj_idx_chunk), nz, ny + 2 * pady, nx + 2 * padx)
    logger.info("Iobs shape: %s [%6.2fGB]" % (str(iobs_shape), np.prod(iobs_shape) * 4 / 1024 ** 3))
    iobs_array = Array(ctypes.c_float, int(np.prod(iobs_shape)))

    dt = timeit.default_timer() - t0
    logger.info("Finished creating shared memory arrays: dt = %6.2fs" % dt)
    shared_nbytess = ph.nbytes + np.prod(iobs_shape) * 4
    logger.info(f"Total size of shared memory arrays: {shared_nbytess / 1024 ** 3:.1f} GB")
    ################################################################
    # Setup job other queues & workers
    ################################################################
    queue_load_align_in = Queue()
    queue_load_align_out = Queue()
    p_load_align = Process(target=worker_load_align_zoom,
                           args=(queue_load_align_in, queue_load_align_out, iobs_array))
    logger.info("Load/align worker: " + str(p_load_align))
    p_load_align.start()
    vp.append(p_load_align)
    vq.append(queue_load_align_in)

    # Algorithms for multiple GPUs
    queue_algo_in, queue_algo_out, p_algo = [], [], []
    for igpu in range(params.ngpu):
        queue_algo_in.append(Queue())
        queue_algo_out.append(Queue())
        p_algo.append(Process(target=worker_algorithms, args=(queue_algo_in[-1], queue_algo_out[-1], iobs_array,
                                                              ph_array, params.profile, igpu)))
        p_algo[-1].start()
        logger.info(f"Algorithm worker [GPU #{igpu}]: " + str(p_algo[-1]))
    vp += p_algo
    vq += queue_algo_in

    # Save projections to hdf5 file
    queue_save_projections_hdf5_in = JoinableQueue()
    p_phases2hdf5 = Process(target=worker_save_projections_hdf5, args=(queue_save_projections_hdf5_in, ph_array))
    p_phases2hdf5.start()
    logger.info(f"Phases -> hdf5 worker: " + str(p_phases2hdf5))
    vp.append(p_phases2hdf5)
    vq.append(queue_save_projections_hdf5_in)

    # Save projections to NXtomo file
    queue_save_projections_nxtomo_in = JoinableQueue()
    p_phases2nxtomo = Process(target=worker_save_projections_nxtomo, args=(queue_save_projections_nxtomo_in, ph_array))
    p_phases2nxtomo.start()
    logger.info(f"Phases -> nxtomo worker: " + str(p_phases2nxtomo))
    vp.append(p_phases2nxtomo)
    vq.append(queue_save_projections_nxtomo_in)

    # Save projections to edf files
    queue_save_projections_edf_in = JoinableQueue()
    p_phases2edf = Process(target=worker_save_projections_edf, args=(queue_save_projections_edf_in, ph_array))
    p_phases2edf.start()
    logger.info(f"Phases -> edf worker: " + str(p_phases2edf))
    vp.append(p_phases2edf)
    vq.append(queue_save_projections_edf_in)

    # FBP
    queue_fbp_in = Queue()
    queue_fbp_out = Queue()
    p_fbp = Process(target=worker_fbp, args=(queue_fbp_in, queue_fbp_out, ph_array, params.profile))
    p_fbp.start()
    logger.info(f"FBP worker: " + str(p_fbp))
    vp.append(p_fbp)
    vq.append(queue_fbp_in)

    ################################################################
    # Distribute tasks for chunks
    ################################################################
    # Distribute chunks so that all even frames are done first, and a CoR search
    # can be performed using that half sinogram
    v_i0 = list(range(0, nb_chunk, 2))
    if nb_chunk > 1:
        v_i0 += list(range(1, nb_chunk, 2))
    # Launch first input job
    proj_idx_chunk = proj_idx[::nb_chunk]
    iobs_shape = (len(proj_idx_chunk), nz, ny + 2 * pady, nx + 2 * padx)

    kwargs = {"ichunk": 0, "data_src": params.data, "prefix": prefix, "proj_idx": proj_idx_chunk,
              "planes": params.planes, "binning": params.binning, "magnification": magnification,
              "padding": (pady, padx), "pad_method": params.pad_method, "rhapp": params.rhapp,
              "reference_plane": params.reference_plane, "spikes_threshold": params.remove_spikes,
              "align_method": params.align, "align_interval": params.align_interval,
              "align_fourier_low_cutoff": params.align_fourier_low_cutoff,
              "align_fourier_high_cutoff": params.align_fourier_high_cutoff,
              "prefix_output": params.prefix_output, "nb_chunk": nb_chunk,
              "motion": params.motion, "proj_urls": proj_urls, "distorted": params.distorted,
              "normalise": params.normalise, "align_max_shift": params.align_max_shift,
              "align_diff": "diff" in params.align, "vsnr": params.vsnr,
              "params": params}
    logger.info("Requesting Iobs data for chunk #1/%d" % nb_chunk)
    queue_load_align_in.put(kwargs)

    for ichunk in range(nb_chunk):
        i0 = v_i0[ichunk]
        # Projections for this chunk
        proj_idx_chunk = proj_idx[i0::nb_chunk]
        # Wait for data (iobs comes from shared memory)
        ref, dx, dy, log = queue_load_align_out.get()
        # GPU for this chunk
        igpu = ichunk % params.ngpu

        logger, stream = new_logger_stream("Master", flush_to=logfile)
        logfile.write(log)
        logfile.flush()
        if ref is None:
            # This only happens if something went wrong loading data
            logger.error(f"Load / Align step failed for chunk {ichunk + 1}")
            stop_all(vp, vq, logfile)
        logger.info("Received Iobs and reference frames for chunk #%d/%d" % (ichunk + 1, nb_chunk))
        # Start algorithms
        # TODO: see if sino_filter can be applied on padded array with half-tomo ?
        kwargs = {"ichunk": ichunk, "iobs_shape": iobs_shape, "ref": ref,
                  "pixel_size": pixel_size, "wavelength": wavelength,
                  "detector_distance": detector_distance, "dx": dx, "dy": dy, "proj_idx": proj_idx_chunk,
                  "padding": (pady, padx), "delta_beta": params.delta_beta,
                  "algorithm": params.algorithm, "algorithm_real": params.algorithm_real,
                  "ph_shape": ph_shape, "i0": i0, "nb_chunk": nb_chunk,
                  "obj_smooth": params.obj_smooth, "obj_inertia": params.obj_inertia,
                  "obj_min": params.obj_min, "obj_max": params.obj_max, "nb_probe": params.nb_probe,
                  # "sino_filter": params.sino_filter if padx and (params.halftomo is None) else None,
                  "sino_filter": params.sino_filter if padx else None,
                  "halftomo": params.halftomo,
                  "return_probe": params.save_phase_chunks, "normalise": params.normalise,
                  "probe0": params.probe, "magnification": magnification,
                  "constrain_probe_direct_beam": params.constrain_probe_direct_beam,
                  "psf": params.psf, "binning": params.binning}
        logger.info("Launching algorithms for chunk #%d/%d" % (ichunk + 1, nb_chunk))
        queue_algo_in[igpu].put(kwargs)
        # Get first message that HolotomoData is ready and the next set of data can be read
        res = queue_algo_out[igpu].get()
        if "Error" in res:
            stop_all(vp, vq, logfile)
        # Read next set of data now that iobs data is free
        if ichunk < (nb_chunk - 1):
            proj_idx_chunk_next = proj_idx[v_i0[ichunk + 1]::nb_chunk]
            iobs_shape = (len(proj_idx_chunk_next), nz, ny + 2 * pady, nx + 2 * padx)

            kwargs = {"ichunk": ichunk + 1, "data_src": params.data, "prefix": prefix, "proj_idx": proj_idx_chunk_next,
                      "planes": params.planes, "binning": params.binning, "magnification": magnification,
                      "padding": (pady, padx), "pad_method": params.pad_method, "rhapp": params.rhapp,
                      "reference_plane": params.reference_plane, "spikes_threshold": params.remove_spikes,
                      "align_method": params.align, "align_interval": params.align_interval,
                      "align_fourier_low_cutoff": params.align_fourier_low_cutoff,
                      "align_fourier_high_cutoff": params.align_fourier_high_cutoff,
                      "prefix_output": params.prefix_output, "nb_chunk": nb_chunk,
                      "motion": params.motion, "proj_urls": proj_urls,
                      "distorted": params.distorted, "normalise": params.normalise,
                      "align_max_shift": params.align_max_shift, "vsnr": params.vsnr,
                      "align_diff": "diff" in params.align, "params": params}
            logger.info("Requesting Iobs data for chunk #%d/%d" % (ichunk + 2, nb_chunk))
            queue_load_align_in.put(kwargs)
        if (ichunk + 1) >= params.ngpu:
            # Get results (phases are in shared memory) for chunk ichunk - ngpu +1
            logger.info("Wait for phases and probe for chunk %d/%d" % (ichunk + 1 - params.ngpu + 1, nb_chunk))
            idx, probe, log = queue_algo_out[(ichunk - params.ngpu + 1) % params.ngpu].get()
            logfile.write(log)
            logfile.flush()
            if idx is None:
                # Something went wrong - aborting
                logger.error("Something went wrong during the algorithms stage - aborting")
                stop_all(vp, vq, logfile)

            if not np.allclose(idx, proj_idx[v_i0[ichunk - params.ngpu + 1]::nb_chunk]):
                logger.error("Whoops: idx != proj_idx_chunk:\n" + str(idx) + str(proj_idx_chunk))
            logger, stream = new_logger_stream("Master", flush_to=logfile)
            logger.info("Received unwrapped phases and probe for chunk %d/%d" %
                        (ichunk + 1 - params.ngpu + 1, nb_chunk))
            if tomo_rot_center is None:
                if (ichunk - params.ngpu + 2 == nb_chunk // 2) or (nb_chunk == 1):
                    logger.info("Evaluating rotation centre using: %s (method=%s)" %
                                (params.data, params.tomo_cor_method))
                    izref = list(params.planes).index(params.reference_plane)
                    queue_cor_in.put({"data_src": params.data, "binning": params.binning,
                                      "reference_plane": params.reference_plane,
                                      "method": params.tomo_cor_method, "motion": params.motion,
                                      "side": params.halftomo, "ph_shape": ph_shape,
                                      "ph_array_dtype": ph_array_dtype, "proj_idx": proj_idx,
                                      "reference_plane_magnification": magnification[0] / magnification[izref],
                                      "tomo_rot_center_near": params.tomo_rot_center_near})
            else:
                params.tomo_cor_method = "manual"

    # Get results for the last (ngpu-1) chunks
    if params.ngpu > 1:
        for ichunk in range(nb_chunk - params.ngpu + 1, nb_chunk):
            idx, probe, log = queue_algo_out[ichunk % params.ngpu].get()
            logfile.write(log)
            logfile.flush()
            if idx is None:
                # Something went wrong - aborting
                stop_all(vp, vq)

            if not np.allclose(idx, proj_idx[v_i0[ichunk]::nb_chunk]):
                print("Whoops: idx != proj_idx_chunk:")
                print(idx)
                print(proj_idx_chunk)
            logger, stream = new_logger_stream("Master", flush_to=logfile)
            logger.info("Received unwrapped phases and probe for chunk %d/%d" % (ichunk + 1, nb_chunk))

    tmp = ",".join(["%d" % p.pid for p in p_algo])
    logger.info("Joining load[%d] & algorithm[%s] processes" % (p_load_align.pid, tmp))
    for p, q in [(p_load_align, queue_load_align_in)] + [(p_algo[i], queue_algo_in[i]) for i in range(params.ngpu)]:
        q.put('STOP')
        p.join()
        vp.remove(p)
        vq.remove(q)

    # Free some shared memory before FBP
    del iobs_array
    gc.collect()

    # Save projections
    if params.save_phase_chunks:
        fname = params.prefix_output + "_i0=%04d.h5" % proj_idx[0]
        logger.info("Saving holotomo projections to: %s" % fname)
        kwargs = {"filename": fname, "idx": idx, "pixel_size": pixel_size,
                  "obj_phase_shape": ph_shape,
                  "probe": probe, "process_parameters": params_dic}
        queue_save_projections_hdf5_in.put(kwargs)

    if params.nxtomo:
        fname = params.prefix_output + ".nx"
        logger.info("Saving holotomo projections to: %s" % fname)
        logger.info("Angular step for nxtomo: %8.3f" % tomo_angle_step)
        kwargs = {"filename": fname, "pixel_size": pixel_size,
                  "obj_phase_shape": ph_shape, "chunk_idx": idx[0],
                  "probe": probe, "process_parameters": params_dic,
                  "angles": np.deg2rad(np.arange(nb_proj) * tomo_angle_step) + tomo_angle_start,
                  "sample_name": prefix, "x_rotation_axis": tomo_rot_center,
                  "detector_distance": detector_distance,
                  "wavelength": wavelength, "append": False}
        queue_save_projections_nxtomo_in.put(kwargs)

    if params.save_edf:
        logger.info("#" * 48)
        logger.info(" Saving phased images to edf files: " + params.prefix_output + "_%04d.edf")
        logger.info("#" * 48)
        kwargs = {"idx": proj_idx, "ph_shape": ph_shape, "prefix": params.prefix_output}
        queue_save_projections_edf_in.put(kwargs)

    # Get back rotation centre if needed
    if params.tomo_rot_center is None:
        logger.info("Getting back CoR estimations")
        tomo_rot_center, log = queue_cor_out.get()
        logger, stream = new_logger_stream("Master", flush_to=logfile)
        logfile.write(log)
        logfile.flush()
        if tomo_rot_center is None:
            # This only happens if something went wrong getting the CoR
            if params.save_fbp_vol:
                logger.error(f"Failed determining the center of rotation - aborting")
                stop_all(vp, vq, logfile)
            else:
                logger.error(f"Failed determining the center of rotation - "
                             f"since we are not saving the FBP volume, let's continue")
        else:
            logger.info("Got back CoR: %6.1f" % tomo_rot_center)

    # Run FBP & export reconstructed volume
    if params.save_fbp_vol:
        sino_filter = params.sino_filter
        if sino_filter is not None:
            if padx:  # and params.halftomo is None:
                # If sino_filter was supplied and padding is used, then the pre-backprojection
                # filtering was performed when getting the unwrapped phases at the end of the algorithms
                # TODO: make this possible also with half-tomo ?
                sino_filter = None
        else:
            # No arguments where given, so perform a standard FBP filtering
            sino_filter = "ram-lak"

        kwargs = {"data_src": params.data, "reference_plane": params.reference_plane,
                  "prefix_output": params.prefix_output,
                  "ph_shape": ph_shape, "idx": proj_idx,
                  "tomo_rot_center": tomo_rot_center / params.binning, "pixel_size": pixel_size,
                  "tomo_angle_step": tomo_angle_step, "tomo_angle_start": tomo_angle_start,
                  "save_fbp_vol": params.save_fbp_vol, "save_3sino": params.save_3sino,
                  "ngpu": params.ngpu, "params": params, "sino_filter": sino_filter,
                  "halftomo": params.halftomo}
        queue_fbp_in.put(kwargs)
        log = queue_fbp_out.get()
        logger, stream = new_logger_stream("Master", flush_to=logfile)
        logfile.write(log)
        logfile.flush()

    logger.info("Waiting for all parallel processes to stop...")
    for p, q in [(p_cor, queue_cor_in), (p_fbp, queue_fbp_in),
                 (p_phases2hdf5, queue_save_projections_hdf5_in),
                 (p_phases2nxtomo, queue_save_projections_nxtomo_in),
                 (p_phases2edf, queue_save_projections_edf_in)]:
        q.put('STOP')
        p.join()
    for q in [queue_save_projections_hdf5_in, queue_save_projections_edf_in, queue_save_projections_nxtomo_in]:
        q.join()
    logger.info("            all parallel processes have returned")

    # Flush logger
    logger, stream = new_logger_stream("Master", flush_to=logfile)


algo_aliases = {'paganin-norm': 'AP**5,delta_beta=0,AP**10,Paganin',
                'paganin-short': 'AP**5,obj_smooth=0,delta_beta=0,AP**10,Paganin,obj_smooth=4',
                'paganin-long': 'AP**20,obj_smooth=0,delta_beta=0,AP**40,Paganin,obj_smooth=4',
                'paganin-long-dm':
                    'AP**10,obj_smooth=0,delta_beta=0,AP**10,DM**40,AP**5,dm_alpha=0.1,obj_smooth=4,Paganin',
                'ctf-norm': 'AP**5,delta_beta=0,AP**10,CTF',
                'ctf-short': 'AP**5,obj_smooth=0,delta_beta=0,AP**10,CTF,obj_smooth=4',
                'ctf-long': 'AP**20,obj_smooth=0,delta_beta=0,AP**40,CTF,obj_smooth=4',
                'ctf-long-dm':
                    'AP**10,obj_smooth=0,delta_beta=0,AP**10,DM**40,AP**5,dm_alpha=0.1,obj_smooth=4,CTF',
                }


def make_parser(instrument=None):
    epilog_id16b = """Example usage:

* quick (binned x2) 4-distance reconstruction, Paganin, save tiff volume:
   ``pynx-holotomo-id16b --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff \
--algorithm AP**5,delta_beta=0,AP**10,Paganin --nz 4 --padding 100 --binning 2 --slurm``

* 4-distance reconstruction, CTF, save tiff volume, distortion correction:
   ``pynx-holotomo-id16b --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff \
--algorithm AP**5,delta_beta=0,AP**20,CTF --nz 4 --padding 200 --distorted --slurm``

* 4-distance reconstruction, CTF, save tiff volume, distortion correction, \
use normalised projections (disables probe optimisation):
   ``pynx-holotomo-id16b --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff \
--algorithm AP**5,delta_beta=0,AP**20,CTF --nz 4 --padding 200 --distorted --normalise --slurm``

* single distance reconstruction, Paganin, save tiff volume:
   ``pynx-holotomo-id16b --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff \
--algorithm AP**5,delta_beta=0,AP**20,Paganin --nz 1 --padding 200 --slurm``

* long 4-distance optimisation using progressive smoothing (improves artefact removal, \
without actually smoothing the object), distortion correction:
   ``pynx-holotomo-id16b --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff \
--algorithm AP**5,delta_beta=0,obj_smooth=0,AP**20,obj_smooth=1,AP**20,obj_smooth=4,AP**20,obj_smooth=8,Paganin \
--nz 4 --padding 200 --distorted --slurm``
"""

    epilog_id16a = """Example usage:

* 4-distance reconstruction, CTF, save tiff volume, random motion:
   ``pynx-holotomo-id16a --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff --motion \
--algorithm AP**5,AP**20,CTF --nz 1 --padding 200 --slurm``

* 4 distances, Paganin, distortion + spikes + motion correction:
  ``pynx-holotomo-id16a --data sample_020nm_1_ --delta_beta 100 --save_fbp_vol tiff \
--algorithm AP**5,delta_beta=0,AP**10,Paganin --nz 4 --distorted --remove_spikes 0.04 --padding 200 \
--motion  --slurm``

* 4-distance reconstruction, CTF, save tiff volume, random motion, \
use phased images for center of rotation:
   ``pynx-holotomo-id16a --data data/alcu_25nm_8000adu_1_ --delta_beta 530 --save_fbp_vol tiff --motion \
--algorithm AP**5,AP**20,CTF --nz 1 --padding 200 --tomo_cor_method phase_registration --slurm``
       """
    if instrument is None:
        instrument = 'id16b' if 'id16b' in sys.argv[0] else 'id16a'

    if instrument == 'id16b':
        parser = argparse.ArgumentParser(prog="pynx-holotomo-id16b",
                                         description="Script for single or multi-distance holo-tomography "
                                                     "reconstruction from data recorded on the ESRF ID16B "
                                                     "beamline.",
                                         epilog=epilog_id16b,
                                         formatter_class=argparse.RawTextHelpFormatter)
    elif instrument == 'id16a':
        parser = argparse.ArgumentParser(prog="pynx-holotomo-id16a",
                                         description="Script for single or multi-distance holo-tomography "
                                                     "reconstruction from data recorded on the ESRF ID16A "
                                                     "beamline.",
                                         epilog=epilog_id16a,
                                         formatter_class=argparse.RawTextHelpFormatter)
    else:
        raise RuntimeError("What program was launched ? Should be pynx-holotomo-id16a or pynx-holotomo-id16b ?")

    parser.add_argument("--data", action='store', required=True,
                        help="Path to the dataset directory, e.g. path/to/sample_1, "
                             "For nz>1 this should be one of the distances directories. "
                             "This can also use a more explicit format with "
                             "e.g. path/to/sample_%%02d_ in case the usual naming changed."
                             "This can also be a NXTomo file from nxtomomill, e.g. "
                             "path/to/sample_50nm_0000.nx")
    # parser.add_argument("--img", action='store', required=False,
    #                     help="Name of the images, inside each distance directory, e.g. "
    #                          "prefix_%d_%04d.edf - this should include two fields, the first"
    #                          "for the distance index, the second for the image.")
    parser.add_argument("--delta_beta", action='store', type=float, required=True,
                        help="delta/beta value for CTF/Paganin")

    group = parser.add_argument_group("Input parameters")
    group.add_argument("--nz", action='store', type=int, default=None,
                       help="Number of distances (planes) for the analysis")
    group.add_argument("--planes", action='store', nargs='+', type=int, default=None,
                       help="Planes to use for the analysis. This is normally set by --nz, "
                            "e.g. '--nz 4' is equivalent to '--planes 1 2 3 4'. "
                            "This option can be used to exclude specific planes e.g. "
                            "the first distance using '--planes 2 3 4'")
    group.add_argument("--binning", action='store', type=int, default=1,
                       help="Binning parameter.")
    group.add_argument("--first", action='store', type=int, default=0,
                       help="Index of the first projection to analyse. Default=0 (first "
                            "available projection)")
    group.add_argument("--last", action='store', type=int, default=None,
                       help="Index of the last (included) projection to analyse. Defaults "
                            "to the last recorded projection, automatically detected.")
    group.add_argument("--step", action='store', type=int, default=None,
                       help="Step to pick the projections to analyse. Defaults to the "
                            "binning value. Can also be >1 to distribute projections to "
                            "independent jobs for phasing, or to separate odd and even "
                            "frames. For example using '--step 2 --first 0' and "
                            "'--step 2 --first 1' will respectively use even and odd "
                            "frames.")

    group = parser.add_argument_group("Pre-processing")
    group.add_argument("--remove_spikes", action='store', type=float, default=0,
                       help="Replace any pixel which differs from the median-filtered (3x3) image "
                            "by more than this (relative) value. Typical value 0.04. "
                            "Filtering is done after dark subtraction.")
    group.add_argument("--double-flat-field", "--double_flat_field", action='store',
                       nargs='*', type=float, default=None,
                       help="Apply a double flat-field correction, i.e. normalise "
                            "the radios by their average value after normalisation "
                            "to the white/empty beam images. Optional filtering options "
                            "can be given, e.g. '--double-flat-field 0.1 0.05': "
                            "the first number corresponds to the high-pass filter "
                            "which will filter out frequencies below 0.1 in the 2D image"
                            "plane, and the second is a low-pass filter along the "
                            "projections axis (this second filter requires much longer "
                            "3D calculations )."
                            "By default the values are (0, 0), which means that the radios"
                            "are normalised by <radios/ref>_z (where the average is taken "
                            "along the projections index). If the high-pass filter is "
                            "used, then the normalisation is instead 1+F(radios/ref) where "
                            "F is the filter used. Good values can be between 0.001 to "
                            "0.05 for the first (relative) frequency (larger values remove"
                            "most of the correction), and up to 0.1 "
                            "for the second (0.1 would remove ring artefacts "
                            "that span about 1/10th of a full ring.)")
    group.add_argument("--distorted", action='store', default=0, const=100, type=int, nargs='?',
                       help="Correct for the distortion of frames vs reference images. "
                            "Can be used for local tomography. Optionally, one "
                            "argument with the tile size can be given (default is 100 "
                            "if '--distorted' is passed without argument). Valid values "
                            "range from 40 to 400.")
    group.add_argument("--normalise", action='store_true', default=False,
                       help="Use this keyword to normalise the observed projections by "
                            "the reference images, immediately upon loading. The normalised "
                            "radios are scaled to keep the average intensity, so that "
                            "the statistical figure of merits are unaffected. "
                            "Note that currently (EXPERIMENTAL), this triggers the "
                            "pre-shifting of all frames before phasing, and so disables "
                            "probe optimisation.")

    class ActionVSNR(argparse.Action):
        def __call__(self, parser_, namespace, values, option_string=None):
            if len(values) % 4 != 2 or len(values) < 6:
                raise argparse.ArgumentError(self,
                                             "You must give 4*N+2 parameters for --vsnr, e.g. "
                                             "'--vsnr 0.1 0.04 0.005 34 0.1 0.04 0.005 38 10 5'")
            try:
                vv = [float(v) for v in values[:-1]]
                vv.append(int(values[-1]))
            except ValueError:
                raise argparse.ArgumentError(self, "Error interpreting --vsnr values")
            for n in range(len(values) // 4):
                qr0, sigmar, sigmat, angle = vv[4 * n: 4 * (n + 1)]
                if qr0 < 0 or qr0 > 0.2:
                    raise argparse.ArgumentError(self, "--vsnr [qr0 sigma_r sigma_t angle] beta nb: "
                                                       "need 0<=qr0<=0.2")
                if sigmar < 0.005 or sigmar > 0.2:
                    raise argparse.ArgumentError(self, "--vsnr [qr0 sigma_r sigma_t angle] beta nb: "
                                                       "need 0.005<=sigma_r<=0.2")
                if sigmat < 0.005 or sigmat > 0.05:
                    raise argparse.ArgumentError(self, "--vsnr [qr0 sigma_r sigma_t angle] beta nb: "
                                                       "need 0.005<=sigma_t<=0.05")
                if angle < -90 or angle > 90:
                    raise argparse.ArgumentError(self, "--vsnr [qr0 sigma_r sigma_t angle] beta nb: "
                                                       "need -90<=angle<=90")
            if vv[-2] < 1 or vv[-2] > 20:
                raise argparse.ArgumentError(self, "--vsnr [qr0 sigma_r sigma_t angle] beta nb: "
                                                   "need 1<= beta <=20")
            if vv[-1] < 1 or vv[-1] > 20:
                raise argparse.ArgumentError(self, "--vsnr [qr0 sigma_r sigma_t angle] beta nb: "
                                                   "need 1<= nb <=20")
            setattr(namespace, self.dest, vv)

    h = argparse.SUPPRESS if 'id16b' not in sys.argv[0] else \
        "Variational Stationary Noise Remover. You can give a list " \
        "of 4*N+1 parameters, with first groups of 4 parameters " \
        "[qr0 sigma_r sigma_t angle] defining the filter, and finally " \
        "the beta value (10) for the ADMM as well as " \
        "the number of cycles (<=20) for the iterative refinement. " \
        "qr0 should be between 0 and 0.2, " \
        "sigma_r (radial) must be between 0.005 and 0.2, sigma_t " \
        "(tangential) must be 0.005< <0.05, and -90<=angle<=90" \
        "Typical values for ID16B with two filters at +34°, -38° " \
        "and 2 cycles: '--vsnr 0.1 0.04 0.005 34 0.1 0.04 0.005 -38 10 2'" \
        " [EXPERIMENTAL]"
    group.add_argument("--vsnr", action=ActionVSNR, default=None, nargs='+',
                       help=h)

    h = argparse.SUPPRESS if 'id16b' not in sys.argv[0] else \
        "detector area for the VSNR correction [ymin ymax xmin xmax], " \
        "e.g. --vsnr_roi 240 1500 540 1800. Values should correspond " \
        "to the unbinned data. The VSNR correction will be smoothed " \
        "using an apodisation of 50 pixels."
    group.add_argument("--vsnr_roi", action='store',
                       default=[240, 1500, 540, 1800], nargs=4, help=h)

    group = parser.add_argument_group("Experimental parameters")
    group.add_argument("--nrj", action='store', type=float,
                       help="Energy (keV). By default will be read from the info or edf files.")
    group.add_argument("--distance", action='store', nargs='+', type=float,
                       help="Sample-detector distance(s) (m), after magnification. "
                            "Multiple values should be given if nz>1 using "
                            "'--distance 1e-3 1.2e-3 1.5e-3 1.9e-3' . By default "
                            "will be automatically calculated from .info and/or motor positions.")
    group.add_argument("--pixel_size", action='store', nargs='+', type=float,
                       help="Pixel size(s) (m), for each distance, after magnification. "
                            "Multiple values should be given if nz>1 using "
                            "'--pixel_size 1e-8 1.2e-8 1.6e-8 2.2e-8' . "
                            "By default this is computed from .info and/or edf files. "
                            "Will be multiplied by binning.")
    group.add_argument("--sx0", action='store', type=float,
                       help="focus offset parameter in mm: this is used to compute the magnified pixel sizes "
                            "and sample-detector distances will be calculated from cx and sx. "
                            "By default this value is automatically deduced from the data files, "
                            "using sx, cx, and the magnified and non-magnified pixel sizes.")
    group.add_argument("--pixel_size_detector", action='store', type=float, default=None,
                       help="Detectr pixel size (m), before magnification. "
                            "This can be used only in combination with sx0, to compute"
                            "the magnified distances. If not given, this is read from the "
                            "'optic_used' field in the .info or edf files.")
    group.add_argument("--cx", action='store', type=float,
                       help="detector motor position (mm), used if sx0 is also given. "
                            "By default the value is read from edf files, so this should only "
                            "be used if there was something wrong in the recorded positions.")
    group.add_argument("--sx", action='store', nargs='+', type=float,
                       help="sample motor position (mm), used only if sx0 is also given. "
                            "Multiple values should be given (--sx 3 4 5 6) if nz>1. "
                            "By default the value is read from edf files, so this should only "
                            "be used if there was something wrong in the recorded positions.")
    group.add_argument("--halftomo", action='store', default=None,
                       type=str, choices=['left', 'right'],
                       help="Use this option for a half-tomo dataset. This will trigger "
                            "the search for the center of rotation on the correct side, "
                            "and expand the field of view during the tomography "
                            "reconstruction with Nabu.")

    group = parser.add_argument_group("Alignment parameters")
    group.add_argument("--rhapp", action='store', type=str, default=None,
                       help="Filename for the rhapp alignment matrix written by holoCT. "
                            "If set, this automatically sets the alignment method to using "
                            "the rhapp matrix, unless --align is also used (the rhapp matrix "
                            "can then still be used for comparison).")
    group.add_argument("--align", action='store', type=str, default=None,
                       choices=['fft', 'fft_diff', 'rhapp'] +
                               [f'fft{s}_spline{n}' for s in ['', '_diff'] for n in range(3, 7)] +
                               [f'fft{s}_poly{n}' for s in ['', '_diff'] for n in range(2, 7)] +
                               [f'fft{s}_median{n}' for s in ['', '_diff'] for n in [3, 5, 7]],
                       help="Alignment method, 'fft' (Fourier cross-correlation), "
                            "'fft_splineN' (same as fft plus a 3rd-order spline fit with N=3-6 knots), "
                            "'fft_polyN' (same as fft plus a Nth order polynomial fit (N=2-6), "
                            "'fft_medianN' (same as fft with  only a median filter of width N=3,5 or 7)), "
                            "'fft_diff', 'fft_diff_splineN', 'fft_diff_polyN', 'fft_diff_medianN': same "
                            "as the previous methods, but align using the difference between the"
                            "projection and the rolling 3-projection average-"
                            "this approach is very robust but incompatible with random motion. "
                            "The default is 'fft_diff_spline3' or 'fft_spline3' if --motion is used. "
                            "Alternatively, 'rhapp' will be used automatically if --rhapp is given.")
    group.add_argument("--align_interval", action='store', type=int, default=1,
                       help="if > 1, align only every N projection, and interpolate "
                            "for the other projections. N>1 can be useful if loading & "
                            "alignment is slower than algorithms when using multiple GPUs.")
    group.add_argument("--align_max_shift", action='store', type=int, default=200,
                       help="Maximum shift (+/-) expected in pixels between projections. This is used"
                            "to limit the shift range and to zero-pad when aligning. The "
                            "value will be divided by the binning value, if any.")

    def valid_align_fourier_low_cutoff(arg):
        try:
            v = float(arg)
        except ValueError:
            raise argparse.ArgumentTypeError("Must be a floating point number between 0 and 0.05")
        if v < 0 or v > 0.05:
            raise ValueError(f"Argument must be with [0; 0.05]")
        return v

    group.add_argument("--align_fourier_low_cutoff", action='store',
                       type=valid_align_fourier_low_cutoff, default=None,
                       help="Cutoff (default: None), a value (0< <=0.05) to cutoff "
                            "low frequencies during Fourier registration.")

    def valid_align_fourier_high_cutoff(arg):
        try:
            v = float(arg)
        except ValueError:
            raise argparse.ArgumentTypeError("Must be a floating point number between 0.25 and 0.5")
        if v < 0.25 or v > 0.5:
            raise ValueError(f"Argument must be with [0.25; 0.5]")
        return v

    group.add_argument("--align_fourier_high_cutoff", action='store',
                       type=valid_align_fourier_high_cutoff, default=None,
                       help="Cutoff (default: None), a value (0.25<=  < 0.5) to cutoff "
                            "high frequencies during Fourier registration. Can be useful for "
                            "noisy datasets.")
    group.add_argument("--reference_plane", action='store', type=int, default=None,
                       choices=[1, 2, 3, 4],
                       help="Reference plane for the alignment, for nz>1. "
                            "Default is the first plane. Numbering begins at 1, same as "
                            "for the --planes argument")
    group.add_argument("--motion", action='store_true', default=False,
                       help="Take into account the random motion 'correct.txt' data during alignment")

    group = parser.add_argument_group("Reconstruction parameters")

    class ActionAlgorithm(argparse.Action):
        """Argparse Action to validate the algorithm string"""

        def __call__(self, parser_, namespace, value, option_string=None):
            if value in algo_aliases:
                pass
            else:
                for algo in value.split(",")[::-1]:
                    if "=" in algo:
                        try:
                            k, v = algo.split("=")
                            if k in ['beta', 'delta_beta', 'dm_alpha', 'obj',
                                     'obj_inertia', 'obj_max', 'obj_min', 'obj_smooth',
                                     'probe', 'probe_inertia', 'psf', 'verbose',
                                     'constrain_probe_direct_beam']:
                                junk = float(v)
                            else:
                                raise RuntimeError("Incorrect algorthm step")
                        except Exception:
                            raise argparse.ArgumentError(self, f"Did not understand algorithm step: '{algo}'")
                    elif algo.lower() in ['paganin', 'ctf']:
                        pass
                    else:
                        try:
                            dm, ap, apn, raar, drap = 1, 1, 1, 1, 1
                            junk = eval(algo.lower())
                        except Exception:
                            raise argparse.ArgumentError(self, f"Did not understand algorithm step: '{algo}'")
            setattr(namespace, self.dest, value)

    group.add_argument("--algorithm", action=ActionAlgorithm, type=str,
                       default="AP**5,delta_beta=0,AP**10,CTF",
                       help=f"Algorithm string, executed from right to left.\n"
                            f"It is recommended to use the following aliases for standard "
                            f"algorithm chains:\n\n"
                            f"* ``paganin-norm``: equivalent to ``{algo_aliases['paganin-norm']}`` "
                            f"(when using --normalise, which makes long iterations ineffective)\n"
                            f"* ``ctf-norm``: same as paganin-norm, but with CTF\n"
                            f"* ``paganin-short``: eq. to "
                            f"``{algo_aliases['paganin-short']}``\n"
                            f"* ``ctf-short``: same as paganin-short, but with CTF\n"
                            f"* ``paganin-long``: eq. to "
                            f"``{algo_aliases['paganin-long']}``\n"
                            f"* ``ctf-long``: same as paganin-long, but with CTF\n"
                            f"* ``paganin-long-dm``: eq. to "
                            f"``{algo_aliases['paganin-long-dm']}``\n"
                            f"* ``ctf-long-dm``: same as paganin-long-dm, but with CTF\n\n"
                            f"Note that in the above, the 'obj_smooth' is a regularisation term which "
                            f"forces the probe to incorporate as much as possible of the high frequencies, "
                            f"thus reducing ring artefacts in the object. It does not smooth the final "
                            f"object as long as 5-10 cycles without smoothing are done at the end.\n\n"
                            f"Examples of explicit algorithms:\n\n"
                            f"* ``AP**5,delta_beta=0,AP**10,Paganin``\n"
                            f"* ``AP**10,obj_smooth=0.5,AP**20,obj_smooth=8,CTF``\n"
                            f"* ``AP**10,obj_smooth=0.5,Paganin``\n"
                            f"* ``AP**10,obj_smooth=0.5,DRAP**10,obj_smooth=1,CTF``\n\n"
                            f"Finally, note that large number of cycles (e.g. >70 AP for nz=4 and "
                            f"~3200 projections and padding=200) will take more than one hour,"
                            f"and will require using either the 'gpu-long' or 'low-gpu' partition,"
                            f"with the --slurm_partition option.")
    group.add_argument("--padding", action='store', type=int, default=0,
                       help="Number of pixels to pad on every side of the array. 20%% of the "
                            "size is normally enough. The value will be rounded to the closest "
                            "value which allows a radix FFT. 2 can be given as a special value "
                            "and will correspond to a doubling of the array size, i.e. a padding "
                            "of ny/2 and nx/2 along each dimension (not recommended, this wastes "
                            "memory and computing time). NOTE: this is only the padding used "
                            "during phasing - when computing the filtered back-projection, "
                            "the padding used always doubles the array size.")
    group.add_argument("--pad_method", action='store', type=str, default='reflect_linear',
                       choices=['edge', 'mean', 'median', 'reflect', 'reflect_linear', 'reflect_erf',
                                'reflect_sine', 'symmetric', 'wrap'],
                       help="Padding method. See methods description in numpy.pad() online doc, with "
                            "'reflec' additional methods, using a linear combination of wrapped "
                            "reflected arrays across the border, to guarantee a continuity across"
                            "the edges of the padded array.")
    group.add_argument("--stack_size", action='store', type=int, default=None,
                       help="Stack size-how many projections are stored simultaneously in memory. "
                            "This is normally automatically adapted to the available GPU memory.")
    group.add_argument("--obj_smooth", action='store', type=float, default=0,
                       help="Object smoothing parameter in pixels. Can be overridden in the "
                            "algorithm string. Requires obj_inertia>0 to work. Note that it does "
                            "not smooth directly the object by a gaussian of the given sigma, but "
                            "rather regularises the updated object to be close to the previous "
                            "gaussian-blurred iteration.")
    group.add_argument("--obj_inertia", action='store', type=float, default=0.1,
                       help="Object inertia parameter. Can be overridden in the algorithm string. "
                            "Required > 0 to exploit smoothing, typical values 0.1 to 0.5.")
    group.add_argument("--obj_min", action='store', type=float, default=None,
                       help="Object minimum amplitude (e.g. 1)")
    group.add_argument("--obj_max", action='store', type=float, default=None,
                       help="Object maximum amplitude (e.g. 1)")
    group.add_argument("--beta", action='store', type=float, default=0.9,
                       help="Beta coefficient for the RAAR or DRAP algorithm. Not to be mistaken"
                            "with the refraction index beta")
    group.add_argument("--nb_probe", action='store', type=int, default=1, choices=[1, 2, 3, 4],
                       help="Number of coherent probe modes. If >1, the (constant) probe mode coefficients"
                            "will linearly vary as a function of the projection index. "
                            "EXPERIMENTAL")
    group.add_argument("--probe", action='store', type=str, default=None,
                       nargs='+',
                       help=argparse.SUPPRESS)
    group.add_argument('--constrain_probe_direct_beam', action='store_true',
                       help="Use this so that each time the probe is updated, it "
                            "is constrained to match the direct beam frame "
                            "[EXPERIMENTAL]")
    group.add_argument('--psf', action='store', type=float, default=None,
                       help=argparse.SUPPRESS)
    # "Probe to use for the optimisation, coming from a near-field"
    #      "ptycho analysis. If used, this will disable the probe optimisation. "
    #      "Since one probe is required for each distance, this should either "
    #      "be a pattern e.g. 'result%%02.cxi', or a list of nz files."
    #      "EXPERIMENTAL.")

    # group.add_argument("--liveplot", action='store', type=float, default=1,
    #                    help="Liveplot of phasing step. Currently unused")
    group.add_argument('--one_probe', '--one-probe', action='store_true',
                       help="Use this option to use the same probe for all distances "
                            "(taking into account propagation) [EXPERIMENTAL]")

    group = parser.add_argument_group("Tomography reconstruction parameters")
    # group.add_argument("--tomo_angle_step", action='store', type=float, default=None,
    #                    help="Angular step (degrees), not taking into account binning. "
    #                         "By default, is determined automatically from data files.")
    group.add_argument("--tomo_rot_center", action='store', type=float, default=None,
                       help="Rotation centre (pixels), not taking into account binning. "
                            "By default, this is determined automatically.")
    group.add_argument("--tomo_rot_center_near", action='store', type=float, default=None,
                       help="Estimated position for the rotation centre (pixels), "
                            "not taking into account binning. This value is passed to "
                            "Nabu's CoR methods for CoR estimation (near_pos).")
    group.add_argument("--tomo_cor_method", action='store', type=str, default='radios_sliding_window',
                       choices=["radios_sliding_window", "radios_growing_window",
                                "radios_accurate", "global",
                                "phase_registration", "phase_registration_diff",
                                "phase_registration_high", "phase_registration_diff_high"],
                       help="Method used for the determination of the centre of rotation, either "
                            "sliding/growing window/accurate on radios (from Nabu), "
                            "'sino' using Nabu's SinoCOR, or 'phase_registration' using "
                            "cross-correlation of phased projections at +180°")
    group.add_argument("--sino_filter", action='store', type=str, default='ram-lak',
                       choices=['ram-lak', 'shepp-logan', 'cosine', 'hamming', 'hann',
                                'tukey', 'lanczos', 'none'],
                       help="Name for the sinogram 1D filter before back-projection. "
                            "The filtering will be done using the padded array if "
                            "padding is used.")

    group = parser.add_argument_group("Output parameters")
    group.add_argument("--prefix_output", action='store', default=None,
                       help="Prefix for the output (can include a full directory path), "
                            "Defaults to data_prefix_resultNN")
    group.add_argument('--processed_data', '--processed-data',
                       action='store_true',
                       help='Use this to automatically save the results '
                            'in ``path/to/PROCESSED_DATA/other/path`` if the'
                            'original data is in ``path/to/RAW_DATA/other/path/``. '
                            'This can be combined with ``--prefix_output`` to specify the '
                            'files prefix (and even add a subfolder).')
    group.add_argument("--save_phase_chunks", action='store_true', default=False,
                       help="Save the phased projections to an hdf5 file. "
                            "This is deprecated, use --nxtomo instead")
    group.add_argument("--save_edf", action='store_true', default=False,
                       help="Save the phased projections individual edf files.")
    group.add_argument("--save_fbp_vol", action='store', default=False, type=str, choices=['hdf5', 'tiff', 'cuts'],
                       help="Save the reconstructed volume, either using the hdf5 format (float16), "
                            "uint16 tiff (single file if volume<4GB, individual slices otherwise),"
                            "or only with a png with the 3 cuts (mostly for testing)")
    group.add_argument("--nxtomo", action='store_true', default=False,
                       help="Save the phased projections using the NXTomo format.")
    group.add_argument("--save_3sino", action='store_true', default=False,
                       help="Save the 3 sinograms at 25, 50 and 75%% of height for testing.")
    # Hidden option to output volume from holoct reconstructed projections. The supplied string must
    # be the file pattern for holoct result files. The program will exit immediately after
    # reconstructing the 'holoct' volume.
    group.add_argument("--holoct", action='store', default=None, type=str,
                       help=argparse.SUPPRESS)

    group = parser.add_argument_group("Job parameters")
    group.add_argument("--ngpu", action='store', default=None, type=int,
                       help="number of GPUs the job should be distributed to (default=2"
                            "for a slurm job, otherwise defaults to 1)")
    group.add_argument("--slurm", action='store_true', default=False,
                       help="If given, the job will be submitted as an ESRF slurm job (p9gpu)")

    default_partition = ['gpu', 'p9gpu']
    try:
        # Development hook
        if 'favre/.conda/envs' in os.environ['CONDA_PREFIX'] and 'scisoft' in platform.node():
            if 'p9' in os.environ['CONDA_DEFAULT_ENV']:
                default_partition = ['p9gpu']
            elif 'x86' in os.environ['CONDA_DEFAULT_ENV']:
                default_partition = ['gpu']
    except OSError:
        pass
    except KeyError:
        pass
    group.add_argument("--slurm_partition", action='store', nargs='*', default=default_partition,
                       choices=['gpu', 'p9gpu', 'gpu-long', 'p9gpu-long', 'low-gpu', 'low-p9gpu'],
                       help=f"Specify the slurm partition(s) to use (default is "
                            f"{','.join(default_partition)}, 1 hour max)")

    group.add_argument("--slurm_nodelist", action='store', default=None,
                       help="Specify a nodelist for the slurm job submission (mostly for debugging or "
                            "very specific needs e.g. memory). The string will be passed directly to "
                            "sbatch --nodelist=...")
    group.add_argument("--slurm_time", action='store', default=None,
                       help="Specify a time limit for the slurm job (mostly for debugging). "
                            "The string will be passed directly to sbatch --time=... . "
                            "By default the time limit is automatically chosen.")
    # group.add_argument("--serial", action='store_true', default=False,
    #                    help="If given, the analysis will be run with multiprocessing (for debugging)")
    group.add_argument("--max_mem_chunk", action='store', default=64, type=int,
                       help="Maximum memory (GB) available per chunk (ngpu chunks are processed in //)."
                            "The actual memory used should be about 1.5x this value. This should be kept"
                            "reasonably low because the process uses a lot of shared memory, and the"
                            "analysis can slow down significantly if too much is used.")
    group.add_argument("--profile", action='store_true', default=False,
                       help="Enable GPU profiling (for development only)")
    group.add_argument("--dry_run", "--dry-run", action='store_true',
                       help="Use this to print what would be executed: if --slurm is "
                            "used, the slurm job submission is printed. Otherwise, the "
                            "execution will stop after printing the parameters.")
    return parser


def make_parser_sphinx_id16b():
    """Returns the argparse for sphinx documentation"""
    return make_parser('id16b')


def main():
    params = make_parser().parse_args()
    if params.data[-1] == "/":
        params.data = params.data[:-1]
    # Use planes or nz
    if params.planes is not None:
        if params.nz is not None:
            if len(params.planes) != params.nz:
                raise RuntimeError("Parameters are inconsistent: --nz=%d  and --planes=%s"
                                   % (params.nz, str(params.planes)))
        params.nz = len(params.planes)
    else:
        if params.nz is None:
            params.nz = 1  # Default value if neither is set
        if params.nz == 1 and params.planes is None:
            if re.search('%[0-9]*?d', params.data) is None:  # Matches %d, %2d, %02d
                params.planes = [int(params.data[-2])]  # "...dataset_2_d/..."
        if params.planes is None:
            params.planes = list(range(1, params.nz + 1))
    # Make sure --cx and --sx are only supplied along --sx0
    if params.sx0 is None and params.sx is not None:
        raise RuntimeError(f"--sx {params.sx} was given but --sx0 was not used. Did you "
                           f"mistake the two ? Note that --sx should only be used when the "
                           f"sx value is incorrect in the edf files (offset or "
                           f"encoder error). Usually, only --sx0 is needed")
    if params.sx0 is None and params.cx is not None:
        raise RuntimeError(f"--cx {params.cx} was given but --sx0 was not used. Did you "
                           f"mistake the two ? Note that --cx should only be used when the "
                           f"cx value is incorrect in the edf files (offset or "
                           f"encoder error). Usually, only --sx0 is needed")
    if params.sx is not None:
        if len(params.sx) != params.nz:
            raise RuntimeError(f"--sx has been given, but with only {len(params.sx)} "
                               f"values instead of nz={params.nz}")

    # Change data to use %d for the plane
    # TODO: change data to use %d for the plane in the parser
    if params.data.endswith('.nx'):
        if re.search('%[0-9]*?d', params.data) is None:  # Matches %d, %2d, %02d
            params.data = f"{params.data[:-4]}%d.nx"
        # Check that data exists
        for iz in params.planes:
            if not os.path.exists(params.data % (iz - 1)):
                raise RuntimeError(f"NX data file does not exists for iz(plane)={iz}: {params.data % (iz - 1)}")
    else:
        if re.search('%[0-9]*?d', params.data) is None:  # Matches %d, %2d, %02d
            params.data = params.data[:-2] + '%d_'

        # Check that data exists
        for iz in params.planes:
            if not os.path.isdir(params.data % iz):
                raise RuntimeError(f"Data folder does not exists for iz(plane)={iz}: {params.data % iz}")

    # Step defaults to binning
    if params.step is None:
        params.step = params.binning

    # Assign or check reference plane
    if params.reference_plane is None:
        params.reference_plane = params.planes[0]
    elif params.reference_plane not in params.planes:
        raise RuntimeError("Reference plane (%d) is not in the given list of planes: %s"
                           % (params.reference_plane, str(params.planes)))

    # Change prefix_output if necessary
    data_dir, prefix = os.path.split(params.data)
    prefix = prefix.split('%')[0]  # Remove the trailing '%d_'

    path_output = ''
    if params.processed_data:
        path_output = replace_dir(data_dir)
        if path_output != data_dir:
            if path_output[-1] == '_':
                path_output = path_output[:-1]
        else:
            warnings.warn(f"Using --processed_data but RAW_DATA was not found in "
                          f"the real data path - option is ignored")
            path_output = ''

    if params.prefix_output is None:
        i = 1
        while True:
            if params.nz > 1:
                params.prefix_output = prefix + f"result{i:02d}"
            else:
                iz = params.planes[0]
                params.prefix_output = prefix + f"{iz}_result{i:02d}"
            try:
                l = os.listdir('.' if path_output in ['', None] else path_output)
            except FileNotFoundError:
                break  # Output directory does not exist yet
            if len([fn for fn in l if fn.startswith(params.prefix_output)]) == 0:
                break
            i += 1
    params.prefix_output = os.path.join(path_output, params.prefix_output)
    print(f"using prefix_output:  {params.prefix_output}\n")

    # Check filtering options
    if params.double_flat_field is not None:
        if len(params.double_flat_field) > 2:
            raise RuntimeError("Only up to two parameters can be supplied for --double-flat-field")
        if len(params.double_flat_field) == 0:
            params.double_flat_field = [0, 0]
        elif len(params.double_flat_field) == 1:
            params.double_flat_field.append(0)
        for r in params.double_flat_field:
            if r < 0 or r > 0.25:
                raise RuntimeError("Relative frequencies for --double-flat-field must be between 0 and 0.25")
    if params.sino_filter == "none":
        params.sino_filter = None

    # Check distorted value
    if params.distorted:
        if params.distorted < 40 or params.distorted > 400:
            raise RuntimeError(f"Error: --distorted value ({params.distorted}) must be within [40,400]")

    # Change alignment method if --rhapp is used
    if params.rhapp is not None:
        params.align = 'rhapp'
    elif params.align is None:
        if params.motion:
            params.align = 'fft_spline3'
        else:
            params.align = 'fft_diff_spline3'

    if params.slurm:
        # Submit job to the ESRF slurm cluster

        if params.ngpu is None:
            params.ngpu = 2

        # Create logfile
        path = os.path.split(params.prefix_output)[0]
        if len(path):
            os.makedirs(path, exist_ok=True)
        s = "%s.slurm" % params.prefix_output
        slurm_partition = ','.join(params.slurm_partition)
        with StringIO() if params.dry_run else open(s, 'w') as slurmfile:
            slurmfile.write("#!/bin/bash -l\n")
            slurmfile.write("#SBATCH --partition=%s\n" % slurm_partition)
            if params.ngpu > 1:
                slurmfile.write("#SBATCH --gres=gpu:%d\n" % params.ngpu)
                slurmfile.write("#SBATCH --mem=0\n")
                # slurmfile.write("#SBATCH --cpus-per-task=128\n")
                slurmfile.write("#SBATCH --exclusive\n")
            else:
                slurmfile.write("#SBATCH --gres=gpu:1\n")
                slurmfile.write("#SBATCH --mem=256g\n")
                slurmfile.write("#SBATCH --cpus-per-task=64\n")
            slurmfile.write("#SBATCH --ntasks=1\n")
            if params.slurm_time is not None:
                slurmfile.write("#SBATCH --time=%s\n" % params.slurm_time)
            else:
                if all([('low' in s or 'long' in s) for s in params.slurm_partition]):
                    slurmfile.write("#SBATCH --time=600\n")
                else:
                    slurmfile.write("#SBATCH --time=60\n")
            if params.slurm_nodelist is not None:
                slurmfile.write("#SBATCH --nodelist=%s\n" % params.slurm_nodelist)
            slurmfile.write("#SBATCH --output=%s.out\n" % params.prefix_output)
            slurmfile.write("#SBATCH --job-name=pynx-holotomo\n")
            # Removing environment variables is important e.g. if you submit a job
            # from within another (interactive) slurm job, as there may be conflicts
            # between the old and new job variables (for example MPI/PMI parameters),
            # especially if the platforms differ (x86/p9). We still export the PATH
            # so essential packages can be found (gcc..)
            slurmfile.write("#SBATCH --export=PATH\n")

            slurmfile.write("module purge\n")
            slurmfile.write("env | grep CUDA\n")
            slurmfile.write("scontrol --details show jobs $SLURM_JOBID |grep RES\n")
            slurmfile.write(("cat /proc/sys/kernel/shmmax\n"))
            slurmfile.write(("cat /proc/sys/kernel/shmall\n"))
            slurmfile.write(("ipcs -m --human\n"))
            slurmfile.write(("ipcs -pm --human\n"))
            try:
                # Can fail e.g. if launched from an interactive slurm job
                login_name = os.getlogin()
            except OSError:
                if 'USER' in os.environ:
                    login_name = os.environ['USER']
                else:
                    login_name = 'unknown'
            try:
                if 'favre/.conda/envs' in os.environ['CONDA_PREFIX'] and 'scisoft' in platform.node():
                    # Development hook - should be removed from production code
                    slurmfile.write("module load mamba\n")
                    slurmfile.write(f"conda activate {os.environ['CONDA_DEFAULT_ENV']}\n")
                else:
                    slurmfile.write("module load pynx/devel\n")
            except OSError:
                slurmfile.write("module load pynx/devel\n")
            except KeyError:
                slurmfile.write("module load pynx/devel\n")
            if 'id16b' in sys.argv[0]:
                c = "pynx-holotomo-id16b"
            else:
                c = "pynx-holotomo-id16a"
            for i in range(1, len(sys.argv)):
                a = sys.argv[i]
                # Don't pass some command-lines arguments already handled
                if '--slurm_' not in sys.argv[i - 1] and '--slurm' not in a and \
                        a not in ['--dry-run', '--dry_run', '--processed_data', '--processed-data']:
                    c += " %s" % a
            if '--ngpu' not in c:
                c += " --ngpu %s" % params.ngpu
            if "prefix_output" not in c:
                c += " --prefix_output %s" % params.prefix_output
            slurmfile.write(c)
            if params.dry_run:
                slurmfile.flush()
                slurmfile.seek(0)
                print(f"DRY-RUN: would have submitted the following slurm job:"
                      f"\n\n{slurmfile.read()}\n")
        if not params.dry_run:
            res = os.system("sbatch %s" % s)
            if res == 0:
                print("Submitted slurm job with: %s\n" % s)
            else:
                raise RuntimeError(f"An error occurred while submitting the batch job ({s}): "
                                   f"is sbatch or the requested slurm partition "
                                   f"[{slurm_partition}] available ?")
    else:
        if params.ngpu is None:
            params.ngpu = 1
        if params.algorithm.lower() in algo_aliases.keys():
            params.algorithm_real = algo_aliases[params.algorithm.lower()]
        else:
            params.algorithm_real = None
        master(params)


if __name__ == '__main__':
    if sys.version_info < (3, 8):
        raise RuntimeError("The holotomo script requires python>=3.8")
    # set_start_method('spawn')
    main()
