# -*- 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

__all__ = ['HoloTomoData', 'HoloTomo', 'OperatorHoloTomo', 'algo_string', 'save_obj_probe_chunk',
           'save_nxtomo']

import logging
import time
import sys
import warnings
import ctypes
import numpy as np
from ..utils.h5py import *
from ..operator import Operator, OperatorException
from ..utils.history import History
from ..version import get_git_version
from multiprocessing import Array, sharedctypes

_pynx_version = get_git_version()


class HoloTomoDataStack:
    """
    Class to hold HoloTomo data arrays for a stack of projections in a GPU. Several are needed to achieve optimal
    throughput using parallel processing as well as data transfers to and from the GPU.
    Some objects may be left empty because they are not needed (e.g. probe which is the same for all projections, or
    psi which is only kept for some algorithms).
    This can also be used to store data in the host computer for optimised transfer using pinned/page-locked memory.
    """

    def __init__(self, istack=None, iproj=None, iobs=None, obj=None, psi=None, nb=None, obj_phase0=None,
                 pinned_memory=False):
        """

        :param istack: the index for this stack
        :param iproj: the index of the first projection in this stack
        :param iobs: observed intensity, with the center in (0,0) for FFT purposes. Shape: (stack_size, nb_z, ny, nx)
        :param obj: the object. Shape: (stack_size, ny, nx)
        :param psi: the current psi array, which has to be stored for some algorithms.
                    Shape: (stack_size, nb_z, nb_probe, ny, nx)
        :param nb: number of used projections in this stack, in case the total number of projections is not
                   a multiple of stack_size
        :param obj_phase0: the initial estimate for the phase, usually obtained with the Paganin operator. This
                           is used to determine the phase wrapping in ulterior estimates of the object, when
                           applying a delta/beta constraint. This has the same structure as obj, but is stored
                           with a float16|32 type. The phase should be positive, i.e. the object argument is -obj_phase0
        :param pinned_memory: True if arrays have been pre-allocated to pinned memory.
        """
        self.istack = istack
        self.iproj = iproj
        self.iobs = iobs
        self.obj = obj
        self.psi = psi
        self.nb = nb
        # True if arrays are currently using pinned memory
        self.pinned_memory = pinned_memory
        self.obj_phase0 = obj_phase0


class HoloTomoData:
    """
    Class to hold phase contrast imaging data collected at several propagation distances, and several projections.
    The input data will be reorganised to optimise memory transfers between GPU and host.
    """

    def __init__(self, iobs, iobs_empty, pixel_size_detector, wavelength, detector_distance, dx=None, dy=None,
                 scale_factor=None, stack_size=1, idx=None, padding=0, pu=None):
        """
        Init function.
        :param iobs: array of observed intensities (scaled to the same pixel size and aligned) with 4 dimensions:
                     - nb_proj different projections along the first axis (dimension can be =1).
                     - the nb_z different distances along the second axis (dimension can be =1)
                     - the 2D images along the last two axis. These may be cropped to be adapted to suitable sizes
                       for GPU FFT.
                     - values < 0 are masked
                     - the input iobs array should be centered on the 2D arrays, as measured
                     Notes:
                     - the iobs data should already be padded, since different propagation distances will have
                        different padded regions.
                     - all masked/padded intensities should have been interpolated and stored
                       as -1-I_interp, to ensure strictly negative values. These values will only
                       be used for Paganin/CTF calculations, and will be considered free otherwise.
        :param iobs_empty: same as iobs, but only with the images measured without a sample (direct beam),
                    with the same array structure. The shape should also be (nb_iobs_empty, nz, ny, nx) in
                    case several images (e.g. begin/end of measurement) are available.
        :param pixel_size_detector: the detector pixel size in meters, which must be the same for all images.
        :param wavelength: the experiment wavelength in meters.
        :param detector_distance: the detector distance as an array or list of nb_z distance in meters
        :param dx, dy: shifts of the object with respect to the illumination, in pixels. The shape of the shifts must
                       be (nb_proj, nb_z).
        :param scale_factor: float array with a scale factor for each image. Shape (nb_proj, nb_dist)
        :param stack_size: data will be processed by stacks including one or several projections, and all
                           propagation distances. This is done to manage the GPU memory footprint of the algorithms.
                           stack_size is the number of projections per stack (default 1).
        :param idx: the vector with the indices of the projections analysed. If None, range(nb_proj) will be used
            A negative value can be used for empty beam images, but they will be ignored.
        :param padding: the number of pixels used for padding on each border. This value is only used for
            the final export, to crop padded areas. This can either be a single value or a tuple (pady, padx).
        :param pu: the ProcessingUnit to use to directly allocate pinned (page-locked) memory. If None,
            the arrays will be reallocated later.
        """
        self.nproj, self.nz, self.ny, self.nx = iobs.shape
        if idx is None:
            self.idx = np.arange(self.nproj, dtype=np.int16)
        else:
            self.idx = np.array(idx).astype(np.int16)
        self.stack_size = stack_size
        self.pixel_size_detector = pixel_size_detector
        # TODO: take into account non-monochromaticity ?
        self.wavelength = wavelength
        self.detector_distance = np.array(detector_distance, dtype=np.float32)

        if dx is None or dy is None:
            self.dx = np.zeros((self.nproj, self.nz), dtype=np.float32)
            self.dy = np.zeros((self.nproj, self.nz), dtype=np.float32)
        else:
            assert dx.shape == dy.shape
            self.dx, self.dy = dx.astype(np.float32), dy.astype(np.float32)

        if scale_factor is None:
            self.scale_factor = np.ones(self.nproj, dtype=np.float32)
        else:
            self.scale_factor = scale_factor.astype(np.float32)

        # Padding used on each border, for final cropping
        if np.isscalar(padding):
            self.padding = (padding, padding)
        else:
            assert len(padding) == 2
            self.padding = padding

        # Reorganize data in a list of data stack, which can be later optimized for faster memory transfers
        self.stack_v = []
        n = stack_size
        for i in range(0, self.nproj, n):
            # Numpy is smart enough to only take the available data, even if i+n > len(iobs)
            tmp = iobs[i:i + n].astype(np.float32)
            if pu is not None:
                self.stack_v.append(HoloTomoDataStack(istack=i // n, iproj=i, iobs=pu.to_pagelocked(tmp),
                                                      nb=len(tmp), pinned_memory=True))
            else:
                self.stack_v.append(HoloTomoDataStack(istack=i // n, iproj=i, iobs=tmp, nb=len(tmp)))

        # Keep iobs_empty as a separate stack
        self.iobs_empty = iobs_empty.astype(np.float32)
        if pu is not None:
            self.iobs_empty = pu.to_pagelocked(self.iobs_empty)
        # Iobs_empty means - used to normalise iobs images by iobs_empty/<iobs_empty>
        self.iobs_empty_mean = np.array([d[d >= 0].mean() for d in self.iobs_empty])

        self.nb_obs = iobs.size

    def replace_data(self, iobs, iobs_empty, dx=None, dy=None, idx=None, init_obj=False, pu=None):
        """
        Replace the existing Iobs data by a new one, re-using existing allocated pinned
        data arrays to save time. All other parameters including stack size remain the same.
        :param iobs: array of observed intensities (scaled to the same pixel size and aligned) with 4 dimensions:
                     - nb_proj different projections along the first axis (dimension can be =1).
                     - the nb_z different distances along the second axis (dimension can be =1)
                     - the 2D images along the last two axis. These may be cropped to be adapted to suitable sizes
                       for GPU FFT.
                     - values < 0 are masked (and stored as -1e38, or -1-I_interp if interpolated)
                     - the input iobs array should be centered on the 2D arrays, as measured
                     Notes:
                     - the iobs data should already be padded, since different propagation distances will have
                        different padded regions. Those should be masked using negative values (-1e38).
        :param iobs_empty: same as iobs, but only with the images measured without a sample (direct beam),
                    with the same array structure. The shape should also be (nb_iobs_empty, nz, ny, nx) in
                    case several images (e.g. begin/end of measurement) are available.
        :param dx, dy: shifts of the object with respect to the illumination, in pixels. The shape of the shifts must
                       be (nb_proj, nb_z).
        :param idx: the vector with the indices of the projections analysed. If None, range(nb_proj) will be used
            A negative value can be used for empty beam images, but they will be ignored.
        :param init_obj: if True, the object and its phase will be initialised to 1 and 1000.
        :param pu: the ProcessingUnit to use to directly allocate pinned (page-locked) memory. If None,
            the arrays will be reallocated later.
        """
        self.nproj, self.nz, self.ny, self.nx = iobs.shape
        if idx is None:
            self.idx = np.arange(self.nproj, dtype=np.int16)
        else:
            self.idx = np.array(idx).astype(np.int16)

        if dx is None or dy is None:
            self.dx = np.zeros((self.nproj, self.nz), dtype=np.float32)
            self.dy = np.zeros((self.nproj, self.nz), dtype=np.float32)
        else:
            assert dx.shape == dy.shape
            self.dx, self.dy = dx.astype(np.float32), dy.astype(np.float32)

        # Reorganize data in a list of data stack, which can be later optimized for faster memory transfers
        n = self.stack_size
        if len(self.stack_v) > np.ceil(self.nproj / n):
            self.stack_v = self.stack_v[:int(np.ceil(self.nproj / n))]
        istack = 0
        for i in range(0, self.nproj, n):
            # Numpy is smart enough to only take the available data, even if i+n > len(iobs)
            tmp = iobs[i:i + n]
            if len(self.stack_v) > istack:
                s = self.stack_v[istack]
                if len(tmp) == s.nb:
                    s.iobs[:] = tmp
                else:
                    s.iobs = pu.to_pagelocked(tmp.astype(np.float32))
                    # Also update object
                    s.obj = pu.pagelocked_empty((self.stack_size, self.ny, self.nx), np.complex64)
                    s.obj_phase0 = pu.pagelocked_empty((self.stack_size, self.ny, self.nx), np.float32)
                    s.psi = None
                s.nb = len(tmp)
            else:
                if pu is not None:
                    self.stack_v.append(HoloTomoDataStack(istack=istack, iproj=i,
                                                          iobs=pu.to_pagelocked(tmp.astype(np.float32)),
                                                          nb=len(tmp), pinned_memory=True))
                    s = self.stack_v[-1]
                    s.obj = pu.pagelocked_empty((self.stack_size, self.ny, self.nx), np.complex64)
                    s.obj_phase0 = pu.pagelocked_empty((self.stack_size, self.ny, self.nx), np.float32)
                else:
                    self.stack_v.append(HoloTomoDataStack(istack=istack, iproj=i, iobs=tmp.astype(np.float32),
                                                          nb=len(tmp)))
                    s = self.stack_v[-1]
                    s.obj = np.ones((self.stack_size, self.ny, self.nx), np.complex64)
                    s.obj_phase0 = np.ones((self.stack_size, self.ny, self.nx), np.float32)
            if init_obj:
                s.obj[:] = 1
                s.obj_phase0[:] = 1000
            istack += 1

        # Keep iobs_empty as a separate stack
        self.iobs_empty[:] = iobs_empty.astype(np.float32)

        # Iobs_empty means - used to normalise iobs images by iobs_empty/<iobs_empty>
        self.iobs_empty_mean = np.array([d[d >= 0].mean() for d in self.iobs_empty])
        self.nb_obs = iobs.size


class HoloTomo:
    """ Reconstruction class for two-dimensional phase coherent imaging data, using multiple propagation
    distances and multiple projections, i.e. holo-tomography.
    This class is designed to handle large datasets which cannot fit on a GPU. Complex operations will loop
    over projections by continuously swapping data between host and GPU.
    """

    def __init__(self, data: HoloTomoData, obj: np.array = None, probe: np.array = None, coherent_probe_modes=False,
                 pu=None):
        """

        :param data: a HoloTomoData object with observed data
        :param obj: an object which will be fitted to the data, which should have 3 dimensions:
                    - projections (considered independent) along the first axis
                    - xy along the last two dimensions, with the same size as the data.
                    The data center should be the center of the array.
        :param probe: the illumination corresponding to the data, which should have 4 dimensions:
                    - one illumination per detector distance along the first axis, or just 1 if it is the same
                    - illumination/probe modes along the first axis (>=1)
                    - illumination 2D data along the last two dimensions
                    Since images can be taken with a shifted illumination, the 2D probe size should be larger than
                    the observed intensity and object arrays. This will be automatically corrected, and the
                    dx and dy shifts will also be re-centered so that the middle dx and dy
                    values are equal to zero.
                    The data center should be the center of the array.
        :param coherent_probe_modes: probe modes are applied coherently, similarly to the
            "orthogonal probe relaxation" method: for each projection, only one coherent illumination
            which is a linear combination of probe modes is used. Note that the modes are
            not generally orthonormal.
            An array of shape (nproj, nz, nprobe) of coefficients should be supplied
            to give the coefficients of each mode.
        :param pu: the processing_unit to use to directly allocate arrays to pinned memory.
            If None, the arrays will be re-allocated later.
        """
        # TODO: add incoherent background, flat field ?
        self.data = data

        # Empty objects to hold GPU data. This will be HoloTomoDataStack objects when used
        # The 'in' and 'out' stacks are used for memory transfers in // to calculations with the main stack.
        # The probe GPU array (identical for all projections) is kept in the active stack
        self._cl_stack, self._cl_stack_in, self._cl_stack_out = None, None, None
        self._cu_stack, self._cu_stack_in, self._cu_stack_out = None, None, None

        # The timestamp counters record when the holotomo data was last altered, either in the host or GPU memory.
        self._timestamp_counter = 1
        self._cl_timestamp_counter = 0
        self._cu_timestamp_counter = 0

        # Holds the Psi array for the current stack of projection(s)
        self._psi = None

        # Probe array
        self._probe = None

        # Default number of coherent probe modes
        self.nb_probe = 1
        # Coherent probe modes coefficients. This holds
        # a 3-dimensional array (nb_proj, nb_z, nb_mode) holding the coefficients
        # of the linear combination of probe modes for each projection and distance
        self.probe_mode_coeff = None

        self._init_obj(obj, pu=pu)
        self._init_probe(probe, coherent_probe_modes, pu=pu)
        self._init_psi(pu=pu)

        self.llk_poisson = 0
        self.llk_gaussian = 0
        self.llk_euclidian = 0
        self.nb_photons_calc = 0

        # Experimental - not compatible with coherent probe modes
        self.constrain_probe_direct_beam = False

        # Experimental - gaussian fwhm (pixels) for the detector psf
        self._psf = None

        # Record the number of cycles (ML, AP, DM, etc...), for history purposes
        self.cycle = 0
        # History record
        self.history = History()

    def _init_obj(self, obj, pu=None):
        """
        Init the object array by storing it in the HoloTomoData stack. We do not keep a complete 3D array of the object
        in a single array. The shape of each stack is (stack_size, nyo, nxo).

        :return:
        """
        # TODO: use calc_obj_shape to take into account lateral shifts and expand object size
        # For now, assume object and probe have the same size
        n = self.data.stack_size
        if obj is None:
            s = n, self.data.ny, self.data.nx
            for i in range(len(self.data.stack_v)):
                if pu is None:
                    self.data.stack_v[i].obj = np.ones(s, dtype=np.complex64)
                else:
                    self.data.stack_v[i].obj = pu.pagelocked_empty(s, np.complex64)
                    self.data.stack_v[i].obj[:] = 1
        else:
            assert obj.shape[-2] == self.data.ny
            assert obj.shape[-1] == self.data.nx
            for i in range(len(self.data.stack_v)):
                if (i + 1) * n <= len(obj):
                    if pu is None:
                        self.data.stack_v[i].obj = obj[i * n:(i + 1) * n]
                    else:
                        self.data.stack_v[i].obj = pu.to_pagelocked(obj[i * n:(i + 1) * n])
                else:
                    dn = len(obj) - i * n
                    if pu is None:
                        self.data.stack_v[i].obj[:dn] = obj[i * n:]
                    else:
                        self.data.stack_v[i].obj[:dn] = pu.to_pagelocked(obj[i * n:])
        # Init obj_phase0 to dummy values
        for i in range(len(self.data.stack_v)):
            s = n, self.data.ny, self.data.nx
            # self.data.stack_v[i].obj_phase0 = -np.ones(s, dtype=np.float16) * 1000
            if pu is None:
                self.data.stack_v[i].obj_phase0 = np.ones(s, dtype=np.float32) * 1000
            else:
                self.data.stack_v[i].obj_phase0 = pu.pagelocked_empty(s, np.float32)
                self.data.stack_v[i].obj_phase0[:] = 1000
        assert self.data.stack_v[0].obj.shape == (n, self.data.ny, self.data.nx)

    def _init_probe(self, probe, coherent_probe_modes, pu=None):
        """
        Init & resize (if necessary) the probe according to the minimum size calculated for the HoloTomoData.

        :return: nothing.
        """
        if probe is None:
            self.nb_probe = 1
            self._probe = np.ones((self.data.nz, 1, self.data.ny,
                                   self.data.nx), dtype=np.complex64)
        else:
            self.nb_probe = probe.shape[1]
            self._probe = probe.astype(np.complex64)
        assert self._probe.shape == (self.data.nz, self.nb_probe, self.data.ny, self.data.nx)
        if isinstance(coherent_probe_modes, np.ndarray):
            assert coherent_probe_modes.shape == (self.data.nproj, self.data.nz, self.nb_probe)
            self.probe_mode_coeff = coherent_probe_modes.astype(np.float32)
        elif coherent_probe_modes:
            self.probe_mode_coeff = np.ones((self.data.nproj, self.data.nz, self.nb_probe), dtype=np.float32)
        else:
            self.probe_mode_coeff = None
        if pu is not None:
            self._probe = pu.to_pagelocked(self._probe)

    def _init_psi(self, pu=None):
        """
        Init the psi array
        :return: nothing.
        """
        if pu is None:
            self._psi = np.empty((self.data.stack_size, self.data.nz,
                                  self.data.ny, self.data.nx), dtype=np.complex64)
        else:
            self._psi = pu.pagelocked_empty((self.data.stack_size, self.data.nz,
                                             self.data.ny, self.data.nx), dtype=np.complex64)

    def get_x_y(self):
        """
        Get 1D arrays of x and y coordinates, taking into account the pixel size.
        x is an horizontal vector and y vertical.

        :return: a tuple (x, y) of 1D numpy arrays
        """
        ny, nx = self.data.ny, self.data.nx
        px = self.data.pixel_size_detector
        x, y = np.arange(-nx // 2, nx // 2, dtype=np.float32), \
            np.arange(-ny // 2, ny // 2, dtype=np.float32)[:, np.newaxis]
        return x * px, y * px

    def set_probe(self, probe, probe_mode_coefficients=None):
        """
        Give a new value for the probe (illumination) array.

        :param probe: the 3D array to be used
        :param probe_mode_coefficients: if None, coherent probe modes will not be used.
            Otherwise an array of shape (nproj, nz, nb_probe) can be supplied.
        :return: nothing
        """
        self._from_pu()
        self._init_probe(probe, probe_mode_coefficients)
        self._timestamp_counter += 1

    def get_probe(self):
        """
        Get the probe
        :return: the complex array of the probe, with 3 dimensions (1 for the modes)
        """
        self._from_pu()
        return self._probe

    def set_obj(self, obj):
        """
        Give a new value for the object projections. The supplied array shape should be: (nb_proj, ny, nx)

        :param obj: the 3D array to be used
        :return: nothing
        """
        self._from_pu()
        nz, ny, nx = obj.shape
        assert nx == self.data.nx and ny == self.data.ny and nz == self.data.nproj
        for s in self.data.stack_v:
            # We use s.obj[:] because the array may be using pinned memory
            s.obj[:] = obj[s.iproj:s.iproj + s.nb].reshape((s.nb, ny, nx)).astype(np.complex64)
        self._timestamp_counter += 1

    def set_positions(self, dx, dy):
        """
        Set the positions fot all projections and distances
        :param dx, dy: shifts of the object projections with respect to the illumination, in pixels.
            The shape of the shifts must be (nb_proj, nb_z).
        :return:
        """
        self._from_pu()
        self.data.dx = dx.astype(np.float32)
        self.data.dy = dy.astype(np.float32)
        self._timestamp_counter += 1

    def _from_pu(self, psi=False):
        """
        Get object, probe, and optionally psi array from opencl or CUDA memory to host memory for the current stack,
        only if they are more recent than the arrays in the host memory.

        Access to psi is only for development purposes

        :return: nothing
        """
        # TODO: should this rather use an imported FromPU operator ? Would be cleaner, but would importing it..

        # Careful when copying back data: arrays may already be allocated with pinned memory, so we must not
        # change the location of the data - hence the [:] copy
        if self._timestamp_counter < self._cl_timestamp_counter and self._cl_stack is not None:
            from .cl_operator import default_processing_unit
            default_processing_unit.finish()
            self.data.stack_v[self._cl_stack.istack].obj[:] = self._cl_stack.obj.get()
            self.data.stack_v[self._cl_stack.istack].obj_phase0[:] = self._cl_stack.obj_phase0.get()
            self._probe = self._cl_stack.probe.get()
            self._timestamp_counter = self._cl_timestamp_counter

        if self._timestamp_counter < self._cu_timestamp_counter and self._cu_stack is not None:
            from .cu_operator import default_processing_unit
            default_processing_unit.finish()
            self.data.stack_v[self._cu_stack.istack].obj[:] = self._cu_stack.obj.get()
            self.data.stack_v[self._cu_stack.istack].obj_phase0[:] = self._cu_stack.obj_phase0.get()
            self._probe[:] = self._cu_probe.get()
            if self.probe_mode_coeff is not None:
                self.probe_mode_coeff = self._cu_probe_mode_coeff.get()
            self._timestamp_counter = self._cu_timestamp_counter

        if psi:
            if self._cl_stack is not None:
                if self._cl_stack.psi is not None:
                    self._psi = self._cl_stack.psi.get()
            if self._cu_stack is not None:
                if self._cu_stack.psi is not None:
                    self._psi = self._cu_stack.psi.get()

    def __rmul__(self, x):
        """
        Multiply object (by a scalar).

        This is a placeholder for a function which will be replaced when importing either CUDA or OpenCL operators.
        If called before being replaced, will raise an error

        :param x: the scalar by which the wavefront will be multiplied
        :return:
        """
        if np.isscalar(x):
            raise OperatorException(
                "ERROR: attempted Op1 * Op2, with Op1=%s, Op2=%s. Did you import operators ?" % (str(x), str(self)))
        else:
            raise OperatorException("ERROR: attempted Op1 * Op2, with Op1=%s, Op2=%s." % (str(x), str(self)))

    def __mul__(self, x):
        """
        Multiply object (by a scalar).

        This is a placeholder for a function which will be replaced when importing either CUDA or OpenCL operators.
        If called before being replaced, will raise an error

        :param x: the scalar by which the wavefront will be multiplied
        :return:
        """
        if np.isscalar(x):
            raise OperatorException(
                "ERROR: attempted Op1 * Op2, with Op1=%s, Op2=%s. Did you import operators ?" % (str(self), str(x)))
        else:
            raise OperatorException("ERROR: attempted Op1 * Op2, with Op1=%s, Op2=%s." % (str(self), str(x)))

    def __str__(self):
        return "HoloTomo"

    def reset_history(self):
        """
        Reset history, and set current cycle to zero
        :return: nothing
        """
        self.history = History()
        self.cycle = 0

    def update_history(self, mode='llk', update_obj=False, update_probe=False, update_backgroung=False,
                       update_pos=False, verbose=False, logger=None, **kwargs):
        """
        Update the history record.
        :param mode: either 'llk' (will record new log-likelihood and number of photons)
                     or 'algorithm' (will only update the algorithm) - for the latter case, algorithm
                     should be given as a keyword argument
        :param verbose: if True, print some info about current process (only if mode=='llk')
        :param kwargs: other parameters to be recorded, e.g. probe_inertia=xx, dt=xx, algorithm='DM'
        :return: nothing
        """
        if mode == 'llk':
            if logger is None:
                logger = logging.getLogger("HoloTomo")
                logger.setLevel(logging.INFO)
            algo = ''
            dt = 0
            if 'algorithm' in kwargs:
                algo = kwargs['algorithm']
            if 'dt' in kwargs:
                dt = kwargs['dt']
            if verbose:
                s = algo_string(algo, self, update_obj, update_probe, update_backgroung, update_pos)
                logger.info("%-10s #%3d LLK= %8.3f(p) %8.3f(g) %8.3f(e), nb photons=%e, dt/cycle=%5.3fs"
                            % (s, self.cycle, self.llk_poisson / self.data.nb_obs, self.llk_gaussian / self.data.nb_obs,
                               self.llk_euclidian / self.data.nb_obs, self.nb_photons_calc, dt))

            self.history.insert(self.cycle, llk_poisson=self.llk_poisson / self.data.nb_obs,
                                llk_gaussian=self.llk_gaussian / self.data.nb_obs,
                                llk_euclidian=self.llk_euclidian / self.data.nb_obs,
                                nb_photons_calc=self.nb_photons_calc, nb_probe=self.nb_probe,
                                **kwargs)
        elif 'algo' in mode:
            if 'algorithm' in kwargs:
                self.history.insert(self.cycle, algorithm=kwargs['algorithm'])

    def save_obj_probe_chunk(self, chunk_prefix, save_obj_phase=True, save_obj_complex=False,
                             save_probe=True, dtype=np.float16, verbose=False, crop_padding=True,
                             process_notes=None, process_parameters=None, nxtomo=True,
                             sample_name=None, angles=None, append=True):
        """
        Save the chunk (the projections included in this object) in an hdf5 file
        :param chunk_prefix: the prefix, e.g. "my_sample_%04d" for the filename the data will be saved to.
            The %04d field will be replaced by the first projection. A '.h5' prefix will be appended
            if no "%" is included in the prefix, only '.h5' will be appended to the filename.
            If nxtomo=True, this should just be the filename e.g. 'dataset.nx'.
        :param save_obj_phase: if True, save the object phase. Mandatory for nxtomo=True
        :param save_obj_complex: if True, also save the different complex projections of the object.
            Not compatible with nxtomo.
        :param save_probe: if True, save the complex probe (illumination)
        :param dtype: the floating point dtype to save the phase
        :param verbose: if True, will print some information.
        :param crop_padding: if True (the default), the padding area is not saved
        :param process_notes: a dictionary which will be saved in '/entry_N/process_1'. A dictionary entry
        can also be a 'note' as keyword and a dictionary as value - all key/values will then be saved
        as separate notes. Example: process={'program': 'PyNX', 'note':{'llk':1.056, 'nb_photons': 1e8}}
        :param process_parameters: a dictionary of parameters which will be saved as a NXcollection
        :param nxtomo: if True, this will be saved with the nxtomo format. This also requires
            giving an angular step.
        :param sample_name: string for the sample description
        :param angles: angles (in radians) for the projections. Required for nxtomo=True.
        :param append: for nxtomo only, if True will append the result if the file already exists
        :return:
        """
        pady, padx = self.data.padding
        if padx == 0 and pady == 0:
            crop_padding = False
        if nxtomo:
            fname = chunk_prefix
            if angles is None:
                raise RuntimeError("save_obj_probe_chunk: nxtomo=True requires the list of angles "
                                   "for the projections")
            else:
                if len(angles) != self.data.nproj:
                    raise RuntimeError("save_obj_probe_chunk: the number of angles "
                                       "does not match the number of projections")
        else:
            fname = chunk_prefix + ".h5"
            if '%' in fname:
                fname = fname % self.data.idx[0]
        if verbose:
            print("Saving holotomo projections chunk & illumination to: %s" % fname)
        idx = None
        if save_obj_phase:
            idx, obj_phase = self.get_obj_phase_unwrapped(crop_padding=crop_padding, dtype=dtype)
        else:
            obj_phase = None

        if save_obj_complex:
            # Assemble object phase
            nproj = self.data.nproj
            dtype_cplx = np.complex64
            ny, nx = self.data.ny, self.data.nx
            obj = np.empty((nproj, ny, nx), dtype=dtype_cplx)
            i = 0
            idx = []
            for s in self.data.stack_v:
                for ii in range(s.nb):
                    idx.append(self.data.idx[s.iproj + ii])
                    ob = s.obj[ii]
                    if crop_padding:
                        if padx:
                            ob = ob[..., padx:-padx]
                        if pady:
                            ob = ob[:, pady:-pady, ]
                    obj[i] = ob
                    i += 1
            obj_complex = ob
        else:
            obj_complex = None

        if save_probe:
            probe = self.get_probe()
            if crop_padding:
                if padx:
                    probe = probe[..., padx:-padx]
                if pady:
                    probe = probe[:, :, pady:-pady]
            probe_mode_coeff = self.probe_mode_coeff
        else:
            probe = None
            probe_mode_coeff = None

        if nxtomo:
            save_nxtomo(fname, pixel_size=self.data.pixel_size_detector, obj_phase=obj_phase,
                        sample_name=sample_name, chunk_idx=self.data.idx[0], angles=angles,
                        probe=self.get_probe(), process_notes=process_notes,
                        process_parameters=process_parameters,
                        detector_distance=self.data.detector_distance,
                        wavelength=self.data.wavelength, append=append)
        else:
            save_obj_probe_chunk(fname, self.data.idx, self.data.pixel_size_detector, obj_phase=obj_phase,
                                 obj_complex=obj_complex,
                                 probe=probe, probe_mode_coeff=probe_mode_coeff, process_notes=process_notes,
                                 process_parameters=process_parameters)

    def get_obj_phase_unwrapped(self, crop_padding=True, dtype=np.float16, idx=None, sino_filter=None):
        """
        Get an array of the object phase, unwrapped based on the initial phase from Paganin/CTF
        :param crop_padding: if True (the default), the padded area is not saved
        :param dtype: the numpy dtype to use for the phase (defaults to float16)
        :param idx: if None, all the projections are returned. If idx is a number or
            a list/array of projections, only those are returned. These are
            the projections index, as given to the HoloTomoData object, i.e.
            HoloTomoData.idx .
        :param sino_filter: a string with the name of a 1D sinogram filter to be applied to
            the phases. This is used to perform the pre-filtering of the filtered back-projection
            (FBP), before cropping the padded areas. Available filters (from silx): Ram-Lak,
            Shepp-Logan, Cosine, Hamming, Hann, Tukey, Lanczos. Case-insensitive.
            If None (the default), no filtering is done.

        :return: idx, ph where idx is the index of the analysed frames, and ph, a 3D array
            with the phases for each projection
        """
        self._from_pu()
        # TODO: we have to import here to avoid early GPU (pycuda/pyopencl) import (?)
        from .utils import sino_filter_pad
        pady, padx = self.data.padding
        if padx == pady == 0:
            crop_padding = False
        ny, nx = self.data.stack_v[0].obj.shape[-2:]
        if isinstance(idx, int) or isinstance(idx, np.integer):
            idx = [idx]
        elif idx is None:
            idx = self.data.idx
        nproj = len(idx)
        if crop_padding:
            ny, nx = ny - 2 * pady, nx - 2 * padx
        obj_phase = np.empty((nproj, ny, nx), dtype=dtype)
        i = 0
        vidx = []
        # TODO: parallelise phase extraction
        for s in self.data.stack_v:
            for ii in range(s.nb):
                idxtmp = self.data.idx[s.iproj + ii]
                if idxtmp not in idx:
                    continue
                vidx.append(idxtmp)
                if s.obj_phase0 is not None:
                    # This should have been updated to the most recent unwrapped phase
                    op = s.obj_phase0[ii]
                else:
                    # Can this happen ? We should always have obj_phase0
                    op = np.angle(s.obj[ii])
                if crop_padding:
                    op = op[pady:-pady]
                if sino_filter is not None:
                    # NB: this is done more efficiently by the GPU SinoFilter operator
                    op = sino_filter_pad(op, sino_filter, padding=padx)
                if crop_padding:
                    op = op[:, padx:-padx]
                obj_phase[i] = op
                i += 1
        return np.array(vidx, dtype=np.int32), obj_phase


class OperatorHoloTomo(Operator):
    """
    Base class for an operator on CDI objects, not requiring a processing unit.
    """

    def timestamp_increment(self, pci: HoloTomo):
        pci._timestamp_counter += 1


def algo_string(algo_base, p: HoloTomo, update_object, update_probe, update_background=False, update_pos=False):
    """
    Get a short string for the algorithm being run, e.g. 'DM/o/3p' for difference map with 1 object and 3 probe modes.

    :param algo_base: 'AP' or 'ML' or 'DM'
    :param p: the holotomo object
    :param update_object: True if updating the object
    :param update_probe: True if updating the probe
    :param update_background: True if updating the background
    :param update_pos: True if updating the positions
    :return: a short string for the algorithm
    """
    s = algo_base

    if update_object:
        s += "/"
        s += "o"

    if update_probe:
        s += "/"
        if p.nb_probe > 1:
            s += "%d" % p.nb_probe
        s += "p"

    if update_background:
        s += "/b"

    if update_pos:
        s += "/t"

    return s


def save_nxtomo(filename, pixel_size, obj_phase, sample_name: str, chunk_idx: int, angles, probe=None,
                process_notes=None, process_parameters=None, obj_phase_shape=None,
                x_rotation_axis=None, y_rotation_axis=None, detector_distance=None,
                wavelength=None, append=True):
    """
    Save the projection phases to a NeXus fil with the nxtomo format.
    :param filename: name of the hdf5 file to save to.
    :param pixel_size: the object pixel size
    :param obj_phase: the object phase to save, either as a numpy array of a multiprocessing.Array. For
        the latter, the shape must be supplied in obj_phase_shape
    :param sample_name: the name or description of the sample
    :param chunk_idx: the index of the chunk to save, in practice the number of the first projection
        included in this chunk.
    :param angles: the list/array of angles for each object projection in this chunk
    :param probe: the complex probe (illumination) to save, or None. Addition to nxtomo format
    :param process_notes: a dictionary which will be saved in '/entry_N/process_1'. A dictionary entry
    can also be a 'note' as keyword and a dictionary as value - all key/values will then be saved
    as separate notes. Example: process={'program': 'PyNX', 'note':{'llk':1.056, 'nb_photons': 1e8}}
    :param process_parameters: a dictionary of parameters which will be saved as a NXcollection
    :param obj_phase_shape: the ph array shape if using multiprocessing.Array
    :param x_rotation_axis: optional axis position in pixel (if vertical axis)
    :param y_rotation_axis: optional axis position in pixel (if horizontal axis)
    :param detector_distance: in meters
    :param wavelength: X-ray beam wavelength in meters
    :param append: for nxtomo only, if True will append the result if the file already exists
    :return: nothing
    """
    if isinstance(obj_phase, sharedctypes.SynchronizedArray):
        # TODO: handle float32 and float16 cases
        obj_phase = np.frombuffer(obj_phase.get_obj(), dtype=np.float16).reshape(obj_phase_shape)

    if append and os.path.exists(filename):
        with h5py.File(filename, 'a') as h:
            # Just update the phased projection arrays
            # Chunk data
            detector = h['/entry_1/instrument/detector']
            nproj, ny, nx = obj_phase.shape
            chunks = min(32, nproj), min(4, ny), nx
            ds = detector.create_dataset(f"_chunk{chunk_idx}", data=obj_phase.astype(np.float32), chunks=chunks,
                                         shuffle=True, **hdf5plugin.LZ4())

            # Virtual dataset
            sh = list(detector['data'].shape)
            sh[0] += len(obj_phase)
            del detector['data']
            layout = h5py.VirtualLayout(shape=tuple(sh), dtype=np.float32)
            n0 = 0
            for i in range(chunk_idx + 1):  # We assume chunks arrive in order
                dsi = detector[f'_chunk{i}']
                n1 = n0 + len(dsi)
                layout[n0:n1] = h5py.VirtualSource(".", dsi.name, dsi.shape, dsi.dtype)
                n0 = n1
            detector.create_virtual_dataset("data", layout)
            del detector['image_key']
            detector.create_dataset("image_key", data=np.zeros(len(detector['data']), dtype=np.int8))

            sample = h['/entry_1/sample']
            rot = sample['rotation_angle'][()]
            del sample['rotation_angle']
            sample.create_dataset("rotation_angle", data=np.append(rot, angles, axis=0))
    else:
        with h5py.File(filename, 'w') as h:
            entry = h.create_group("/entry_1")
            h.attrs['default'] = 'entry_1'
            h.attrs['creator'] = 'PyNX'
            # f.attrs['NeXus_version'] = '2018.5'  # Should only be used when the NeXus API has written the file
            h.attrs['HDF5_Version'] = h5py.version.hdf5_version
            h.attrs['h5py_version'] = h5py.version.version
            entry.attrs['NX_class'] = 'NXentry'
            entry.create_dataset("program_name", data="PyNX %s" % _pynx_version)
            entry.create_dataset("definition", data="NXtomo")
            entry.attrs["default"] = "data"

            instrument = entry.create_group("instrument")
            instrument.attrs["NX_class"] = "NXinstrument"

            beam = instrument.create_group("beam")
            beam.attrs["NX_class"] = "NXbeam"
            if wavelength is not None:
                beam.create_dataset("incident_energy", data=12.3984e-10 / wavelength)
                beam["incident_energy"].attrs['units'] = 'keV'
            detector = instrument.create_group("detector")
            detector.attrs["NX_class"] = "NXdetector"

            px = pixel_size
            ny, nx = obj_phase.shape[-2:]
            unit_scale = np.log10(max(nx * px, ny * px))
            if unit_scale < -6:
                unit_name = "nm"
                unit_scale = 1e9
            elif unit_scale < -3:
                unit_name = u"µm"
                unit_scale = 1e6
            elif unit_scale < 0:
                unit_name = "mm"
                unit_scale = 1e3
            else:
                unit_name = "m"
                unit_scale = 1

            # Pixel size & distance
            detector.create_dataset("x_pixel_size", data=px)
            detector.create_dataset("y_pixel_size", data=px)
            if detector_distance is not None:
                detector.create_dataset("distance", data=detector_distance)

            # Chunk data
            chunks = min(32, len(obj_phase)), min(4, ny), nx
            ds = detector.create_dataset(f"_chunk{chunk_idx}", data=obj_phase.astype(np.float32), chunks=chunks,
                                         shuffle=True, **hdf5plugin.LZ4())
            ds.attrs['interpretation'] = 'image'

            # Virtual dataset
            layout = h5py.VirtualLayout(shape=obj_phase.shape, dtype=np.float32)
            layout[:] = h5py.VirtualSource(".", ds.name, ds.shape, ds.dtype)
            detector.create_virtual_dataset("data", layout)

            detector.create_dataset("image_key", data=np.zeros(len(obj_phase), dtype=np.int8))

            sample = entry.create_group("sample")
            sample.attrs["NX_class"] = "NXsample"
            sample.create_dataset("name", data=sample_name)
            sample.create_dataset("rotation_angle", data=np.rad2deg(angles))
            sample['rotation_angle'].attrs['units'] = 'degree'

            data = entry.create_group("data")
            data.attrs["NX_class"] = "NXdata"
            data['data'] = h5py.SoftLink('/entry_1/instrument/detector/data')
            data['rotation_angle'] = h5py.SoftLink('/entry_1/sample/rotation_angle')
            data['image_key'] = h5py.SoftLink('/entry_1/instrument//detector/image_key')
            # X & Y axis data for NeXuS plotting
            data.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
            data.attrs['signal'] = 'data'
            data.attrs['interpretation'] = 'image'
            # Flip to have origin at top, left
            yc = np.flip((np.arange(ny) * px - ny * px / 2) * unit_scale)
            data.create_dataset('row_coords', data=yc)
            data['row_coords'].attrs['units'] = unit_name
            data['row_coords'].attrs['long_name'] = 'Y (%s)' % unit_name
            xc = (np.arange(nx) * px - nx * px / 2) * unit_scale
            data.create_dataset('col_coords', data=xc)
            data['col_coords'].attrs['units'] = unit_name
            data['col_coords'].attrs['long_name'] = 'X (%s)' % unit_name
            data.create_dataset("image_size", data=[px * nx, px * ny])

            process_1 = entry.create_group("process_1")
            process_1.attrs['NX_class'] = 'NXprocess'
            process_1.create_dataset("program", data='PyNX')  # NeXus spec
            process_1.create_dataset("version", data="%s" % _pynx_version)  # NeXus spec

            if process_notes is not None:  # Notes
                for k, v in process_notes.items():
                    if isinstance(v, str):
                        if k not in process_1:
                            process_1.create_dataset(k, data=v)
                    elif isinstance(v, dict) and k == 'note':
                        # Save this as notes:
                        for kk, vv in v.items():
                            i = 1
                            while True:
                                note_s = 'note_%d' % i
                                if note_s not in process_1:
                                    break
                                i += 1
                            note = process_1.create_group(note_s)
                            note.attrs['NX_class'] = 'NXnote'
                            note.create_dataset("description", data=kk)
                            # TODO: also save values as floating-point if appropriate
                            note.create_dataset("data", data=str(vv))
                            note.create_dataset("type", data="text/plain")

            # Configuration & results of process: custom ESRF data policy
            # see https://gitlab.esrf.fr/sole/data_policy/blob/master/ESRF_NeXusImplementation.rst
            config = process_1.create_group("configuration")
            config.attrs['NX_class'] = 'NXcollection'

            command = ""
            for arg in sys.argv:
                command += arg + " "
            process_1.create_dataset("command", data=command)  # CXI spec

            if process_parameters is not None:
                for k, v in process_parameters.items():
                    if v is not None:
                        if type(v) is dict:
                            # This can happen if complex configuration is passed on
                            if len(v):
                                kd = config.create_group(k)
                                kd.attrs['NX_class'] = 'NXcollection'
                                for kk, vv in v.items():
                                    kd.create_dataset(kk, data=vv)
                        else:
                            config.create_dataset(k, data=v)
    if x_rotation_axis is not None or y_rotation_axis is not None:
        with h5py.File(filename, 'a') as h:
            detector = h['/entry_1/instrument/detector']
            if 'x_rotation_axis_pixel_position' in detector:
                detector['x_rotation_axis_pixel_position'] = x_rotation_axis
            else:
                detector.create_dataset('x_rotation_axis_pixel_position', data=x_rotation_axis)
            if 'y_rotation_axis_pixel_position' in detector:
                detector['y_rotation_axis_pixel_position'] = y_rotation_axis
            else:
                detector.create_dataset('y_rotation_axis_pixel_position', data=y_rotation_axis)


def save_obj_probe_chunk(filename, idx, pixel_size, obj_phase=None, obj_complex=None,
                         probe=None, probe_mode_coeff=None,
                         process_notes=None, process_parameters=None, obj_phase_shape=None):
    """
    Save the chunk (the projections included in this object) in a hdf5 file
    :param filename: name of the hdf5 file to save to.
    :param idx: the index of projections to be saved
    :param pixel_size: the object pixel size
    :param obj_phase: the object phase to save, either as a numpy array of a multiprocessing.Array. For
        the latter, the shape must be supplied in obj_phase_shape
    :param obj_complex: the complex projections of the object to save, or None
    :param probe: the complex probe (illumination) to save, or None
    :param probe_mode_coeff: the probe mode coefficients, or None
    :param process_notes: a dictionary which will be saved in '/entry_N/process_1'. A dictionary entry
    can also be a 'note' as keyword and a dictionary as value - all key/values will then be saved
    as separate notes. Example: process={'program': 'PyNX', 'note':{'llk':1.056, 'nb_photons': 1e8}}
    :param process_parameters: a dictionary of parameters which will be saved as a NXcollection
    :param obj_phase_shape: the ph array shape if using multiprocessing.Array
    :return:
    """
    if isinstance(obj_phase, sharedctypes.SynchronizedArray):
        # TODO: handle float32 and float16 cases
        obj_phase = np.frombuffer(obj_phase.get_obj(), dtype=np.float16).reshape(obj_phase_shape)

    f = h5py.File(filename, "w")
    # f.create_dataset("cxi_version", data=150)
    entry = f.create_group("/entry_1")
    entry_path = "/entry_1"
    f.attrs['default'] = 'entry_1'
    f.attrs['creator'] = 'PyNX'
    # f.attrs['NeXus_version'] = '2018.5'  # Should only be used when the NeXus API has written the file
    f.attrs['HDF5_Version'] = h5py.version.hdf5_version
    f.attrs['h5py_version'] = h5py.version.version
    entry.attrs['NX_class'] = 'NXentry'

    entry.create_dataset("program_name", data="PyNX %s" % _pynx_version)
    entry.create_dataset("start_time", data=time.strftime("%Y-%m-%dT%H:%M:%S%z", time.localtime(time.time())))

    # X & Y axis data for NeXuS plotting
    if obj_phase is not None:
        ny, nx = obj_phase.shape[-2:]
    elif obj_complex is not None:
        ny, nx = obj_complex.shape[-2:]
    elif probe is not None:
        ny, nx = probe.shape[-2:]
    else:
        raise RuntimeError("save_obj_probe_chunk: need to give at least one of obj_phase, obj_complex or probe")
    px = pixel_size
    unit_scale = np.log10(max(nx * px, ny * px))
    if unit_scale < -6:
        unit_name = "nm"
        unit_scale = 1e9
    elif unit_scale < -3:
        unit_name = u"µm"
        unit_scale = 1e6
    elif unit_scale < 0:
        unit_name = "mm"
        unit_scale = 1e3
    else:
        unit_name = "m"
        unit_scale = 1

    if obj_phase is not None:
        entry.attrs['default'] = 'object_phase'
        result_1 = entry.create_group("result_1")
        entry["object_phase"] = h5py.SoftLink(entry_path + '/result_1')
        result_1['title'] = 'Object Phase'
        result_1.attrs['NX_class'] = 'NXdata'
        result_1.attrs['signal'] = 'data'

        result_1.create_dataset("projection_idx", data=np.array(idx, dtype=np.int16))
        result_1.create_dataset("data", data=obj_phase, chunks=(1, ny, nx), shuffle=True, **hdf5plugin.LZ4())
        result_1["data"].attrs['interpretation'] = 'image'
        result_1.create_dataset("data_space", data="real")
        result_1.create_dataset("image_size", data=[px * nx, px * ny])
        # Store probe pixel size (not in CXI specification)
        result_1.create_dataset("x_pixel_size", data=px)
        result_1.create_dataset("y_pixel_size", data=px)
        # X & Y axis data for NeXuS plotting
        result_1.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
        # Flip to have origin at top, left
        yc = np.flip((np.arange(ny) * px - ny * px / 2) * unit_scale)
        result_1.create_dataset('row_coords', data=yc)
        result_1['row_coords'].attrs['units'] = unit_name
        result_1['row_coords'].attrs['long_name'] = 'Y (%s)' % unit_name
        xc = (np.arange(nx) * px - nx * px / 2) * unit_scale
        result_1.create_dataset('col_coords', data=xc)
        result_1['col_coords'].attrs['units'] = unit_name
        result_1['col_coords'].attrs['long_name'] = 'X (%s)' % unit_name

    if obj_complex is not None:
        if obj_phase is None:
            entry.attrs['default'] = 'object'
        result_2 = entry.create_group("result_2")
        entry["object"] = h5py.SoftLink(entry_path + '/result_2')
        result_2['title'] = 'Object (complex)'
        result_2.attrs['NX_class'] = 'NXdata'
        result_2.attrs['signal'] = 'data'
        # Assemble object phase
        nproj = len(idx)
        dtype_cplx = np.complex64
        obj = np.empty((nproj, ny, nx), dtype=dtype_cplx)

        result_2.create_dataset("projection_idx", data=np.array(idx, dtype=np.int16))
        result_2.create_dataset("data", data=obj, chunks=True, shuffle=True, **hdf5plugin.LZ4())
        result_2["data"].attrs['interpretation'] = 'image'
        result_2.create_dataset("data_space", data="real")
        result_2.create_dataset("image_size", data=[px * nx, px * ny])
        # Store pixel size (not in CXI specification)
        result_2.create_dataset("x_pixel_size", data=px)
        result_2.create_dataset("y_pixel_size", data=px)
        # X & Y axis data for NeXuS plotting
        result_2.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
        # Flip to have origin at top, left
        yc = np.flip((np.arange(ny) * px - ny * px / 2) * unit_scale)
        result_2.create_dataset('row_coords', data=yc)
        result_2['row_coords'].attrs['units'] = unit_name
        result_2['row_coords'].attrs['long_name'] = 'Y (%s)' % unit_name
        xc = (np.arange(nx) * px - nx * px / 2) * unit_scale
        result_2.create_dataset('col_coords', data=xc)
        result_2['col_coords'].attrs['units'] = unit_name
        result_2['col_coords'].attrs['long_name'] = 'X (%s)' % unit_name

    if probe is not None:
        result_3 = entry.create_group("result_3")
        entry["probe"] = h5py.SoftLink(entry_path + '/result_3')
        result_3['title'] = 'Probe'
        result_3.attrs['NX_class'] = 'NXdata'
        result_3.attrs['signal'] = 'data'
        result_3.create_dataset("data", data=probe, chunks=True, shuffle=True, **hdf5plugin.LZ4())
        result_3["data"].attrs['interpretation'] = 'image'
        result_3.create_dataset("data_space", data="real")
        result_3.create_dataset("image_size", data=[px * nx, px * ny])
        # Store probe pixel size (not in CXI specification)
        result_3.create_dataset("x_pixel_size", data=px)
        result_3.create_dataset("y_pixel_size", data=px)
        # X & Y axis data for NeXuS plotting
        result_3.attrs['axes'] = np.array(['row_coords', 'col_coords'], dtype=h5py.special_dtype(vlen=str))
        # Flip to have origin at top, left
        yc = np.flip((np.arange(ny) * px - ny * px / 2) * unit_scale)
        result_3.create_dataset('row_coords', data=yc)
        result_3['row_coords'].attrs['units'] = unit_name
        result_3['row_coords'].attrs['long_name'] = 'Y (%s)' % unit_name
        xc = (np.arange(nx) * px - nx * px / 2) * unit_scale
        result_3.create_dataset('col_coords', data=xc)
        result_3['col_coords'].attrs['units'] = unit_name
        result_3['col_coords'].attrs['long_name'] = 'X (%s)' % unit_name
        if probe_mode_coeff is not None:
            result_3.create_dataset('probe_mode_coeff', data=probe_mode_coeff)

    command = ""
    for arg in sys.argv:
        command += arg + " "
    process_1 = entry.create_group("process_1")
    process_1.attrs['NX_class'] = 'NXprocess'
    process_1.create_dataset("program", data='PyNX')  # NeXus spec
    process_1.create_dataset("version", data="%s" % _pynx_version)  # NeXus spec
    process_1.create_dataset("command", data=command)  # CXI spec

    if process_notes is not None:  # Notes
        for k, v in process_notes.items():
            if isinstance(v, str):
                if k not in process_1:
                    process_1.create_dataset(k, data=v)
            elif isinstance(v, dict) and k == 'note':
                # Save this as notes:
                for kk, vv in v.items():
                    i = 1
                    while True:
                        note_s = 'note_%d' % i
                        if note_s not in process_1:
                            break
                        i += 1
                    note = process_1.create_group(note_s)
                    note.attrs['NX_class'] = 'NXnote'
                    note.create_dataset("description", data=kk)
                    # TODO: also save values as floating-point if appropriate
                    note.create_dataset("data", data=str(vv))
                    note.create_dataset("type", data="text/plain")

    # Configuration & results of process: custom ESRF data policy
    # see https://gitlab.esrf.fr/sole/data_policy/blob/master/ESRF_NeXusImplementation.rst
    config = process_1.create_group("configuration")
    config.attrs['NX_class'] = 'NXcollection'
    if process_parameters is not None:
        for k, v in process_parameters.items():
            if v is not None:
                if type(v) is dict:
                    # This can happen if complex configuration is passed on
                    if len(v):
                        kd = config.create_group(k)
                        kd.attrs['NX_class'] = 'NXcollection'
                        for kk, vv in v.items():
                            kd.create_dataset(kk, data=vv)
                else:
                    config.create_dataset(k, data=v)
