#! /opt/local/bin/python
# -*- coding: utf-8 -*-

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

import os
import platform
import sys
import gc
import copy
import time
import timeit
import warnings
import traceback
from PIL import Image

from ...utils import h5py
import fabio
import numpy as np
from scipy.fftpack import fftn, ifftn, fftshift
from scipy.ndimage.measurements import center_of_mass
from scipy.io import loadmat
from .. import *
from ...processing_unit import Backend
from ..cdi import _as_dim
from ...utils.math import smaller_primes
from ...utils.array import rebin
from ...mpi import MPI
from ..plot import show_cdi
from .parser import make_parser

"""Default parameters.
These may be modified in specific runners
Other default values are defined in make_parser().
"""
default_params = {
    'algorithm': None,
    'auto_center_resize': False,
    'beta': 0.9,
    'confidence_interval_factor_mask_min': 0.5,
    'confidence_interval_factor_mask_max': 1.2,
    'crop': None,
    'crop_output': 4,
    'data': None,
    'data2cxi': False,
    'detwin': False,
    'detector_distance': None,
    'fig_num': 1,
    'flatfield': None,
    'free_pixel_mask': None,
    'gps_inertia': 0.05,
    'gps_t': 1.0,
    'gps_s': 0.9,
    'gps_sigma_f': 0,
    'gps_sigma_o': 0,
    'gpu': None,
    'instrument': None,
    'iobs_saturation': None,
    'live_plot': False,
    'mask': None,
    'mask_interp': None,
    'max_size': None,
    'mpi': 'splitscan',
    'nb_run': 1,
    'nb_run_keep': None,
    'nb_run_keep_max_obj2_out': 0.10,
    'note': None,
    'object': None,
    'output_format': 'cxi',
    'pixel_size_detector': None,
    'plot_axis': 0,
    'psf': False,
    'psf_filter': None,
    'rebin': None,
    'roi': None,
    'roi_final': None,
    'roi_user': None,
    'sample_name': None,
    'save': 'final',
    'save_plot': False,
    'support': 'auto',
    'support_fraction_min': 1e-4,
    'support_fraction_max': 0.7,
    'support_threshold_auto_tune_factor': 1.1,
    'support_autocorrelation_threshold': (0.09, 0.11),
    'support_formula': None,
    'support_only_shrink': False,
    'support_smooth_width_begin': 2,
    'support_smooth_width_end': 0.5,
    'support_smooth_width_relax_n': 500,
    'support_size': None,
    'support_threshold_method': 'rms',
    'support_post_expand': None,
    'support_threshold': 0.25,
    'support_update_border_n': 0,
    'support_update_period': 50,
    'positivity': False,
    'verbose': 50,
    'wavelength': None,
    'zero_mask': False,
}


class CDIRunnerException(Exception):
    pass


class CDIRunnerScan:
    """
    Abstract class to handle CDI data. Must be derived to be used.
    """

    def __init__(self, params, scan, timings=None):
        self.params = params
        self.scan = scan
        self.mask = None
        self.rebinf = 1
        self.cdi = None
        self.iobs = None
        self.support = None
        self.processing_unit = None
        self._algo_s = ""
        # Keep track of saved results, for later sorting & merging
        # Entries are dictionaries with file name and properties including log-likelihood,...
        self.saved_results = []
        if timings is not None:
            self.timings = timings
        else:
            self.timings = {}

    def load_data(self):
        """
        Loads data. This function only loads data (2D or 3D) from generic files, and should be derived for
        beamline-specific imports (from spec+individual images, etc..)

        Returns: nothing. Loads data in self.iobs, and initializes mask and initial support.

        """
        t0 = timeit.default_timer()
        filename = self.params['data']
        print('Loading data: ', filename)

        # Extract hdf5path if supplied as 'filename:path/to/support'
        if ':' in os.path.splitdrive(filename)[1] and ('h5' in filename or 'cxi' in filename):
            # Looks like an hdf5 file with the hdf5 path
            drive, filename = os.path.splitdrive(filename)
            filename, h5paths = filename.split(':')
            filename = drive + filename
            h5paths = [h5paths]
        else:
            h5paths = ['entry_1/instrument_1/detector_1/data', 'entry_1/image_1/data']

        ext = os.path.splitext(filename)[-1]
        if ext == '.npy':
            self.iobs = np.load(filename).astype(np.float32)
        elif ext == '.edf':
            self.iobs = fabio.open(filename).data.astype(np.float32)
        elif ext == '.tif' or ext == '.tiff':
            frames = Image.open(filename)
            nz, ny, nx = frames.n_frames, frames.height, frames.width
            self.iobs = np.empty((nz, ny, nx), dtype=np.float32)
            for i in range(nz):
                frames.seek(i)
                self.iobs[i] = np.array(frames)
        elif ext == '.npz':
            d = np.load(filename)
            if 'data' in d.keys():
                self.iobs = d['data'].astype(np.float32)
            elif 'iobs' in d.keys():
                self.iobs = d['iobs'].astype(np.float32)
            else:
                # Assume only the data is in the datafile
                for k, v in d.items():
                    self.iobs = v.astype(np.float32)
                    break
        elif ext == '.h5':
            with h5py.File(filename, 'r') as h:
                import_nok = True
                for hp in h5paths:
                    if hp in h:
                        self.iobs = h[hp][()].astype(np.float32)
                        import_nok = False
                        break
                if import_nok:
                    raise CDIRunnerException("Could not find data=%s" % self.params['data'])
        elif ext == '.cxi':
            cxi = h5py.File(filename, 'r')
            if '/entry_1/instrument_1/source_1/energy' in cxi:
                nrj = cxi['/entry_1/instrument_1/source_1/energy'][()] / 1.60218e-16
                self.params['wavelength'] = 12.3984 / nrj * 1e-10
                print("  CXI input: Energy = %8.2fkeV" % nrj)
            if '/entry_1/instrument_1/detector_1/distance' in cxi:
                self.params['detector_distance'] = cxi['/entry_1/instrument_1/detector_1/distance'][()]
                print("  CXI input: detector distance = %8.2fm" % self.params['detector_distance'])
            if '/entry_1/instrument_1/detector_1/x_pixel_size' in cxi:
                self.params['pixel_size_detector'] = cxi['/entry_1/instrument_1/detector_1/x_pixel_size'][()]
                print("  CXI input: detector pixel size = %8.2fum" % (self.params['pixel_size_detector'] * 1e6))
            print("  CXI input: loading iobs")
            if 'entry_1/instrument_1/detector_1/data' in cxi:
                self.iobs = cxi['entry_1/instrument_1/detector_1/data'][()].astype(np.float32)
            elif 'entry_1/image_1/data' in cxi:
                self.iobs = cxi['entry_1/image_1/data'][()].astype(np.float32)
                if 'entry_1/image_1/is_fft_shifted' in cxi:
                    if cxi['entry_1/image_1/is_fft_shifted'][()] > 0:
                        self.iobs = fftshift(self.iobs)
            else:
                self.iobs = cxi['entry_1/data_1/data'][()].astype(np.float32)
            if 'entry_1/instrument_1/detector_1/mask' in cxi:
                self.mask = cxi['entry_1/instrument_1/detector_1/mask'][()].astype(np.int8)
                nb = self.mask.sum()
                print("  CXI input: loading mask, with %d pixels masked (%6.3f%%)" % (nb, nb * 100 / self.mask.size))

            if 'entry_1/data_1/process_1/configuration' in cxi:
                # Load specific parameters from CXI file and carry on some parameters
                # TODO: this is a bit of a kludge, find a more maintainable way to decide which parameters to keep..
                for k, v in cxi['entry_1/data_1/process_1/configuration'].items():
                    vv = None
                    if k in self.params:
                        vv = self.params[k]
                    if (k not in self.params.keys() or (k in ['instrument', 'sample_name', 'scan', 'imgname']
                                                        and vv is None)) and v is not None:
                        if type(v) is h5py.Group:
                            self.params[k] = {}
                            for dk, dv in v.items():
                                self.params[k][dk] = dv[...]
                        else:
                            self.params[k] = v[...]

        print("Finished loading iobs data, with size:", self.iobs.size)

        self.timings['load_data'] = timeit.default_timer() - t0

    def rebin_data(self):
        """
        Will rebin the data (iobs and mask) and update pixel_size_detector according to self.params['rebin'].
        this must be called at the end of load_data, after init_mask. If a ROI has been given, cropping
        will occur here.
        :return:
        """
        if self.params['roi_user'] is not None:
            print("Cropping to user ROI: ", self.params['roi_user'])
            if self.iobs.ndim == 3:
                izmin, izmax, iymin, iymax, ixmin, ixmax = self.params['roi_user']
                self.iobs = self.iobs[izmin:izmax, iymin:iymax, ixmin:ixmax]
                if self.mask is not None:
                    self.mask = self.mask[izmin:izmax, iymin:iymax, ixmin:ixmax]
            elif self.iobs.ndim == 2:
                iymin, iymax, ixmin, ixmax = self.params['roi_user']
                self.iobs = self.iobs[iymin:iymax, ixmin:ixmax]
                if self.mask is not None:
                    self.mask = self.mask[iymin:iymax, ixmin:ixmax]

        if self.params['rebin'] is not None:
            print("Rebinning Iobs with rebin=(%s)" % self.params['rebin'])
            rs = self.params['rebin']
            print(self.params)
            if len(rs) == 1:
                self.params['rebin'] = int(rs[0])
                if self.params['rebin'] == 1:
                    print('Ignoring rebin=1')
                    return
            else:
                self.params['rebin'] = []
                for rr in rs:
                    self.params['rebin'].append(int(rr))
                if self.params['rebin'][:self.iobs.ndim] == [1] * self.iobs.ndim:
                    print('Ignoring rebin=1')
                    return
            self.iobs = rebin(self.iobs, self.params['rebin'])
            if self.mask is not None:
                self.mask = rebin(self.mask, self.params['rebin'])

    def init_mask(self):
        """
        Load mask if the corresponding parameter has been set, or just initialize an array of 0.
        This must be called after iobs has been imported, and before cropping, so at the end of load_data()

        Returns:
            Nothing. self.mask is created.
        """
        if self.params['mask'] is None or self.params['mask'] in ['no', 'nan', 'NaN']:
            if self.mask is None:
                # Else the mask has probably been pre-loaded, e.g. from a CXI file
                self.mask = np.zeros_like(self.iobs, dtype=np.int8)

        if self.mask is None:
            self.mask = np.isnan(self.iobs).astype(np.int8)
        else:
            # If there was already a mask loaded from the file, we cumulate with NaN
            self.mask = self.mask.astype(np.int8) + np.isnan(self.iobs).astype(np.int8)

        if self.params['mask'] is not None and self.params['mask'] not in ['no', 'nan', 'NaN']:
            # Allow to add multiple masks
            for mask_string in self.params['mask']:
                if mask_string.lower() == 'zero':
                    # All zero-valued pixels are masked
                    print('Masking all pixels with iobs <= 0')
                    self.mask += (self.iobs <= 0).astype(np.int8)
                elif mask_string.lower() == 'negative':
                    # All pixels <0 are masked
                    print('Masking all pixels with iobs < 0')
                    self.mask += (self.iobs < 0).astype(np.int8)
                elif mask_string.lower() == 'tomo_beamstop':
                    # Mask all half-lines along all directions which are always equal to zero
                    print('Masking all zero-valued half-lines')
                    self.mask += np.zeros_like(self.iobs, dtype=np.int8)
                    nz, ny, nx = self.iobs.shape
                    # Along z
                    iy, ix = np.nonzero(self.iobs[:nz // 2].sum(axis=0) == 0)
                    self.mask[:nz // 2, iy, ix] += 1
                    iy, ix = np.nonzero(self.iobs[nz // 2:].sum(axis=0) == 0)
                    self.mask[nz // 2:, iy, ix] += 1
                    # Along y
                    iz, ix = np.nonzero(self.iobs[:, :ny // 2].sum(axis=1) == 0)
                    self.mask[iz, :ny // 2, ix] += 1
                    iz, ix = np.nonzero(self.iobs[:, ny // 2:].sum(axis=1) == 0)
                    self.mask[iz, ny // 2:, ix] += 1
                    # Along x
                    iz, iy = np.nonzero(self.iobs[:, :, :nx // 2].sum(axis=2) == 0)
                    self.mask[iz, iy, :nx // 2] += 1
                    iz, iy = np.nonzero(self.iobs[:, :, nx // 2:].sum(axis=2) == 0)
                    self.mask[iz, iy, nx // 2:] += 1
                    # Also mask pixels around center which are equal to zero
                    self.mask[3 * nz // 8:5 * nz // 8, 3 * ny // 8:5 * ny // 8, 3 * nx // 8:5 * nx // 8] += \
                        self.iobs[3 * nz // 8:5 * nz // 8, 3 * ny // 8:5 * ny // 8, 3 * nx // 8:5 * nx // 8] <= 0

                elif 'maxipix' in mask_string.lower():
                    if (self.iobs.shape[-2] % 258 != 0) or (self.iobs.shape[-1] % 258 != 0):
                        raise CDIRunnerException(
                            "ERROR: used mask= maxipix, but x and y dimensions are not multiple of 258.")
                    print('Initializing MAXIPIX mask')
                    ny, nx = self.iobs.shape[-2:]
                    dn = 3
                    if mask_string.lower() == 'maxipix1':
                        # Use the big gap pixel (3x surface) intensity in a single pixel,
                        # normalised in corr_flat_field()
                        dn = 2
                    elif mask_string.lower() == 'maxipix6':
                        # The big pixel intensity will be spread over the entire gap in corr_flat_field()
                        dn = 0
                    if dn > 0:
                        if self.iobs.ndim == 2:
                            for i in range(258, ny, 258):
                                self.mask[i - dn:i + dn] += 1
                            for i in range(258, nx, 258):
                                self.mask[:, i - dn:i + dn] += 1
                        else:
                            for i in range(258, ny, 258):
                                self.mask[:, i - dn:i + dn] += 1
                            for i in range(258, nx, 258):
                                self.mask[:, :, i - dn:i + dn] += 1
                else:
                    filename = mask_string
                    h5path = None
                    if ':' in os.path.splitdrive(filename)[1] and ('h5' in filename or 'cxi' in filename):
                        # Looks like an hdf5 file with the hdf5 path
                        drive, filename = os.path.splitdrive(filename)
                        filename, h5path = filename.split(':')
                        filename = drive + filename

                    print('Loading mask from: ', filename)
                    ext = os.path.splitext(filename)[-1].lower()
                    if ext == '.npy':
                        self.mask += np.load(filename).astype(np.int8)
                    elif ext == '.tif' or ext == '.tiff':
                        frames = Image.open(filename)
                        frames.seek(0)
                        self.mask += (np.array(frames) != 0).astype(np.int8)
                    elif ext == '.mat':
                        a = list(loadmat(filename).values())
                        for v in a:
                            if np.size(v) > 1000:
                                # Avoid matlab strings and attributes, and get the array
                                self.mask += np.array(v).astype(np.int8)
                                break
                    elif ext == '.npz':
                        d = np.load(filename)
                        if 'mask' in d.keys():
                            self.mask += d['mask'].astype(np.int8)
                        else:
                            # Assume only the mask is in the datafile
                            for k, v in d.items():
                                self.mask += v.astype(np.int8)
                                break
                    elif ext == '.edf':
                        self.mask += fabio.open(filename).data.astype(np.int8)
                    elif ext in ['.h5', '.cxi']:
                        with h5py.File(filename, 'r') as h:
                            if h5path is None:
                                # Try to find a dataset with 'mask' in the name and an adequate shape
                                def find_mask(name, obj):
                                    if 'mask' in name.lower():
                                        if isinstance(obj, h5py.Dataset):
                                            print(obj.shape, self.iobs.shape)
                                            if obj.shape == self.iobs.shape or \
                                                    (self.iobs.ndim == 3 and obj.shape == self.iobs[0].shape):
                                                return name

                                h5path = h.visititems(find_mask)
                            if False:
                                # Too risky
                                if h5path is None:
                                    # Find the first dataset with an adequate shape
                                    def find_mask2(name, obj):
                                        if isinstance(obj, h5py.Dataset):
                                            if obj.shape == self.iobs.shape or \
                                                    (self.iobs.ndim == 3 and obj.shape == self.iobs[0].shape):
                                                return name

                                    h5path = h.visititems(find_mask2)
                            if h5path is None:
                                raise CDIRunnerException("Supplied mask=%s, but could not find a suitable mask"
                                                         "dataset in the hdf5 file." % filename)
                            print("Loading mask from hdf5 file %s at path: %s" % (filename, h5path))
                            self.mask += h[h5path][()]
                    else:
                        raise CDIRunnerException(
                            "Supplied mask=%s, but format not recognized (recognized .npy,"
                            ".npz, .tif, .edf, .mat, .h5, .cxi)"
                            % filename)

        if self.params['iobs_saturation'] is not None:
            if self.params['iobs_saturation'] > 0:
                m = self.iobs > self.params['iobs_saturation']
                ms = m.sum()
                if ms > 0:
                    print('Masking %d saturated pixels (%6.3f%%)' % (ms, ms / self.iobs.size))
                    self.mask = ((self.mask + m) > 0).astype(np.int8)

        # Keep only 1 and 0
        self.mask = self.mask > 0

        nb = self.mask.sum()
        p = nb * 100 / self.mask.size
        if p > 20:
            print('WARNING: large number of masked pixels !')
        print("Initialized mask, with %d pixels masked (%6.3f%%)" % (nb, p))
        if self.iobs.ndim == 3 and self.mask.ndim == 2:
            print("Got a 2D mask, expanding to 3D")
            self.mask = np.tile(self.mask, (len(self.iobs), 1, 1))
        if self.iobs.shape != self.mask.shape:
            raise CDIRunnerException('Iobs', self.iobs.shape, ' and mask', self.mask.shape, ' must have the same shape')
        # Set to zero masked pixels (need a bool array for this)
        self.iobs[self.mask] = 0

        if self.params['mask'] not in ['zero', 'negative']:
            # Make sure no intensity is < 0, which could come from background subtraction
            tmp = self.iobs < 0
            nb = tmp.sum()
            if nb:
                print("CDI runner: setting %d pixels with iobs < 0 to zero" % nb)
                self.iobs[tmp] = 0

        # Keep mask as int8
        self.mask = self.mask.astype(np.int8)

    def corr_flat_field(self):
        """
        Correct the image by a known flat field. Currently only applies correction for maxipix large pixels
        :return:
        """
        if self.params['flatfield'] is not None:
            filename = self.params['flatfield']
            print('Loading flatfield from: ', filename)
            h5path = None
            if ':' in os.path.splitdrive(filename)[1] and ('h5' in filename or 'cxi' in filename):
                # Looks like an hdf5 file with the hdf5 path
                drive, filename = os.path.splitdrive(filename)
                filename, h5path = filename.split(':')
                filename = drive + filename
            ext = os.path.splitext(filename)[-1]
            flatfield = None
            if ext == '.npy':
                flatfield = np.load(filename).astype(np.float32)
            elif ext == '.edf':
                flatfield = fabio.open(filename).data.astype(np.float32)
            elif ext == '.tif' or ext == '.tiff':
                frames = Image.open(filename)
                nz, ny, nx = frames.n_frames, frames.height, frames.width
                flatfield = np.empty((nz, ny, nx), dtype=np.float32)
                for i in range(nz):
                    frames.seek(i)
                    flatfield[i] = np.array(frames)
            elif ext == '.npz':
                d = np.load(filename)
                if 'flatfield' in d.keys():
                    flatfield = d['flatfield'].astype(np.float32)
                elif 'flat' in d.keys():
                    flatfield = d['flat'].astype(np.float32)
                elif 'data' in d.keys():
                    flatfield = d['data'].astype(np.float32)
                else:
                    # Assume only the flatfield is in the datafile and hope for the best...
                    for k, v in d.items():
                        flatfield = v.astype(np.float32)
                        break
            elif ext in ['.h5', '.cxi']:
                with h5py.File(filename, 'r') as h:
                    if h5path is None:
                        # Try to find a dataset with 'flat' in the name and an adequate shape
                        def find_flat(name, obj):
                            print(name)
                            if 'flat' in name.lower():
                                if isinstance(obj, h5py.Dataset):
                                    print(obj.shape, self.iobs.shape)
                                    if obj.shape == self.iobs.shape or \
                                            (self.iobs.ndim == 3 and obj.shape == self.iobs[0].shape):
                                        return name

                        h5path = h.visititems(find_flat)
                    if h5path is None:
                        raise CDIRunnerException("Supplied flatfield=%s, but could not find a suitable flatfield"
                                                 "dataset in the hdf5 file." % filename)
                    print("Applying flatfield from hdf5 file %s at path: %s" % (filename, h5path))
                    flatfield = h[h5path][()]
            if flatfield is None:
                raise CDIRunnerException(
                    "Supplied flatfield=%s, but format not recognized (recognized .npy, "
                    ".npz, .tif, .edf, .mat, .h5, .cxi)"
                    % filename)
            if flatfield.shape != self.iobs.shape and (self.iobs.ndim != 3 or flatfield.shape != self.iobs[0].shape):
                raise CDIRunnerException(
                    "Supplied flatfield=%s does not have the correct shape:" % filename,
                    flatfield.shape, self.iobs.shape)
            self.iobs *= flatfield.astype(np.float32)
        elif self.params['mask'] is not None:
            if 'maxipix1' in self.params['mask']:
                ny, nx = self.iobs.shape[-2:]
                self.iobs = self.iobs.astype(np.float32)  # Avoid integer types
                if self.iobs.ndim == 2:
                    for i in range(258, ny, 258):
                        self.iobs[i - 3] /= 3
                        self.iobs[i + 2] /= 3
                    for i in range(258, nx, 258):
                        self.iobs[:, i - 3] /= 3
                        self.iobs[:, i + 2] /= 3
                else:
                    for i in range(258, ny, 258):
                        self.iobs[:, i - 3] /= 3
                        self.iobs[:, i + 2] /= 3
                    for i in range(258, nx, 258):
                        self.iobs[:, :, i - 3] /= 3
                        self.iobs[:, :, i + 2] /= 3
            elif 'maxipix6' in self.params['mask']:
                ny, nx = self.iobs.shape[-2:]
                if self.iobs.ndim == 2:
                    for i in range(258, ny, 258):
                        v0, v1 = self.iobs[i - 3] / 3, self.iobs[i + 2] / 3
                        m = self.mask[i - 3] + self.mask[i + 2]
                        for j in range(-3, 3):
                            self.iobs[i + j] = (v0 * (6 - j + 3) + v1 * (j + 3)) / 6
                            self.mask[i + j] += m
                    for i in range(258, nx, 258):
                        v0, v1 = self.iobs[:, i - 3] / 3, self.iobs[:, i + 2] / 3
                        m = self.mask[:, i - 3] + self.mask[:, i + 2]
                        for j in range(-3, 3):
                            self.iobs[:, i + j] = (v0 * (6 - j + 3) + v1 * (j + 3)) / 6
                            self.mask[:, i + j] += m
                else:
                    for i in range(258, ny, 258):
                        v0, v1 = self.iobs[:, i - 3] / 3, self.iobs[:, i + 2] / 3
                        m = self.mask[:, i - 3] + self.mask[:, i + 2]
                        for j in range(-3, 3):
                            self.iobs[:, i + j] = (v0 * (6 - j + 3) + v1 * (j + 3)) / 6
                            self.mask[:, i + j] += m
                    for i in range(258, nx, 258):
                        v0, v1 = self.iobs[:, :, i - 3] / 3, self.iobs[:, :, i + 2] / 3
                        m = self.mask[:, :, i - 3] + self.mask[:, :, i + 2]
                        for j in range(-3, 3):
                            self.iobs[:, :, i + j] = (v0 * (6 - j + 3) + v1 * (j + 3)) / 6
                            self.mask[:, :, i + j] += m
                # Keep only 1 and 0 for the mask
                self.mask = (self.mask > 0).astype(np.int8)

            # np.savez_compressed('test_maxipix_flat_mask.npz', mask=self.mask, iobs2=self.iobs.sum(axis=0))

    def init_support(self):
        """
        Prepare the initial support. Note that the mask should be initialized first.
        if self.params['support'] == 'auto', nothing is done, and the support will be created after the cdi object,
        to use GPU functions.
        :return: Nothing. self.support is created
        """
        if self.params['support_formula'] is not None:
            self.support = InitSupportShape(formula=self.params['support_formula'], lazy=True)
            return

        elif self.params['support'] in ['square', 'cube', 'circle', 'sphere'] and \
                self.params['support_size'] is not None:
            print("support_size:", self.params['support_size'])
            self.support = InitSupportShape(shape=self.params['support'],
                                            size=self.params['support_size'], lazy=True)
            return

        elif type(self.params['support']) is str:

            if self.params['support'].lower().split(',')[0] in ['object', 'obj']:
                if self.params['object'] is None:
                    raise CDIRunnerException("Error: support=object was used but no object was given")
                if isinstance(self.params['object'], str):
                    if self.params['object'] == "support" or self.params['object'].startswith("support,"):
                        raise CDIRunnerException("Error: support=object was used but no object was given")
                t = self.params['support_threshold']
                s = self.params['support_smooth_width_begin']
                if ',' in self.params['support']:
                    tmp = self.params['support'].split(',')
                    t = float(tmp[1])
                    if len(tmp) > 2:
                        s = float(tmp[2])
                self.support = SupportUpdate(threshold_relative=t, smooth_width=s, method='max', lazy=True,
                                             verbose=True)
                print("support=object: will create the initial support from the supplied object "
                      "using: relative threshold=%4.3f, method=max, smooth=%4.2f" % (t, s))
                return

            if self.params['support'].lower() == 'auto':
                print("No support given. Will use autocorrelation to estimate initial support")
                r = self.params['support_autocorrelation_threshold']
                self.support = AutoCorrelationSupport(r, lazy=True, verbose=True)
                return

            filename = self.params['support']
            print('Loading support from: ', filename)

            # Extract hdf5path if supplied as 'filename:path/to/support'
            if ':' in os.path.splitdrive(filename)[1] and ('h5' in filename or 'cxi' in filename):
                # Looks like an hdf5 file with the hdf5 path
                drive, filename = os.path.splitdrive(filename)
                filename, h5paths = filename.split(':')
                filename = drive + filename
                h5paths = [h5paths]
            else:
                h5paths = ['entry_last/image_1/mask', 'entry_1/image_1/mask']

            ext = os.path.splitext(filename)[-1]
            if ext == '.npy':
                self.support = np.load(filename).astype(np.int8)
            elif ext == '.tif' or ext == '.tiff':
                frames = Image.open(filename)
                frames.seek(0)
                self.support = np.array(frames).astype(np.int8)
            elif ext == '.mat':
                a = list(loadmat(filename).values())
                for v in a:
                    if np.size(v) > 1000:
                        # Avoid matlab strings and attributes, and get the array
                        self.support = np.array(v).astype(np.int8)
                        break
            elif ext == '.npz':
                d = np.load(filename)
                if 'support' in d.keys():
                    self.support = d['support'].astype(np.int8)
                else:
                    # Assume only the support is in the datafile
                    for k, v in d.items():
                        self.support = v.astype(np.int8)
                        break
            elif ext == '.edf':
                self.support = fabio.open(filename).data.astype(np.int8)
            elif ext == '.cxi' or ext == '.h5':
                # this may be a CXI result file, try to find the mask
                with h5py.File(filename, 'r') as h:
                    import_nok = True
                    for hp in h5paths:
                        if hp in h:
                            self.support = h[hp][()].astype(np.int8)
                            print("Loaded support from %s:%s" % (filename, hp))
                            import_nok = False
                            break
                    if import_nok:
                        raise CDIRunnerException(
                            "Supplied support=%s, could not find support e.g. in entry_1/image_1/mask)" % filename)
            else:
                raise CDIRunnerException(
                    "Supplied support=%s, but format not recognized"
                    "(not .npy, .npz, .tif, .edf, .mat, .cxi)" % filename)
            if self.iobs.ndim == 3 and self.support.ndim == 2:
                print("Got a 2D support, expanding to 3D")
                self.support = np.tile(self.support, (len(self.iobs), 1, 1))

            # We can accept a support with a size different from the data, just pad with zeros
            # if smaller than iobs, or crop (assuming it is centered) otherwise
            if self.iobs.shape != self.support.shape:
                print("Reshaping supplied support to fit iobs size",
                      self.support.shape, "->", self.iobs.shape)
                sup = np.zeros_like(self.iobs, dtype=np.int8)
                if self.iobs.ndim == 2:
                    ny, nx = self.iobs.shape
                    nys, nxs = self.support.shape
                    # Crop
                    if ny < nys:
                        self.support = self.support[(nys - ny) // 2:(nys - ny) // 2 + ny]
                    if nx < nxs:
                        self.support = self.support[:, (nxs - nx) // 2:(nxs - nx) // 2 + nx]
                    # Pad
                    nys, nxs = self.support.shape
                    sup[(ny - nys) // 2:(ny - nys) // 2 + nys, (nx - nxs) // 2:(nx - nxs) // 2 + nxs] = self.support
                if self.iobs.ndim == 3:
                    nz, ny, nx = self.iobs.shape
                    nzs, nys, nxs = self.support.shape
                    # Crop
                    if nz < nzs:
                        self.support = self.support[(nzs - nz) // 2:(nzs - nz) // 2 + nz]
                    if ny < nys:
                        self.support = self.support[:, (nys - ny) // 2:(nys - ny) // 2 + ny]
                    if nx < nxs:
                        self.support = self.support[:, :, (nxs - nx) // 2:(nxs - nx) // 2 + nx]
                    # Pad
                    nzs, nys, nxs = self.support.shape
                    sup[(nz - nzs) // 2:(nz - nzs) // 2 + nzs, (ny - nys) // 2:(ny - nys) // 2 + nys,
                    (nx - nxs) // 2:(nx - nxs) // 2 + nxs] = self.support
                self.support = sup
        else:
            print("No support given. Will use autocorrelation to estimate initial support")
            r = self.params['support_autocorrelation_threshold']
            self.support = AutoCorrelationSupport(r, lazy=True, verbose=True)
            return

        nb = self.support.sum()
        p = nb * 100 / self.support.size
        print("Initialized support ", self.support.shape, ", with %d pixels (%6.3f%%)" % (nb, p))
        if p > 20:
            print('WARNING: large number of support pixels !')

    def prepare_data(self, init_mask=True, rebin_data=True, init_support=True, center_crop_data=True,
                     corr_flat_field=True):
        """
        Prepare CDI data for processing.

        :param init_mask: if True (the default), initialize the mask
        :param rebin_data: if True (the default), rebin the data
        :param init_support: if True (the default), initialize the support
        :param center_crop_data: if True (the default), center & crop data
        :param corr_flat_field: if True (the default), correct for the flat field
        :return:
        """
        if init_mask:
            t0 = timeit.default_timer()
            self.init_mask()
            self.timings['init_mask'] = timeit.default_timer() - t0
        if corr_flat_field:
            t0 = timeit.default_timer()
            self.corr_flat_field()
            self.timings['corr_flat_field'] = timeit.default_timer() - t0
        if rebin_data:
            t0 = timeit.default_timer()
            self.rebin_data()
            self.timings['rebin_data'] = timeit.default_timer() - t0
        if init_support:
            t0 = timeit.default_timer()
            self.init_support()
            self.timings['init_support'] = timeit.default_timer() - t0
        if center_crop_data:
            t0 = timeit.default_timer()
            self.center_crop_data()
            self.timings['center_crop_data'] = timeit.default_timer() - t0

    def center_crop_data(self):
        """
        Auto-center diffraction pattern, and resize according to allowed FFT prime decomposition.
        Takes into account ROI if supplied.
        :return:
        """
        self.params['iobs_shape_orig'] = self.iobs.shape
        if self.iobs.ndim == 3:
            nz0, ny0, nx0 = self.iobs.shape
            if self.params['auto_center_resize'] and self.params['roi_user'] is None:
                # Find center of mass
                z0, y0, x0 = center_of_mass(self.iobs)
                print("Center of mass at: %.3f %.3f %.3f" % (z0, y0, x0))
                iz0, iy0, ix0 = int(round(z0)), int(round(y0)), int(round(x0))
            else:
                iz0, iy0, ix0 = int(nz0 // 2), int(ny0 // 2), int(nx0 // 2)
            # Max symmetrical box around center of mass
            nx = 2 * min(ix0, nx0 - ix0)
            ny = 2 * min(iy0, ny0 - iy0)
            nz = 2 * min(iz0, nz0 - iz0)
            if self.params['max_size'] is not None:
                m = self.params['max_size']
                if nx > m:
                    nx = m
                if ny > m:
                    ny = m
                if nz > m:
                    nz = m
            # Crop data to fulfill FFT size requirements
            nz1, ny1, nx1 = smaller_primes((nz, ny, nx), maxprime=self.processing_unit.max_prime_fft_radix(),
                                           required_dividers=(2,))

            print("Centering & reshaping data: (%d, %d, %d) -> (%d, %d, %d)" % (nz0, ny0, nx0, nz1, ny1, nx1))
            self.iobs = self.iobs[iz0 - nz1 // 2:iz0 + nz1 // 2, iy0 - ny1 // 2:iy0 + ny1 // 2,
                        ix0 - nx1 // 2:ix0 + nx1 // 2]
            if self.mask is not None:
                self.mask = self.mask[iz0 - nz1 // 2:iz0 + nz1 // 2, iy0 - ny1 // 2:iy0 + ny1 // 2,
                            ix0 - nx1 // 2:ix0 + nx1 // 2]
            if isinstance(self.support, np.ndarray):
                self.support = self.support[nz0 // 2 - nz1 // 2:nz0 // 2 + nz1 // 2,
                               ny0 // 2 - ny1 // 2:ny0 // 2 + ny1 // 2,
                               nx0 // 2 - nx1 // 2:nx0 // 2 + nx1 // 2]

            # Store the final ROI relative to the uncropped, non-rebinned array, for CXI export
            r = [1, 1, 1]
            if self.params['rebin'] is not None:
                # 'rebin' parameter has already been pre-processed in rebin_data()
                if type(self.params['rebin']) is int:
                    r = [self.params['rebin']] * 3
                elif len(self.params['rebin']) == 1:
                    r = [self.params['rebin']] * 3
                else:
                    r = self.params['rebin']
            izmin, izmax, iymin, iymax, ixmin, ixmax = 0, 0, 0, 0, 0, 0
            if self.params['roi_user'] is not None:
                izmin, izmax, iymin, iymax, ixmin, ixmax = self.params['roi_user']
            self.params['roi_final'] = ((iz0 - nz1 // 2) * r[0] + izmin, (iz0 + nz1 // 2) * r[0] + izmin,
                                        (iy0 - ny1 // 2) * r[1] + iymin, (iy0 + ny1 // 2) * r[1] + iymin,
                                        (ix0 - nx1 // 2) * r[2] + ixmin, (ix0 + nx1 // 2) * r[2] + ixmin)
        else:
            ny0, nx0 = self.iobs.shape
            if self.params['auto_center_resize'] and self.params['roi_user'] is None:
                # Find center of mass
                y0, x0 = center_of_mass(self.iobs)
                print("Center of mass at: %.3f %.3f" % (y0, x0))
                iy0, ix0 = int(round(y0)), int(round(x0))
            else:
                iy0, ix0 = int(ny0 // 2), int(nx0 // 2)
            # Max symmetrical box around center of mass
            nx = 2 * min(ix0, nx0 - ix0)
            ny = 2 * min(iy0, ny0 - iy0)
            if self.params['max_size'] is not None:
                m = self.params['max_size']
                if nx > m:
                    nx = m
                if ny > m:
                    ny = m
            # Crop data to fulfill FFT size requirements
            ny1, nx1 = smaller_primes((ny, nx), maxprime=self.processing_unit.max_prime_fft_radix(),
                                      required_dividers=(2,))

            print("Centering & reshaping data: (%d, %d) -> (%d, %d)" % (ny0, nx0, ny1, nx1))
            self.iobs = self.iobs[iy0 - ny1 // 2:iy0 + ny1 // 2, ix0 - nx1 // 2:ix0 + nx1 // 2]
            if self.mask is not None:
                self.mask = self.mask[iy0 - ny1 // 2:iy0 + ny1 // 2, ix0 - nx1 // 2:ix0 + nx1 // 2]
            if isinstance(self.support, np.ndarray):
                self.support = self.support[ny0 // 2 - ny1 // 2:ny0 // 2 + ny1 // 2,
                               nx0 // 2 - nx1 // 2:nx0 // 2 + nx1 // 2]
                print(self.iobs.shape, self.mask.shape, self.support.shape)
            # Store the final ROI relative to the uncropped, non-rebinned array, for CXI export
            r = [1, 1]
            if self.params['rebin'] is not None:
                # 'rebin' parameter has already been pre-processed in rebin_data()
                if type(self.params['rebin']) is int:
                    r = [self.params['rebin']] * 2
                elif len(self.params['rebin']) == 1:
                    r = [self.params['rebin']] * 2
                else:
                    r = self.params['rebin']
            iymin, iymax, ixmin, ixmax = 0, 0, 0, 0
            if self.params['roi_user'] is not None:
                iymin, iymax, ixmin, ixmax = self.params['roi_user']
            self.params['roi_final'] = ((iy0 - ny1 // 2) * r[0] + iymin, (iy0 + ny1 // 2) * r[0] + iymin,
                                        (ix0 - nx1 // 2) * r[1] + ixmin, (ix0 + nx1 // 2) * r[1] + ixmin)

    def init_object(self):
        """
        Prepare the initial object. It will be random unless a file was supplied for input.
        :return: the initialised object, or an operator which can be used for a lazy initialisation
        """
        if self.iobs is not None:
            shape = self.iobs.shape
        else:
            # Once the cdi object is created, we delete self.iobs to save memory
            shape = self.cdi.iobs.shape

        if self.params['object'] is None:
            # Initialize random object
            if self.params['positivity']:
                obj = InitObjRandom(src="support", amin=0, amax=1, phirange=0)
            else:
                obj = InitObjRandom(src="support", amin=0, amax=1, phirange=2 * np.pi)
            return obj

        elif type(self.params['object']) is str:
            if ',' in self.params['object']:
                filename = self.params['object'].split(',')[0]
                if filename == 'support':
                    # Special case: object=support,0.9,1,0,0.5
                    s = [float(v) for v in self.params['object'].split(',')[1:]]
                    obj = InitObjRandom(src="support", amin=s[0], amax=s[1], phirange=s[2])
                    return obj
            else:
                filename = self.params['object']
            print('Loading object from: ', filename)
            ext = os.path.splitext(filename)[-1]
            if ext == '.npy':
                obj = np.load(filename).astype(np.complex64)
            elif ext == '.mat':
                a = list(loadmat(filename).values())
                for v in a:
                    if np.size(v) > 1000:
                        # Avoid matlab strings and attributes, and get the array
                        obj = np.array(v).astype(np.complex64)
                        break
            elif ext == '.npz':
                d = np.load(filename)
                if 'object' in d.keys():
                    obj = d['object'].astype(np.complex64)
                elif 'obj' in d.keys():
                    obj = d['obj'].astype(np.complex64)
                elif 'data' in d.keys():
                    obj = d['data'].astype(np.complex64)
                else:
                    # Assume only the object is in the datafile
                    for k, v in d.items():
                        obj = v.astype(np.complex64)
                        break
            elif ext == '.cxi':
                # TODO: allow to supply the hdf5 data path
                cxi = h5py.File(filename, 'r')
                if '/entry_1/image_1/data' in cxi:
                    obj = cxi['/entry_1/image_1/data'][()].astype(np.complex64)
                elif '/entry_1/data1/data' in cxi:
                    obj = cxi['/entry_1/data1/data'][()].astype(np.complex64)
                else:
                    raise CDIRunnerException(
                        "Supplied object=%s, but did not find either /entry_1/image_1/data"
                        "or /entry_1/data1/data in the CXI file" % filename)
            elif ext == '.h5':
                # TODO: allow to supply the hdf5 data path
                h = h5py.File(filename, 'r')
                if '/entry_1/data_1/data' in h:
                    obj = h['/entry_1/image_1/data']
                elif '/entry_1/data1/data' in h:
                    obj = h['/entry_1/data1/data']
                else:
                    raise CDIRunnerException(
                        "Supplied object=%s, but did not find either /entry_1/image_1/data"
                        "or /entry_1/data1/data in the CXI file" % filename)
                if obj.ndim == len(shape):
                    obj = obj[()]
                elif obj.ndim == len(shape) + 1:
                    # Likely a modes file, take the first one
                    obj = obj[0]
                else:
                    raise CDIRunnerException(
                        "Supplied object=%s, but the dimensions of the object don't match" % filename)
            else:
                raise CDIRunnerException(
                    "Supplied object=%s, but format not recognized (not .cxi, .npy, .npz, .mat)" % filename)

        # We can accept a smaller object than the data, just expand it as necessary
        if shape != obj.shape:
            for i in range(len(shape)):
                if shape[i] < obj.shape[i]:
                    raise CDIRunnerException('Object', obj.shape, ' must have a <= shape than iobs',
                                             shape, '(smaller objects will be padded with zeros)')
            print("Reshaping object to fit iobs size", obj.shape, "->", shape)
            tmp = np.zeros(shape, dtype=np.complex64)
            if len(shape) == 2:
                ny, nx = shape
                nys, nxs = obj.shape
                tmp[(ny - nys) // 2:(ny - nys) // 2 + nys, (nx - nxs) // 2:(nx - nxs) // 2 + nxs] = obj
            if len(shape) == 3:
                nz, ny, nx = shape
                nzs, nys, nxs = obj.shape
                tmp[(nz - nzs) // 2:(nz - nzs) // 2 + nzs, (ny - nys) // 2:(ny - nys) // 2 + nys,
                (nx - nxs) // 2:(nx - nxs) // 2 + nxs] = obj
            obj = tmp
        if self.cdi is not None:
            self.cdi.set_obj(obj, shift=True)
        if type(self.params['object']) is str:
            if ',' in self.params['object']:
                # Random parameters where supplied in addition to a starting object
                s = [float(v) for v in self.params['object'].split(',')[1:]]
                if self.cdi is not None:
                    obj = InitObjRandom(src="obj", amin=s[0], amax=s[1], phirange=s[2])
                else:
                    obj = InitObjRandom(src=fftshift(obj), amin=s[0], amax=s[1], phirange=s[2])
        print("Finished initializing object ")
        return obj

    def _init_free_pixel_mask(self, iobs):
        if self.params['free_pixel_mask'] is not None:
            filename = self.params['free_pixel_mask']
            print('Loading free pixel mask from: ', filename)
            ext = os.path.splitext(filename)[-1]
            if ext == '.npy':
                free_pixel_mask = np.load(filename).astype(bool)
            elif ext == '.tif' or ext == '.tiff':
                frames = Image.open(filename)
                frames.seek(0)
                free_pixel_mask = np.array(frames).astype(bool)
            elif ext == '.mat':
                a = list(loadmat(filename).values())
                for v in a:
                    if np.size(v) == iobs.size:
                        # Avoid matlab strings and attributes, and get the array
                        free_pixel_mask = np.array(v).astype(bool)
                        break
            elif ext == '.npz':
                d = np.load(filename)
                if 'free_pixel_mask' in d.keys():
                    free_pixel_mask = d['free_pixel_mask'].astype(bool)
                else:
                    # Assume only the mask is in the datafile
                    for k, v in d.items():
                        if v.shape == iobs.shape:
                            free_pixel_mask = v.astype(bool)
                        break
            elif ext == '.edf':
                free_pixel_mask = fabio.open(filename).data.astype(bool)
            elif ext == '.cxi':
                fh5 = h5py.File(filename, mode='r')
                if '/entry_last/image_1/process_1/configuration' in fh5:
                    free_pixel_mask = fh5['/entry_last/image_1/process_1/configuration/free_pixel_mask'][()]
                else:
                    raise CDIRunnerException("Supplied free_pixel_mask=%s, but /entry_last/image_1/process_1/"
                                             "configuration/free_pixel_mask not found !" % filename)
            else:
                raise CDIRunnerException("Supplied mask=%s, but format not recognized (recognized: .npy,"
                                         ".npz, .tif, .edf, .mat)" % filename)
            self.cdi.set_free_pixel_mask(free_pixel_mask, verbose=True, shift=True)
        else:
            self.cdi = InitFreePixels(ratio=0.05, island_radius=3, exclude_zone_center=0.05,
                                      verbose=True, lazy=True) * self.cdi

    def prepare_cdi(self):
        """
        Prepare CDI object from input parameters. If self.cdi already exists, it is re-used, and
        only the initial object and support are updated. To save memory, self.iobs and self.mask
        are set to None.

        :return: nothing. Creates or updates self.cdi object.
        """
        t0 = timeit.default_timer()
        obj = self.init_object()
        self.timings['prepare_cdi: init_object'] = timeit.default_timer() - t0
        t1 = timeit.default_timer()

        sup = None
        if isinstance(self.support, np.ndarray):
            sup = fftshift(self.support)

        if self.cdi is not None:
            # If the object is supplied as an array, the support may be initialised
            # from the support.
            if isinstance(obj, np.ndarray):
                self.cdi.set_obj(obj)

            if isinstance(sup, np.ndarray):
                self.cdi.set_support(sup)
            elif isinstance(self.support, OperatorCDI):
                self.cdi = self.support * self.cdi

            # The object may be initialised from the support
            if isinstance(obj, OperatorCDI):
                self.cdi = obj * self.cdi

            # Reset the PSF
            self.cdi.init_psf(model=None)
        else:
            pix = self.params['pixel_size_detector']
            d = self.params['detector_distance']
            w = self.params['wavelength']
            m = self.mask
            if m is not None:
                m = fftshift(m)
                self.mask = None
            iobs = fftshift(self.iobs)
            self.iobs = None
            gc.collect()
            self.cdi = CDI(iobs=iobs, support=sup, mask=m, pixel_size_detector=pix,
                           wavelength=w, detector_distance=d, obj=None)
            self._algo_s = ""

            # If the object is supplied as an array, the support may be initialised
            # from the support.
            if isinstance(obj, np.ndarray):
                self.cdi.set_obj(obj)

            if isinstance(self.support, OperatorCDI):
                self.cdi = self.support * self.cdi

            # The object may be initialised from the support
            if isinstance(obj, OperatorCDI):
                self.cdi = obj * self.cdi

            gc.collect()
            self.timings['prepare_cdi: create CDI'] = timeit.default_timer() - t1

            t1 = timeit.default_timer()
            self._init_free_pixel_mask(iobs)
            self.timings['prepare_cdi: init_free_pixels'] = timeit.default_timer() - t1

            if self.params['mask_interp'] is not None:
                t1 = timeit.default_timer()
                d, n = self.params['mask_interp']
                print("Interpolating masked pixels with InterpIobsMask(%d, %d)" % (d, n))
                self.cdi = InterpIobsMask(d, n) * self.cdi
                self.timings['prepare_cdi: mask_interp'] = timeit.default_timer() - t1

        if self.params['crop'] is not None:
            if isinstance(obj, OperatorCDI) and isinstance(sup, OperatorCDI):
                # Object and support will be initialised by operators
                self.cdi.set_crop(self.params['crop'], obj_support=False)
            else:
                # This is slower
                self.cdi.set_crop(self.params['crop'])

        self.timings['prepare_cdi'] = timeit.default_timer() - t0

    def prepare_processing_unit(self):
        """
        Prepare processing unit (CUDA, OpenCL, or CPU).

        Returns: nothing. Creates self.processing_unit

        """
        # Avoid calling again select_gpu()
        if default_processing_unit.cu_device is None and default_processing_unit.cl_device is None:
            s = "CDI runner: preparing processing unit"
            if self.params['gpu'] is not None:
                s += " [given GPU name: %s]" % str(self.params['gpu'])
            print(s)
            try:
                default_processing_unit.select_gpu(gpu_name=self.params['gpu'], verbose=True)
            except Exception as ex:
                s0 = "\n  original error: " + str(ex)
                if self.params['gpu'] is not None:
                    s = "Failed initialising GPU. Please check GPU name [%s] or CUDA/OpenCL installation"
                    raise CDIRunnerException(s % str(self.params['gpu']) + s0)
                else:
                    raise CDIRunnerException(
                        "Failed initialising GPU. Please check GPU name or CUDA/OpenCL installation" + s0)
        if default_processing_unit.backend == Backend.CPU:
            force_cpu = self.params['gpu'] in ['cpu', 'CPU']
            if 'PYNX_PU' in os.environ:
                if os.environ['PYNX_PU'].lower() == 'cpu':
                    force_cpu = True
            if force_cpu:
                warnings.warn("Using gpu=CPU - please try to use a GPU !")
            else:
                raise CDIRunnerException("CUDA or OpenCL or GPU not available - you need a GPU to use pynx.CDI !")

        self.processing_unit = default_processing_unit

    def get_psf(self):
        """
        Get the PSF parameters from self.params['psf']

        :return: a tuple (model, fwhm, eta, update_psf) with the PSF parameters,
          or (psf, update_psf) if the PSF array was loaded from a file, or False
          if no psf is used.
        """
        if self.params['psf'] is False:
            return False
        elif self.params['psf'] is True:
            return "pseudo-voigt", 1, 0.05, 5
        v = self.params['psf'].split(',')
        update_psf = 5
        if len(v[0]) > 4:
            psf = None
            if v[0][-4:].lower() == '.npz':
                d = np.load(v[0])
                if 'psf' in d.keys():
                    psf = d['psf']
                elif 'PSF' in d.keys():
                    psf = d['PSF']
                else:
                    # Assume only the PSF is in the datafile
                    for k, v in d.items():
                        psf = v
                        if isinstance(psf, np.ndarray):
                            break
                print("Loaded PSF from %s with shape:" % v[0], psf.shape)
                if psf is None:
                    raise CDIRunnerException("Could not find PSF array in: ", v[0])
                if len(v) > 1:
                    update_psf = int(v[1])
                return psf, update_psf
            elif v[0][-4:].lower() == '.cxi':
                h = h5py.File(v[0], mode='r')
                for e in ['entry_last/image_1/instrument_1/detector_1/point_spread_function0',
                          'entry_1/image_1/instrument_1/detector_1/point_spread_function0']:
                    if e in h:
                        psf = h[e][()]
                print("Loaded PSF from %s with shape:" % v[0], psf.shape)
                if psf is None:
                    raise CDIRunnerException("Could not find PSF array in: ", v[0])
                if len(v) > 1:
                    update_psf = int(v[1])
                return psf, update_psf
        # We got there, so no PSF was imported from a file, use a model
        model = v[0]
        fwhm = float(v[1])
        eta = 0.05
        if 'voigt' in model.lower() or 'pseudo' in model.lower():
            eta = float(v[2])
            if len(v) == 4:
                update_psf = int(v[3])
        elif len(v) == 3:
            update_psf = int(v[2])
        return model, fwhm, eta, update_psf

    def run(self, file_name=None, run_n=1):
        """
        Main work function. Will run selected algorithms according to parameters

        :param file_name: output file_name. If None, the result is not saved.
        :param run_n: the run number
        :return:
        """
        # Starting point of the history of algorithm & timings
        self.cdi.reset_history()
        if self.params['algorithm'] is not None:
            self.run_algorithm(self.params['algorithm'], file_name=file_name, run_n=run_n)
            return

        # We did not get an algorithm string, so create the chain of operators from individual parameters
        nb_raar = self.params['nb_raar']
        nb_hio = self.params['nb_hio']
        nb_er = self.params['nb_er']
        nb_ml = self.params['nb_ml']
        positivity = self.params['positivity']
        support_only_shrink = self.params['support_only_shrink']
        beta = self.params['beta']
        detwin = self.params['detwin']
        psf = self.params['psf']
        live_plot = self.params['live_plot']
        plot_axis = self.params['plot_axis']
        support_update_period = self.params['support_update_period']
        support_smooth_width_begin = self.params['support_smooth_width_begin']
        support_smooth_width_end = self.params['support_smooth_width_end']
        support_threshold = self.params['support_threshold']
        support_threshold_method = self.params['support_threshold_method']
        support_post_expand = self.params['support_post_expand']
        support_update_border_n = self.params['support_update_border_n']
        confmin = self.params['confidence_interval_factor_mask_min']
        confmax = self.params['confidence_interval_factor_mask_max']
        verbose = self.params['verbose']
        zero_mask = self.params['zero_mask']
        min_fraction = self.params['support_fraction_min']
        max_fraction = self.params['support_fraction_max']
        if live_plot is True:
            live_plot = verbose
        elif live_plot is False:
            live_plot = 0

        if nb_ml > 0 and psf:
            print("ML deactivated - PSF is unimplemented in ML")
            nb_ml = 0

        print('No algorithm chain supplied. Proceeding with the following parameters:')
        print(' %30s = ' % 'nb_hio', nb_hio)
        print(' %30s = ' % 'nb_raar', nb_raar)
        print(' %30s = ' % 'nb_er', nb_er)
        print(' %30s = ' % 'nb_ml', nb_ml)
        print(' %30s = ' % 'positivity', positivity)
        print(' %30s = ' % 'support_only_shrink', support_only_shrink)
        print(' %30s = ' % 'beta', beta)
        print(' %30s = ' % 'detwin', detwin)
        print(' %30s = ' % 'live_plot', live_plot)
        print(' %30s = ' % 'support_update_period', support_update_period)
        print(' %30s = ' % 'support_smooth_width_begin', support_smooth_width_begin)
        print(' %30s = ' % 'support_smooth_width_end', support_smooth_width_end)
        print(' %30s = ' % 'support_threshold', support_threshold)
        print(' %30s = ' % 'support_threshold_method', support_threshold_method)
        print(' %30s = ' % 'support_post_expand', support_post_expand)
        if support_update_border_n:
            print(' %30s = ' % 'support_update_border_n', support_update_border_n)
        print(' %30s = ' % 'support_fraction_min/max', min_fraction, max_fraction)
        print(' %30s = ' % 'confidence_interval_factor_mask_min', confmin)
        print(' %30s = ' % 'confidence_interval_factor_mask_max', confmax)
        print(' %30s = ' % 'zero_mask', zero_mask)
        print(' %30s = ' % 'verbose', verbose)
        print(' %30s = ' % 'psf', psf)

        # Write an algorithm string based on the given parameters
        # Begin by listing all cycles for which there is a new operator applied
        cycles = [nb_hio + nb_raar + nb_er + nb_ml]
        if nb_hio > 0:
            cycles.append(nb_hio)
        if nb_raar > 0:
            cycles.append(nb_hio + nb_raar)
        if nb_er > 0:
            cycles.append(nb_hio + nb_raar + nb_er)
        if nb_ml > 0:
            cycles.append(nb_hio + nb_raar + nb_er + nb_ml)
        if support_update_period > 0:
            cycles += list(range(0, nb_hio + nb_raar + nb_er, support_update_period))
        if detwin:
            detwin_cycles = int((nb_hio + nb_raar) // 4)
            cycles.append(detwin_cycles)
        if psf:
            psf_0 = int((nb_raar + nb_hio) * 0.66)
            cycles += [psf_0]

        if type(zero_mask) is str:
            if zero_mask.lower() == 'auto':
                zm = 1
                zm_cycle = int((nb_hio + nb_raar) * 0.6)
                cycles.append(zm_cycle)
            else:
                zm_cycle = None
                if zero_mask.lower() == 'true' or zero_mask.lower() == '1':
                    zm = True
                else:
                    zm = False
        else:
            zm_cycle = None
            zm = zero_mask  # Can be 0, 1, True, False..

        cycles.sort()

        algo = 1  # Start with unity operator
        algo_str = []
        while len(cycles) > 1:
            algo1 = 1
            i0, i1 = cycles[:2]
            n = i1 - i0
            cycles.pop(0)

            if i0 == i1:
                continue
            update_psf = 0
            psf_filter = self.params['psf_filter']
            if psf:
                if i1 == psf_0:
                    psf = self.get_psf()
                    if len(psf) == 2:
                        algo1 = InitPSF(psf=psf[0]) * algo1
                    else:
                        algo1 = InitPSF(model=psf[0], fwhm=psf[1], eta=psf[2])
                    print('Activating PSF after %d cycles, updated every %d cycles' % (psf_0, psf[-1]))
                if i1 >= psf_0:
                    update_psf = psf[-1]

            if zm_cycle is not None:
                if i1 == zm_cycle:
                    # Free masked pixels
                    zm = 0
                    print('Switching from zero_mask=1 to 0 after %d cycles' % zm_cycle)
            if i1 <= nb_hio:
                algo1 = HIO(beta=beta, positivity=positivity, calc_llk=verbose, show_cdi=live_plot,
                            zero_mask=zm, confidence_interval_factor_mask_min=confmin,
                            confidence_interval_factor_mask_max=confmax, update_psf=update_psf,
                            psf_filter=psf_filter, plot_axis=plot_axis) ** n * algo1
            elif nb_raar > 0 and i1 <= (nb_hio + nb_raar):
                algo1 = RAAR(beta=beta, positivity=positivity, calc_llk=verbose, show_cdi=live_plot,
                             zero_mask=zm, confidence_interval_factor_mask_min=confmin,
                             confidence_interval_factor_mask_max=confmax, update_psf=update_psf,
                             psf_filter=psf_filter, plot_axis=plot_axis) ** n * algo1
            elif nb_er > 0 and i1 <= (nb_hio + nb_raar + nb_er):
                algo1 = ER(positivity=positivity, calc_llk=verbose, show_cdi=live_plot, zero_mask=zm,
                           confidence_interval_factor_mask_min=confmin,
                           confidence_interval_factor_mask_max=confmax, update_psf=update_psf,
                           psf_filter=psf_filter, plot_axis=plot_axis) ** n * algo1
            elif nb_ml > 0 and i1 <= (nb_hio + nb_raar + nb_er + nb_ml):
                algo1 = ML(calc_llk=verbose, show_cdi=live_plot, plot_axis=plot_axis) ** n * \
                        FourierApplyAmplitude(calc_llk=True, zero_mask=zm) * algo1

            if support_update_period > 0:
                if i1 % support_update_period == 0 and len(cycles) > 1:
                    s = support_smooth_width_begin, support_smooth_width_end, nb_raar + nb_hio
                    algo1 = SupportUpdate(threshold_relative=support_threshold, smooth_width=s,
                                          force_shrink=support_only_shrink, method=support_threshold_method,
                                          post_expand=support_post_expand,
                                          update_border_n=support_update_border_n, min_fraction=min_fraction,
                                          max_fraction=max_fraction) * algo1

            if detwin:
                if i1 == detwin_cycles:
                    if i1 <= nb_hio:
                        algo1 = DetwinHIO(beta=beta, positivity=False, nb_cycle=10, zero_mask=zm) * algo1
                    else:
                        algo1 = DetwinRAAR(beta=beta, positivity=False, nb_cycle=10, zero_mask=zm) * algo1

            algo_str.append(str(algo1))
            if algo is None:
                algo = algo1
            else:
                algo = algo1 * algo

        # Finish by FourierApplyAmplitude
        confmin = self.params['confidence_interval_factor_mask_min']
        confmax = self.params['confidence_interval_factor_mask_max']
        algo = FourierApplyAmplitude(calc_llk=True, zero_mask=zm,
                                     confidence_interval_factor_mask_max=confmax,
                                     confidence_interval_factor_mask_min=confmin, obj_stats=True) * algo

        # algo = PRTF(file_name='%s-PRTF.png' % os.path.splitext(file_name)[0], fig_title=file_name) * algo

        # print("Algorithm chain: ", algo)

        # More pretty printing:
        s = ''
        while len(algo_str):
            s0 = algo_str.pop(0)
            n = 1
            while len(algo_str):
                if algo_str[0] == s0:
                    del algo_str[0]
                    n += 1
                else:
                    break
            if s != '':
                s = ' * ' + s
            if n > 1:
                s = '(%s)**%d' % (s0, n) + s
            else:
                s = s0 + s

        # Reformat chain so that it can be re-used from the command-line
        s = s.replace('()', '')
        s = s.replace('SupportUpdate', 'Sup')
        if psf:
            s = s.replace('*InitPSF *', ',psf=%d,' % psf[-1])

        print("Algorithm chain: ", s)

        # Now execute the recipe !
        t0 = timeit.default_timer()
        self.cdi = algo * self.cdi

        print("\nTotal elapsed time for algorithms: %8.2fs" % (timeit.default_timer() - t0))
        calc_throughput(self.cdi, verbose=True)
        if file_name is not None:
            if self.params['save'] not in ['none', None]:
                self.save_result(file_name)
            if self.params['save_plot']:
                self.save_plot(os.path.splitext(file_name)[0] + '.png', algo_string=s)

    def run_algorithm(self, algo_string, file_name=None, run_n=1):
        """
        Run a single or suite of algorithms in a given run.

        :param algo_string: a single or suite of algorithm steps to use, which can either correspond to a change
                            of parameters, i.e. 'beta=0.5', 'support_threshold=0.3', 'positivity=1',
                            or operators which should be applied to the cdi object, such as:
                            'hio**20': 20 cycles of HIO
                            'er': a single cycle or error reduction
                            'detwinhio**20': 20 cycles of hio after halving the support
                            'er**20*hio**100': 100 cycles of hio followed by 20 cycles of ER
                            '(sup*er**10*raar**100)**10': 10 times [100 cycles of RAAR + 10 ER + Support update]
        :param file_name: output file_name. If None, the result is not saved.
        :param run_n: the run number
        """
        algo_split = algo_string.split(',')
        algo_split.reverse()
        print(algo_split)
        t0 = timeit.default_timer()
        update_psf = 0
        for algo in algo_split:
            if self._algo_s == "":
                self._algo_s += algo
            else:
                self._algo_s += ',' + algo
            print("\n", "#" * 100, "\n#", "\n#         Run: %g , Algorithm: %s\n#\n" % (run_n, algo), "#" * 100)
            realoptim = False  # Is this a real optimization (HIO, ER, ML, RAAR,...), or just a change of parameter ?
            if algo.lower().find('beta=') >= 0:
                self.params['beta'] = float(algo.lower().split('beta=')[-1])
            elif algo.lower().find('support_threshold=') >= 0:
                self.params['support_threshold'] = float(algo.lower().split('support_threshold=')[-1])
            elif algo.lower().find('support_threshold_mult=') >= 0:
                self.params['support_threshold'] *= float(algo.lower().split('support_threshold_mult=')[-1])
            elif algo.lower().find('gps_inertia=') >= 0:
                self.params['gps_inertia'] = float(algo.lower().split('gps_inertia=')[-1])
            elif algo.lower().find('gps_t=') >= 0:
                self.params['gps_t'] = float(algo.lower().split('gps_t=')[-1])
            elif algo.lower().find('gps_s=') >= 0:
                self.params['gps_s'] = float(algo.lower().split('gps_s=')[-1])
            elif algo.lower().find('gps_sigma_o=') >= 0:
                self.params['gps_sigma_o'] = float(algo.lower().split('gps_sigma_o=')[-1])
            elif algo.lower().find('gps_sigma_f=') >= 0:
                self.params['gps_sigma_f'] = float(algo.lower().split('gps_sigma_f=')[-1])
            elif algo.lower().find('positivity=') >= 0:
                self.params['positivity'] = int(algo.lower().split('positivity=')[-1])
            elif algo.lower().find('zero_mask=') >= 0:
                self.params['zero_mask'] = int(algo.lower().split('zero_mask=')[-1])
            elif algo.lower().find('support_only_shrink=') >= 0:
                self.params['support_only_shrink'] = int(algo.lower().split('support_only_shrink=')[-1])
            elif algo.lower().find('positivity=') >= 0:
                self.params['positivity'] = float(algo.lower().split('positivity=')[-1])
            elif algo.lower().find('support_smooth_width_begin=') >= 0:
                self.params['support_smooth_width_begin'] = float(algo.lower().split('support_smooth_width_begin=')[-1])
            elif algo.lower().find('support_smooth_width_end=') >= 0:
                self.params['support_smooth_width_end'] = float(algo.lower().split('positivity=')[-1])
            elif algo.lower().find('support_post_expand=') >= 0:
                s = algo.lower().split('support_post_expand=')[-1].replace('#', ',')
                self.params['support_post_expand'] = eval(s)
            elif algo.lower().find('support_threshold_method=') >= 0:
                self.params['support_threshold_method'] = algo.lower().split('support_post_expand=')[-1]
            elif algo.lower().find('support_update_border_n=') >= 0:
                self.params['support_update_border_n'] = int(algo.lower().split('support_update_border_n=')[-1])
            elif algo.lower().find('confidence_interval_factor_mask_min=') >= 0:
                self.params['confidence_interval_factor_mask_min'] = \
                    algo.lower().split('confidence_interval_factor_mask_min=')[-1]
            elif algo.lower().find('confidence_interval_factor_mask_max=') >= 0:
                self.params['confidence_interval_factor_mask_max'] = \
                    algo.lower().split('confidence_interval_factor_mask_max=')[-1]
            # elif algo.lower().find('crop=') >= 0:
            #     b = algo.lower().split('crop=')[-1]
            #     if len(b) > 1:
            #         # parameters per axis have been given, e.g. crop=122
            #         b = [int(v) for v in b]
            #     else:
            #         b = int(b)
            #     self.cdi.set_crop(b)
            #     # self.cdi = ScaleObj(method="F", verbose=True) * self.cdi
            # elif algo.lower().find('bin=') >= 0:
            #     b = algo.lower().split('bin=')[-1]
            #     if len(b) > 1:
            #         # parameters per axis have been given, e.g. bin=122
            #         b = [int(v) for v in b]
            #     else:
            #         b = int(b)
            #     self.cdi.set_bin(b)
            #     # TODO: rescale object to take into account bin/skip
            #     # self.cdi = ScaleObj(method="F", verbose=True) * self.cdi
            # elif algo.lower().find('skip=') >= 0:
            #     b = algo.lower().split('skip=')[-1]
            #     if len(b) > 1:
            #         # parameters per axis have been given, e.g. skip=122
            #         b = [int(v) for v in b]
            #         if len(b) == self.cdi.iobs.ndim:
            #             # Skip is used, but the shifts were not given select them randomly
            #             for i in range(self.cdi.iobs.ndim):
            #                 b.append(np.random.randint(0, b[i]))
            #     else:
            #         b = int(b)
            #     print("Using bin(skip) parameters:", b)
            #     self.cdi.set_bin(b)
            #     # TODO: rescale object to take into account bin/skip
            #     # self.cdi = ScaleObj(method="F", verbose=True) * self.cdi
            elif algo.lower().find('upsample=') >= 0:
                self.cdi.set_upsample(eval(algo.lower().split('upsample=')[-1]))
                # self.cdi = ScaleObj(method="F", verbose=True) * self.cdi
            elif algo.lower() == "crop2bin":
                v = self.cdi.get_crop()
                if v is None:
                    raise CDIRunnerException("Using 'crop2bin' but crop was not used first.")
                if self.processing_unit.backend in [Backend.PYCUDA]:
                    self.cdi.set_bin(v, obj_support=False)
                    # Zoom in object & support on GPU
                    bin_f = _as_dim(v, ndim=self.cdi.iobs.ndim)
                    self.cdi = Zoom(bin_f) * self.cdi
                else:
                    self.cdi.set_bin(v)
                self.cdi = ScaleObj(method="F", verbose=True) * self.cdi
            elif algo.lower().find('bin=') >= 0:
                v = eval(algo.lower().split('bin=')[-1])
                self.cdi.set_bin(v)
            elif algo.lower().find('crop=') >= 0:
                v = eval(algo.lower().split('crop=')[-1])
                self.cdi.set_crop(v)
            elif 'psf_init=' in algo.lower():
                tmp = algo.lower().split('psf_init=')[-1].split("@")
                fwhm = float(eval(tmp[1]))
                if len(tmp) > 2:
                    eta = float(eval(tmp[2]))
                    print("Init PSF with model = %s, fwhm=%5.2f pixels" % (tmp[0], fwhm))
                else:
                    eta = 0.05
                    print("Init PSF with model = %s, fwhm=%5.2f pixels, eta = %5.3f" % (tmp[0], fwhm, eta))
                self.cdi = InitPSF(model=tmp[0], fwhm=fwhm, eta=eta) * self.cdi
                if update_psf == 0:
                    update_psf = 5
                    print("Will update PSF every %d cycles" % update_psf)
            elif 'psf=' in algo.lower():
                update_psf = int(eval(algo.lower().split('psf=')[-1]))
            elif 'psf_filter' in algo.lower():
                psf_filter = algo.lower().split('psf_filter=')[-1]
                if psf_filter in ['hann', 'tukey']:
                    self.params['psf_filter'] = psf_filter
                    warnings.warn("psf_filter is experimental, you probably should not use it")
                else:
                    self.params['psf_filter'] = None
            elif algo.lower().find('verbose=') >= 0:
                self.params['verbose'] = int(algo.lower().split('verbose=')[-1])
            elif algo.lower().find('live_plot=') >= 0:
                self.params['live_plot'] = int(algo.lower().split('live_plot=')[-1])
            elif algo.lower().find('fig_num=') >= 0:
                self.params['fig_num'] = int(algo.lower().split('fig_num=')[-1])
            elif algo.lower().find('manual') >= 0:
                pass  # Dummy type of algorithm
            else:
                # Not a keyword, so this must be an algorithm to run. Finally !
                realoptim = True
                positivity = self.params['positivity']
                support_only_shrink = self.params['support_only_shrink']
                beta = self.params['beta']
                support_smooth_width_begin = self.params['support_smooth_width_begin']
                support_smooth_width_end = self.params['support_smooth_width_end']
                support_smooth_width_relax_n = self.params['support_smooth_width_relax_n']
                smooth_width = support_smooth_width_begin, support_smooth_width_end, support_smooth_width_relax_n
                support_threshold = self.params['support_threshold']
                gps_inertia = self.params['gps_inertia']
                gps_t = self.params['gps_t']
                gps_s = self.params['gps_s']
                gps_sigma_o = self.params['gps_sigma_o']
                gps_sigma_f = self.params['gps_sigma_f']
                support_threshold_method = self.params['support_threshold_method']
                support_post_expand = self.params['support_post_expand']
                support_update_border_n = self.params['support_update_border_n']
                min_fraction = self.params['support_fraction_min']
                max_fraction = self.params['support_fraction_max']
                confmin = self.params['confidence_interval_factor_mask_min']
                confmax = self.params['confidence_interval_factor_mask_max']
                verbose = int(self.params['verbose'])
                live_plot = self.params['live_plot']
                plot_axis = self.params['plot_axis']
                fig_num = self.params['fig_num']
                zm = self.params['zero_mask']
                psf_filter = self.params['psf_filter']
                if type(zm) is str:
                    if zm.lower() == 'auto':
                        # TODO: better handle zero_mask='auto' for custom algorithms ?
                        zm = False
                    else:
                        if zm.lower() == 'true' or zm.lower() == '1':
                            zm = True
                        else:
                            zm = False
                if int(live_plot) == 1:  # That's for live_plot=True
                    live_plot = verbose

                # Create elementary operators
                fap = FourierApplyAmplitude(calc_llk=True, zero_mask=zm,
                                            confidence_interval_factor_mask_max=confmax,
                                            confidence_interval_factor_mask_min=confmin)
                er = ER(positivity=positivity, calc_llk=verbose, show_cdi=live_plot, fig_num=fig_num, zero_mask=zm,
                        confidence_interval_factor_mask_min=confmin, plot_axis=plot_axis,
                        confidence_interval_factor_mask_max=confmax, update_psf=update_psf, psf_filter=psf_filter)
                hio = HIO(beta=beta, positivity=positivity, calc_llk=verbose, show_cdi=live_plot, fig_num=fig_num,
                          zero_mask=zm, confidence_interval_factor_mask_min=confmin, plot_axis=plot_axis,
                          confidence_interval_factor_mask_max=confmax, update_psf=update_psf, psf_filter=psf_filter)
                raar = RAAR(beta=beta, positivity=positivity, calc_llk=verbose, show_cdi=live_plot, fig_num=fig_num,
                            zero_mask=zm, confidence_interval_factor_mask_min=confmin, plot_axis=plot_axis,
                            confidence_interval_factor_mask_max=confmax, update_psf=update_psf, psf_filter=psf_filter)
                gps = GPS(inertia=gps_inertia, t=gps_t, s=gps_s, sigma_f=gps_sigma_f, sigma_o=gps_sigma_o,
                          positivity=positivity, calc_llk=verbose, show_cdi=live_plot, fig_num=fig_num, zero_mask=zm,
                          confidence_interval_factor_mask_min=confmin, plot_axis=plot_axis,
                          confidence_interval_factor_mask_max=confmax, update_psf=update_psf, psf_filter=psf_filter)
                detwinhio = DetwinHIO(detwin_axis=0, beta=beta, positivity=positivity, zero_mask=zm)
                detwinhio1 = DetwinHIO(detwin_axis=1, beta=beta, positivity=positivity, zero_mask=zm)
                detwinhio2 = DetwinHIO(detwin_axis=2, beta=beta, positivity=positivity, zero_mask=zm)
                detwinraar = DetwinRAAR(detwin_axis=0, beta=beta, positivity=positivity, zero_mask=zm)
                detwinraar1 = DetwinRAAR(detwin_axis=1, beta=beta, positivity=positivity, zero_mask=zm)
                detwinraar2 = DetwinRAAR(detwin_axis=2, beta=beta, positivity=positivity, zero_mask=zm)
                cf = CF(positivity=positivity, calc_llk=verbose, show_cdi=live_plot, fig_num=fig_num, zero_mask=zm,
                        update_psf=update_psf, psf_filter=psf_filter)
                ml = ML(calc_llk=verbose, show_cdi=live_plot, fig_num=fig_num, plot_axis=plot_axis) * fap
                initpsf = InitPSF()
                sup = SupportUpdate(threshold_relative=support_threshold, smooth_width=smooth_width,
                                    force_shrink=support_only_shrink, post_expand=support_post_expand,
                                    method=support_threshold_method, update_border_n=support_update_border_n,
                                    min_fraction=min_fraction, max_fraction=max_fraction)
                supportupdate = sup
                showcdi = ShowCDI(fig_num=fig_num)

                try:
                    ops = eval(algo.lower())
                    self.cdi = ops * self.cdi
                except SupportTooSmall:
                    raise
                except SupportTooLarge:
                    raise
                except Exception as ex:
                    print(traceback.format_exc())
                    print('\n\n Caught exception for scan %s: %s    \n' % (str(self.scan), str(ex)))
                    print('Could not interpret operator-based algorithm: ', algo)
                    print(f'Use {sys.argv[0]} --help for details notably on --algorithm syntax')
                    sys.exit(1)
            if self.params['save'] == 'all' and realoptim and file_name is not None:
                # Finish with FourierApplyAmplitude
                zm = self.params['zero_mask']
                if type(zm) is str:
                    if zm.lower() == 'auto':
                        zm = False
                    else:
                        if zm.lower() == 'true' or zm.lower() == '1':
                            zm = True
                        else:
                            zm = False
                confmin = self.params['confidence_interval_factor_mask_min']
                confmax = self.params['confidence_interval_factor_mask_max']
                self.cdi = FourierApplyAmplitude(calc_llk=True, zero_mask=zm,
                                                 confidence_interval_factor_mask_max=confmax,
                                                 confidence_interval_factor_mask_min=confmin) * self.cdi
                self.save_result(file_name)

        self.timings['run_algorithm (algorithms)'] = timeit.default_timer() - t0
        t1 = timeit.default_timer()
        if self.params['save'] not in ['all', 'none', None] and file_name is not None:
            # Finish with FourierApplyAmplitude
            zm = self.params['zero_mask']
            if type(zm) is str:
                if zm.lower() == 'auto':
                    zm = False
                else:
                    if zm.lower() == 'true' or zm.lower() == '1':
                        zm = True
                    else:
                        zm = False
            confmin = self.params['confidence_interval_factor_mask_min']
            confmax = self.params['confidence_interval_factor_mask_max']
            self.cdi = FourierApplyAmplitude(calc_llk=True, zero_mask=zm,
                                             confidence_interval_factor_mask_max=confmax,
                                             confidence_interval_factor_mask_min=confmin) * self.cdi
            self.save_result(file_name)

        print("\nTotal elapsed time for algorithms & saving: %8.2fs" % (timeit.default_timer() - t0))

        calc_throughput(self.cdi, verbose=True)
        if file_name is not None and self.params['save_plot']:
            self.save_plot(os.path.splitext(file_name)[0] + '.png', algo_string=algo_string)
        self.timings['run_algorithm (final save)'] = timeit.default_timer() - t1

    def save_result(self, file_name=None, verbose=True):
        """ Save the results (object, support,...) at the end of a run

        :param file_name: the filename to save the data to. If its extension is either npz or cxi, this will
                          override params['output_format']. Otherwise, the extension will be replaced. Note that
                          the full filename will be appended with the llk value.
        :return:
        """
        if self.params['output_format'].lower() == 'none':
            return

        if file_name is None:
            file_name = "Result.%s" % self.params['output_format']

        filename, ext = os.path.splitext(os.path.split(file_name)[1])

        if len(ext) < 2:
            ext = "." + self.params['output_format']

        llk = self.cdi.get_llk(normalized=True)

        if llk is not None and self.params['save'] != 'all':
            filename += "_LLKf%08.4f_LLK%08.4f" % (llk[3], llk[0])

        filename += "_SupportThreshold%7.5f" % self.params['support_threshold']

        if ext not in ['.npz', '.cxi']:
            warnings.warn("CDI.save_result(): output format (%s) not understood, defaulting to CXI" % ext)
            ext = '.cxi'
        filename += ext

        if os.path.isfile(filename) and self.params['save'] != 'all':
            warnings.warn("CDI.save_result(): output file already exists, not overwriting: %s" % filename)
            # raise CDIRunnerException("ERROR: output file already exists, no overwriting: %s" % filename)
            filename = filename[:-4]
            nn = 1
            while os.path.isfile("%s_%02d%s" % (filename, nn, ext)):
                nn += 1
            filename = "%s_%02d%s" % (filename, nn, ext)

        if verbose:
            print("Saving result to: %s" % filename)

        if ext == '.npz':
            np.savez_compressed(filename, obj=self.cdi.get_obj(shift=True))
        else:
            process = {}

            if self.params['algorithm'] is not None:
                process["algorithm"] = self.params['algorithm']

            params_string = ""
            for p in self.params.items():
                k, v = p
                if v is not None and k not in ['output_format']:
                    params_string += "%s = %s\n" % (k, str(v))

            process["note"] = {'configuration parameters used for phasing': params_string}

            append = False
            if self.params['save'] == 'all':
                append = True

            self.cdi.save_obj_cxi(filename, sample_name=self.params['sample_name'], experiment_id=None,
                                  instrument=self.params['instrument'], note=self.params['note'],
                                  crop=int(self.params['crop_output']), process_notes=process,
                                  process_parameters=self.params, append=append)
            if os.name == 'posix':
                try:
                    sf = os.path.split(filename)
                    os.system('ln -sf "%s" %s' % (sf[1], os.path.join(sf[0], 'latest.cxi')))
                except:
                    pass

        llk = self.cdi.get_llk()
        d = {'file_name': filename, 'llk_poisson': llk[0], 'llk_gaussian': llk[1], 'llk_euclidian': llk[2],
             'llk_poisson_free': llk[3], 'llk_gaussian_free': llk[4], 'llk_euclidian_free': llk[5],
             'nb_point_support': self.cdi.nb_point_support, 'obj2_out': self.cdi._obj2_out}
        replaced = False
        for n, v in enumerate(self.saved_results):
            if v['file_name'] == filename:
                self.saved_results[n] = d
                replaced = True
                break
        if replaced is False:
            self.saved_results.append(d)
        if verbose:
            print("To view the result file (HDF5/CXI or npz), use: silx view %s" % filename)

    def save_plot(self, file_name, algo_string=None):
        """
        Save a final plot for the object (only supported for 3D data)
        :param file_name: filename for the plot
        :param algo_string: will be used as subtitle
        :return: nothing
        """
        if self.cdi is None:
            return
        if self.cdi.iobs.ndim != 3:
            print("save_plot is only supported for 3D CDI")
            return
        print("Saving plot to: ", file_name)
        p = {'cwd': os.getcwd()}
        for k, v in self.params.items():
            if v is not None and k not in ['live_plot', 'verbose', 'data2cxi', 'output_format', 'save_plot',
                                           'fig_num', 'user_id01_specfile_header', 'user_id01_specfile_scan_header']:
                p[k] = v
        if isinstance(self.params['save_plot'], str):
            plot_type = self.params['save_plot']
        else:
            if self.params['positivity']:
                plot_type = 'abs'
            else:
                plot_type = 'rgba'
        title = self.params['data']
        if 'specfile' in self.params:
            if self.params['specfile'] is not None and self.scan is not None:
                title = ("%s  #%s" % (self.params['specfile'], str(self.scan)))

        show_cdi(self.cdi, params=p, save_plot=file_name, display_plot=False, figsize=(10, 10),
                 crop=2, plot_type=plot_type, title=title, subtitle=algo_string)

        if os.name == 'posix':
            try:
                sf = os.path.split(file_name)
                os.system('ln -sf "%s" %s' % (sf[1], os.path.join(sf[0], 'latest.png')))
            except:
                pass


class CDIRunner:
    """
    Class to process CDI data, with parameters from the command-line or from a text file.
    """

    def __init__(self, argv, params, *args, **kwargs):
        """

        :param argv: the command-line parameters
        :param params: parameters for the optimization, with some default values.
        :param arg: ignored. Only for backward compatibility
        :param kwargs: ignored. Only for backward compatibility
        """
        self.params = copy.deepcopy(params)
        self.argv = argv
        # Must be overwritten in derived classes with the class derived from CDIRunnerScan
        self.CDIRunnerScan = None

        self.parameters_file_name = None
        self.timings = {}

        self.mpi_master = True  # True even if MPI is not used
        if MPI is not None:
            self.mpic = MPI.COMM_WORLD
            self.mpi_master = self.mpic.Get_rank() == 0
            self.mpi_size = self.mpic.Get_size()
            self.mpi_rank = self.mpic.Get_rank()

        self.warn_argparse = None
        self.parse_arg()
        if 'help' not in self.argv and '--help' not in self.argv:
            self.check_params()

    def print(self, *args, **kwargs):
        """
        MPI-aware print function. Non-master processes will be muted
        :param args: args passed to print
        :param kwargs: kwrags passed to print
        :return: nothing
        """
        if self.mpi_master:
            print(*args, **kwargs)

    @classmethod
    def make_parser(cls, script_name="pynx-cdi", description=None, epilog=None, default_par=None):
        """
        Create the parser for the command-line arguments.
        This should be updated in derived objects.

        :param script_name: name of the script, e.g. 'pynx-cdi-id01'. By default,
          will be read from the command-line arguments.
        :param description: initial description string (top of help text)
        :param epilog: epilog for the parser (usually with examples)
        :param default_par: default parameters. May be changed in derived parsers
        :return: the parser
        """
        if default_par is None:
            default_par = default_params
        if description is None:
            description = "Generic ptycho runner"
        if epilog is None:
            epilog = "Example command-line go here:..."
        return make_parser(script_name, description=description,
                                  epilog=epilog, default_params=default_par)

    def parse_parameters_file(self, filename, argv):
        """
        Read parameters from a text file, written with one parameter per line. The parameters
        are converted to command-line arguments (e.g. ``data = myfile.h5`` is converted to
        ``--data myfile.h5`` and appended to an argv-style list.

        :param filename: the file to read the parameters from.
        :return: nothing. The parameters are appended to the supplied argv list.
        """
        self.parameters_file_name = filename
        ll = open(filename).readlines()
        for l in ll:
            i = l.find('#')
            if i >= 0:
                l = l[:i]
            if len(l.strip()) < 4:
                continue
            if l.strip()[0] == '#':
                continue
            s = l.find('=')
            if 0 < s < (len(l) - 1):
                k = l[:s].lower().strip()
                v = l[s + 1:].strip()
                v = v.replace("'", "")
                v = v.replace('"', '')
                if v not in ['false', 'False', 'none', 'None']:
                    # The 'False'/'None' case may be when the parameters uses 'store_true'
                    # or, in that case we may just skip the input.
                    argv.append('--' + k)
                    if k in ['support_size', 'rebin', 'mask', 'roi',
                             'mask_interp', 'support_threshold',
                             'support_autocorrelation_threshold',
                             'support_post_expand'] and ',' in v:
                        argv += v.replace(' ', '').split(',')
                    elif v not in ['true', 'True']:
                        # If 'True' is used, then this is an argument which uses 'store_true' and may not
                        # accept 'True' as a value (e.g. --data2cxi)
                        argv.append(v)

    def parse_arg(self):
        """
        Parses the arguments given on a command line

        Returns: nothing

        """
        t0 = time.time()

        if self.mpi_master:
            # argparse format or not ?
            execs = os.path.split(sys.argv[0])[-1] + ' '
            cmd = execs + ' '.join(sys.argv[1:])
            if cmd.count('--') < 1:  # KLUDGE...
                # Convert arguments to argparse format
                argv = [sys.argv[0]]
                for arg in sys.argv[1:]:
                    arg = arg.strip()
                    if len(arg):  # weird but ' ' can occur in argv...
                        if arg.find('.txt') > 0:
                            print('Using parameters file: ', arg)
                            self.parse_parameters_file(filename=arg, argv=argv)
                        else:
                            a = arg
                            if '=' in a:
                                k, v = a.split('=', maxsplit=1)
                                argv.append('--' + k)
                                # Handle special cases of arguments which need to be comma-separated
                                if k in ['support_size', 'rebin', 'mask', 'roi',
                                         'mask_interp', 'support_threshold',
                                         'support_autocorrelation_threshold',
                                         'support_post_expand'] and ',' in v:
                                    argv += v.split(',')
                                else:
                                    argv.append(v)
                            else:
                                argv.append('--' + a)
                            self.warn_argparse = \
                                "###########################################################################\n" \
                                "DEPRECATION WARNING:\n" \
                                "  It seems that you are still using the old-style arguments, " \
                                "so your command-line was converted from: \n" \
                                f"   {cmd}\n" \
                                "To:\n" \
                                f"  {execs + ' '.join(argv[1:])}\n\n" \
                                "Please switch to the new format, as the old won't be supported " \
                                "in PyNX versions released >2024.\n\n" \
                                f"Use '{execs}--help' for the command-line help \n" \
                                "###########################################################################"

                            print(f"WARNING: converting old-style command-line from:\n  {cmd} "
                                  f"\nto: \n  {' '.join(argv)}\n")
                sys.argv = argv
            self.params = vars(self.make_parser().parse_args())

            # TODO: remove these from parameters ? They should probably be
            #  member variables instead.
            for k in ['roi_final', 'roi_user', 'crop']:
                if k not in self.params:
                    self.params[k] = default_params[k]
            for k in default_params.keys():
                if k not in self.params:
                    print(f"WARNING: missing parameter in self.params: {k}")

        if MPI is not None:
            self.params = self.mpic.bcast(self.params, root=0)

    def check_params(self):
        """
        Check if self.params includes a minimal set of valid parameters
        If 'help' is given as command-line parameter, the help text is printed and the program exits.

        Returns: Nothing. Will raise an exception if necessary
        """
        if self.params['psf'] is True:
            raise CDIRunnerException(
                'PSF: using "psf" (on the command-line or psf=True (input file) is not longer accepted. '
                'You must specify a model using e.g. psf=gaussian,0.5,10 (command-line) '
                'or  psf="pseudo-voigt,0.5,0.05,10" (using quotes in an input file). '
                'See the helptext for details and recommended values depending on the level of coherence.')

        self.check_params_beamline()

    def check_params_beamline(self):
        """
        Check if self.params includes a minimal set of valid parameters, specific to a beamline

        Returns: Nothing. Will raise an exception if necessary
        """
        pass

    def check_mpi(self, task, timeout=60):
        """

        :param task: the string describing the task after which the check is made
        :param timeout: the timeout in seconds after which the synchronisation will be considered failed,
            and an exception will be raised.
        :return: nothing.
        :raises: a CDIRunnerException
        """

        if MPI is not None:
            if self.mpi_size > 1:
                b = self.mpic.Ibarrier()
                nb = int(timeout * 10)
                while not b.Test() and nb > 0:
                    time.sleep(0.1)
                    nb -= 1
                    if nb % 50 == 0:
                        print("MPI master waiting for %s from all tasks..." % task)
                if nb == 0:
                    raise CDIRunnerException("Failure in MPI communication (60s timeout)."
                                             "%s failed on one node ?" % task)
                else:
                    if self.mpi_master:
                        print("%s for %d MPI processes OK" % (task, self.mpi_size))

    def process_scans(self):
        """
        Run all the analysis on the supplied scan list, unless 'help' is given as a
        command-line argument.

        :return: Nothing
        """
        t0 = timeit.default_timer()

        if 'scan' in self.params:
            if isinstance(self.params['scan'], str):
                if 'range' in self.params['scan']:
                    # use range() function to specify a series of scans
                    vscan = eval(self.params['scan'])
                else:
                    # + is used to combine several scans in one, but there may still be a series of scans, e.g. '1+2,3+4'
                    vscan = self.params['scan'].split(',')
            else:
                # This could be None, the default value for self.params['scan'], to load a CXI, numpy, tiff file, etc...
                vscan = [self.params['scan']]
        else:
            vscan = [None]

        if MPI is not None:
            if 'scan' in self.params['mpi'] and self.mpi_size > 1:
                # Distribute the scan among the different cients, independently
                vscan = vscan[self.mpi_rank::self.mpi_size]
                print("MPI #%2d analysing scans:" % self.mpi_rank, vscan)
            elif 'run' in self.params['mpi'] and self.mpi_master:
                self.print("Using MPI: %s" % self.params['mpi'])

        warn_liveplot = False
        if self.params['live_plot'] and self.params['nb_run'] > 1:
            warn_liveplot = True
        if MPI is not None:
            if self.params['live_plot'] and self.mpi_size > 1:
                warn_liveplot = True
        if warn_liveplot:
            self.print("\n", "#" * 100, "\n#",
                       "\n# WARNING: you are using the LIVE_PLOT option with multiple runs or MPI \n#"
                       "\n#      this will slow down the optimisation !\n"
                       "\n#      Please remember to enable live_plot only for tests \n"
                       + "#" * 100)

        for scan in vscan:
            if scan is not None:
                self.print("\n", "#" * 100, "\n#", "\n#  CDI Scan: ", scan, "\n#\n", "#" * 100)
            try:
                # We can alter data filename with scan number
                data_orig = self.params['data']
                # Get prefix to save CXI and output file
                file_prefix = None
                if 'specfile' in self.params:
                    if self.params['specfile'] is not None:
                        self.print(self.params['specfile'])
                        file_prefix = os.path.splitext(os.path.split(self.params['specfile'])[1])[0]
                if 'h5file' in self.params:
                    if self.params['h5file'] is not None:
                        self.print(self.params['h5file'])
                        file_prefix = os.path.splitext(os.path.split(self.params['h5file'])[1])[0]
                if file_prefix is None and 'data' in self.params:
                    if self.params['data'] is not None:
                        if '%' in self.params['data']:
                            # Inject scan number in data filename
                            try:
                                self.params['data'] = self.params['data'] % int(scan)
                            except:
                                try:
                                    self.params['data'] = self.params['data'] % scan
                                except:
                                    print("Failed to replace %d or %s in data string by: ", scan)
                        file_prefix = os.path.splitext(os.path.split(self.params['data'])[1])[0]

                if file_prefix is None:
                    file_prefix = 'data'

                if isinstance(scan, str):
                    try:
                        s = int(scan)
                        file_prefix = file_prefix + '_S%04d' % s
                    except ValueError:
                        file_prefix = file_prefix + '_S%s' % scan
                elif isinstance(scan, int):
                    file_prefix = file_prefix + '_S%04d' % scan

                self.ws = self.CDIRunnerScan(self.params, scan, timings=self.timings)
                self.ws.load_data()
                need_init_mask = True

                # data2cxi shall only be done by master for mpi=run jobs
                do_data2cxi = False
                if self.params['data2cxi']:
                    do_data2cxi = True
                    if MPI is not None:
                        if 'run' in self.params['mpi'] and not self.mpi_master:
                            do_data2cxi = False
                if do_data2cxi:
                    t1 = timeit.default_timer()
                    # This will save the original data, before cropping & centering
                    file_name = file_prefix + '.cxi'
                    if os.path.isfile(file_name):
                        self.print("Data CXI file already exists, not overwriting: ", file_name)
                    else:
                        self.ws.init_mask()
                        self.ws.corr_flat_field()
                        need_init_mask = False
                        self.print("Saving data to CXI file: ", file_name)
                        save_cdi_data_cxi(file_name, self.ws.iobs, wavelength=self.params['wavelength'],
                                          detector_distance=self.params['detector_distance'],
                                          pixel_size_detector=self.params['pixel_size_detector'],
                                          sample_name=self.params['sample_name'], mask=self.ws.mask,
                                          instrument=self.params['instrument'], note=self.params['note'],
                                          process_parameters=self.params)
                    self.timings['data2cxi'] = timeit.default_timer() - t1

                t1 = timeit.default_timer()
                self.ws.prepare_processing_unit()
                if default_processing_unit.backend == Backend.PYCUDA:
                    gpuname = default_processing_unit.cu_device.name()
                elif default_processing_unit.backend == Backend.PYOPENCL:
                    gpuname = default_processing_unit.cl_device.name
                else:
                    gpuname = "CPU?"
                self.params['processing_hostname'] = platform.node()
                self.params['processing_platform'] = platform.platform()
                self.params['processing_backend'] = default_processing_unit.backend.name.lower()
                self.params['processing_unit_name'] = gpuname
                s = ""
                if MPI is not None:
                    if self.mpi_size > 1:
                        s = "[MPI #%d/%d]" % (self.mpi_rank, self.mpi_size)
                print("Processing Unit init OK on %s [backend=%s, GPU=%s] %s" %
                      (platform.node(), default_processing_unit.backend.name.lower(), gpuname, s))
                # if we are using MPI, check the processing unit init went OK
                self.check_mpi(task="GPU init", timeout=60)
                self.timings['prepare_processing_unit'] = timeit.default_timer() - t1
                self.ws.prepare_data(init_mask=need_init_mask, corr_flat_field=need_init_mask)

                support_threshold = self.params['support_threshold']

                vrun = np.arange(1, self.params['nb_run'] + 1)

                if MPI is not None:
                    if 'run' in self.params['mpi'] and self.mpi_size > 1:
                        # Distribute the runs among the different cients, independently
                        vrun = vrun[self.mpi_rank::self.mpi_size]
                        print("MPI #%2d: performing %d runs" % (self.mpi_rank, len(vrun)))

                for run in vrun:
                    # KLUDGE: modifying the support threshold in params may not be the best way ?
                    st = ""
                    if isinstance(support_threshold, list) or isinstance(support_threshold, tuple):
                        if len(support_threshold) == 1:
                            self.params['support_threshold'] = support_threshold[0]
                        else:
                            self.params['support_threshold'] = np.random.uniform(support_threshold[0],
                                                                                 support_threshold[1])
                        st = "support_threshold = %5.3f" % self.params['support_threshold']
                    nbtry = 5
                    ok = False
                    while not ok and nbtry > 0:
                        nbtry -= 1
                        try:
                            self.print("\n", "#" * 100, "\n#", "\n#  CDI Run: %g/%g  %s\n#\n" %
                                       (run, self.params['nb_run'], st), "#" * 100)
                            s = time.strftime("-%Y-%m-%dT%H-%M-%S", time.localtime())
                            if self.params['nb_run'] > 1:
                                s += "_Run%04d" % run
                            self.ws.prepare_cdi()
                            self.ws.run(file_name=file_prefix + s, run_n=run)
                            ok = True
                        except SupportTooSmall as ex:
                            s = self.params['support_threshold_auto_tune_factor']
                            if nbtry:
                                s2 = "# ... trying to divide support threshold by %4.3f " \
                                     "(test %d/5)\n#\n" % (s, 5 - nbtry + 1)
                            else:
                                s2 = "# ... giving up\n#\n"
                            self.print("\n", "#" * 100, "\n#", "\n#  %s\n" % str(ex), s2, "#" * 100)
                            self.params['support_threshold'] /= s
                            st = "support_threshold = %5.3f" % self.params['support_threshold']
                        except SupportTooLarge as ex:
                            s = self.params['support_threshold_auto_tune_factor']
                            if nbtry:
                                s2 = "# ... trying to multiply support threshold by %4.3f " \
                                     "(test %d/5)\n#\n" % (s, 5 - nbtry + 1)
                            else:
                                s2 = "# ... giving up\n#\n"
                            self.print("\n", "#" * 100, "\n#", "\n#  %s\n" % str(ex), s2, "#" * 100)
                            self.params['support_threshold'] *= s
                            st = "support_threshold = %5.3f" % self.params['support_threshold']

                if self.ws.cdi is not None:
                    # Exotic error: if the number of tasks is larger than the number of runs,
                    # the cdi object may not have been initialised

                    # Free GPU memory
                    self.ws.cdi = FreePU() * self.ws.cdi

                if self.params['nb_run'] > 1 and self.params['nb_run_keep'] is not None:
                    # Keep only some of the best runs, delete others
                    if self.params['nb_run'] > self.params['nb_run_keep']:
                        res = self.ws.saved_results
                        if MPI is not None:
                            if 'run' in self.params['mpi']:
                                vres = self.mpic.gather(self.ws.saved_results, root=0)
                                if self.mpi_master:
                                    for i in range(1, self.mpi_size):
                                        res += vres[i]
                        if self.mpi_master or 'run' not in self.params['mpi']:
                            m = self.params['nb_run_keep_max_obj2_out']
                            if m > 0:
                                # Remove solutions with too large amplitude outside the support
                                res.sort(key=lambda x: x['obj2_out'])
                                if res[0]['obj2_out'] > m:
                                    raise CDIRunnerException("All solutions have too much amplitude outside "
                                                             "the support. All have been kept on disk "
                                                             "but either change nb_run_keep_max_obj2_out "
                                                             "or alter the optimisation parameters to fix this.")
                                if res[-1]['obj2_out'] > m:
                                    print("Removing solutions with too large amplitude outside the support"
                                          " (obj2_out > nb_run_keep_max_obj2_out=%4.3f):" % m)
                                while len(res) > 0 and res[-1]['obj2_out'] > m:
                                    v = res.pop()
                                    self.print("Removing: %s" % v['file_name'])
                                    fn = v['file_name']
                                    fnpng = fn.split("_LLKf")[0] + '.png'
                                    os.system('rm -f "%s" "%s"' % (fn, fnpng))
                            # Keep remaining solutions based on LLK_free
                            res.sort(key=lambda x: x['llk_poisson_free'])
                            if len(res) > self.params['nb_run_keep']:
                                print("Keeping %d solutions with the smallest Poisson free LLK"
                                      % self.params['nb_run_keep'])
                                for i in range(self.params['nb_run_keep'], len(res)):
                                    self.print("Removing: %s" % res[i]['file_name'])
                                    fn = res[i]['file_name']
                                    fnpng = fn.split("_LLKf")[0] + '.png'
                                    os.system('rm -f "%s" "%s"' % (fn, fnpng))
                                self.ws.saved_results = res[:self.params['nb_run_keep']]
                    else:
                        self.print('CDI runner: nb_run=%d <= nb_run_keep=%d !! Keeping all results' % (
                            self.params['nb_run'], self.params['nb_run_keep']))

                if self.params['nb_run'] > 1 and self.params['save'] not in ['none', None] and self.mpi_master:
                    # Move latest links to best result
                    res = self.ws.saved_results
                    res.sort(key=lambda x: x['llk_poisson_free'])
                    print("Moving 'latest.cxi|png' links to the best result")
                    os.system('rm -f latest.cxi latest.png')
                    sf = os.path.split(res[0]['file_name'])
                    os.system('ln -sf "%s" %s' % (sf[1], os.path.join(sf[0], 'latest.cxi')))
                    if self.params['save_plot']:
                        sf = os.path.split(res[0]['file_name'].split("_LLKf")[0] + '.png')
                        os.system('ln -sf "%s" %s' % (sf[1], os.path.join(sf[0], 'latest.png')))
                    self.params['data'] = data_orig
            except Exception as ex:
                if MPI is not None:
                    s = "[MPI #%d/%d]" % (self.mpi_rank, self.mpi_size)
                else:
                    s = ""
                print("ERROR occurred on: %s [%s] %s" % (platform.node(), platform.platform(), s))
                print('\n\n Caught exception for scan %s:\n   %s    \n'
                      'on %s [%s] %s\n' % (str(scan), str(ex), platform.node(), platform.platform(), s))
                print(traceback.format_exc())
                sys.exit(1)

        self.print("Timings:")
        for k, v in self.timings.items():
            if v > 1e-6:
                if MPI is not None:
                    self.print("MPI #%2d: %40s :%6.2fs" % (self.mpi_rank, k, v))
                else:
                    self.print("         %40s :%6.2fs" % (k, v))

        if warn_liveplot:
            self.print("\n", "#" * 100, "\n#",
                       "\n# WARNING: you are using the LIVE_PLOT option with multiple runs or MPI \n#"
                       "\n#      this will slow down the optimisation !\n"
                       "\n#      Please remember to enable live_plot only for tests \n"
                       + "#" * 100)
