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

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

import argparse
from ...mpi import MPI


class RawTextArgumentDefaultsHelpFormatter(argparse.RawTextHelpFormatter,
                                           argparse.ArgumentDefaultsHelpFormatter):
    pass


class ActionPtychoMotors(argparse.Action):
    """Argparse Action to check --ptychomotors is correctly used"""

    def __call__(self, parser_, namespace, values, option_string=None):
        try:
            if len(values) == 1:
                vv = values[0].split(',')
            elif len(values) not in [2, 4]:
                raise argparse.ArgumentError(
                    self, f"--ptycho_motors: need either 1 ('pix,piz' or "
                          f"'pix,piz,-y,x') or two ('pix piz') or 4 arguments"
                          f"('pix piz -y x), not {len(values)}. Did not understand: "
                          f"--ptycho_motors {' '.join(values)}"
                )
            else:
                vv = values
        except:
            raise argparse.ArgumentError(
                self, f"--ptycho_motors: need either 1 ('pix,piz' or "
                      f"'pix,piz,-y,x') or two ('pix piz') or 4 arguments"
                      f"('pix piz -y x), not {len(values)}. Did not understand: "
                      f"--ptycho_motors {' '.join(values)}"
            )

        setattr(namespace, self.dest, vv)


class ActionROI(argparse.Action):
    """Argparse Action to check --roi is correctly used"""

    def __call__(self, parser_, namespace, values, option_string=None):
        if len(values) == 1:
            if values[0].count(',') == 3:
                try:
                    vv = [int(v) for v in values[0].split(',')]
                except ValueError:
                    raise argparse.ArgumentError(
                        self, f"--roi {' '.join(values)}: error "
                              f"converting 4 coordinates to integers")
                nx, ny = vv[1] - vv[0], vv[3] - vv[2]
                if nx <= 0 or ny <= 0 or (nx % 2 != 0) or (ny % 2 != 0) or (nx != ny):
                    raise argparse.ArgumentError(
                        self, f"--roi={values[0]}: need an even "
                              f"and equal size along both dimensions (nx={nx}, ny={ny})")
                setattr(namespace, self.dest, vv)
            elif values[0] not in ['auto', 'full', 'all']:
                raise argparse.ArgumentError(
                    self, "--roi: need either 'auto' (default), 'full'"
                          " or 4 separate numbers with the xmin xmax ymin ymax")
            else:
                setattr(namespace, self.dest, values[0])
        elif len(values) == 4:
            try:
                vv = [int(v) for v in values]
            except ValueError:
                raise argparse.ArgumentError(
                    self, f"--roi {' '.join(values)}: error "
                          f"converting 4 coordinates to integers")
            nx, ny = vv[1] - vv[0], vv[3] - vv[2]
            if nx <= 0 or ny <= 0 or (nx % 2 != 0) or (ny % 2 != 0) or (nx != ny):
                raise argparse.ArgumentError(
                    self, f"--roi {' '.join(values)}: need an even "
                          f"and equal size along both dimensions (nx={nx}, ny={ny})")
            setattr(namespace, self.dest, vv)
        else:
            raise argparse.ArgumentError(
                self, "--roi: need either 'auto' (default), 'full' "
                      "or 4 separate numbers with the xmin xmax ymin ymax")


def make_parser(prog: str, description: str, epilog: str, default_params: dict):
    # TODO: use a default value defined in runner.py for all parameters
    p = default_params
    parser = argparse.ArgumentParser(prog=prog,
                                     description=description,
                                     formatter_class=RawTextArgumentDefaultsHelpFormatter,
                                     epilog=epilog,
                                     conflict_handler='resolve')
    grp_in = parser.add_argument_group("Input parameters")

    class ActionScan(argparse.Action):
        """Argparse Action to check --scan is correctly used"""

        def __call__(self, parser_, namespace, values, option_string=None):
            vv = []
            for v in values:
                v = v.replace('"', '').replace("'", "")
                try:
                    if 'range(' in v or ',' in v:
                        vv += list(eval(v))
                    else:
                        vv.append(int(v))
                except:
                    raise argparse.ArgumentTypeError(
                        f"--scan: failed converting '{v}' to a scan number or list of")
            setattr(namespace, self.dest, vv)

    grp_in.add_argument(
        '--scan', default=None, nargs='+', action=ActionScan,
        help='scan number e.g. --scan 5 to extract a given scan number from a series. '
             'Alternatively a list or range of scans can be given using "--scan 12 23 45"'
             'or --scan "range(12,25)" (note the quotes, necessary when using range with ()).')
    grp_in.add_argument(
        '--max_frame', '--maxframe', type=int, default=p['maxframe'],
        dest='maxframe', help='maximum number of frames to import')
    grp_in.add_argument(
        '--max_size', '--maxsize', type=int, default=p['maxsize'],
        dest='maxsize',
        help='Maximum frame size to use, while keeping center of gravity of diffraction '
             'in the cut frame center.')

    class ActionModuloFrame(argparse.Action):
        """Argparse Action to check --moduloframe is correctly used"""

        def __call__(self, parser_, namespace, values, option_string=None):
            if len(values) == 1 and ',' in values[0]:
                # For conversion from old-style arguments e.g. moduloframe=2,1
                values = values[0].split(',')
            try:
                vv = [int(v) for v in values]
            except ValueError:
                raise argparse.ArgumentError(
                    self, f"--modulo_frame: need integer values: {' '.join(values)}")
            if len(values) == 1:
                setattr(namespace, self.dest, [vv[0], 0])
            elif len(vv) == 2:
                if vv[1] >= vv[0]:
                    raise argparse.ArgumentError(
                        self, f"--modulo_frame N1 N2: need 0 <= N2={vv[1]} <= N1={vv[0]}")
                setattr(namespace, self.dest, [vv[0], vv[1]])
            else:
                raise argparse.ArgumentError(self, "--modulo_frame: need 1 or 2 integer arguments")

    grp_in.add_argument(
        '--modulo_frame', '--moduloframe', type=str, nargs='+',
        dest='moduloframe', action=ActionModuloFrame,
        help='Instead of taking all sequential frames, load one every N. If two numbers are\n'
             'given (--moduloframe N1 N2), N2 must be smaller than N1 and the frames loaded\n'
             'will be those with i %% N1 == N2.')
    grp_in.add_argument(
        '--mask', '--loadmask', '--load_mask', action='store', type=str,
        default=None, required=False, dest='loadmask',
        help='Specify th mask to use for detector data, which should have the same 2D'
             'shape as the raw detector data.'
             'This should be a boolean or integer array with good pixels=0 and bad ones>0'
             '(values are expected to follow the CXI convention)'
             'Acceptable formats/values:\n'
             '- "mask.npy", "mask.npz" (the first data array will be used)\n'
             '- "mask.edf" or "mask.edf.gz" (a single 2D array is expected)\n'
             '- "mask.h5:/entry_1/path/to/mask": hdf5 format with the full path to the '
             '  2D array. hdf5 is also accepted as extension.\n'
             '- "maxipix": if this special name is entered, the masked pixels will be rows '
             '  and columns multiples of 258+/-3\n')
    grp_in.add_argument(
        '--mask_iobs_max', '--maskiobsmax', type=float,
        default=p['mask_iobs_max'],
        help='if given, all pixels with an observed intensity >= the given '
             'value are masked. This is applied frame-by-frame.')

    grp_in.add_argument(
        '--roi', type=str, default=p['roi'], nargs='+', action=ActionROI,
        help='Region-of-interest to be used for actual inversion. This requires 4 values'
             'for xmin xmax ymin ymax. '
             'The area is taken with python conventions, i.e. pixels with indices xmin<= x < xmax and '
             'ymin<= y < ymax.'
             'Additionally, the shape of the area must be square, and '
             'n=xmax-xmin=ymax-ymin must also be a suitable integer number'
             'for OpenCL or CUDA FFT, i.e. it must be a multiple of 2 and the largest number in'
             'its prime factor decomposition must be less or equal to the largest value'
             'acceptable by vkFFT for a radix transform(<=13).'
             'If n does not fulfill these constraints,'
             'it will be reduced using the largest possible integer smaller than n.'
             'This option supersedes "maxsize" unless roi="auto".\n'
             'Other possible values:\n'
             '- "auto": automatically selects the roi from the center of mass '
             '          and the maximum possible size. [default]\n'
             '- "all" or "full": use the entire, uncentered frames. Only useful for pre-processed'
             '         data. Cropping may still be performed to get a square and '
             '         FFT-friendly size. Near-field should default to "full"\n')
    grp_in.add_argument(
        '--rebin', '--binning', '--bin', type=int, default=p['rebin'],
        help='the experimental images can be rebinned (i.e. a group of n x n pixels is replaced by a '
             'single one whose intensity is equal to the sum of all the pixels). Binning is '
             'performed last: the ROI, mask, background, pixel size should all correspond to '
             'full (non-rebinned) frames..\n'
             'This is recommended if the size of the probe (including tails) is smaller than half '
             'of the Fourier (the probe array) window (then use ``--rebin 2``): reconstruction '
             'will be 4x faster without any loss of resolution. It is also recommended if the '
             'extent of the scanned area is smaller than the size of the Fourier (probe) array.')
    grp_in.add_argument(
        '--xy', type=str, default=p['xy'],
        help='expression to be used for the XY positions (e.g. "-x,y",...), to compute the'
             'values from the raw input (scaling, flipping, etc..). Mathematical '
             'operations can also be used, e.g.: ``--xy=0.5*x+0.732*y,0.732*x-0.5*y``')
    grp_in.add_argument(
        '--detector_orientation', '--detectororientation', type=str, nargs=3,
        default=p['detector_orientation'],
        help='three optional flags which, if 1, will do in this order: '
             'array transpose (x/y exchange), flipud, fliplr. '
             'The changes also apply to the mask')
    grp_in.add_argument(
        '--flatfield', '--flat_field', '--flat', type=str, default=None,
        help='flatfield path for the correction to be applied to the detector data. The array must '
             'have the same shape as the frames, which will be multiplied by this '
             'correction.\n'
             'Acceptable formats:  flat.npy or flat.npz (the first data array will be used), '
             ' flat.edf or flat.edf.gz (a single 2D array is expected), '
             '"flat.h5:/entry_1/path/to/flat" (hdf5 format with the full path to the 2D array, '
             'hdf5 is also accepted as extension), flat.mat (from a matlab file. The first '
             'array found is loaded)')
    grp_in.add_argument(
        '--dark', type=str, default=None,
        help='the dark correction file (incoherent background). The array must have '
             'have the same shape as the frames, which will be multiplied by this '
             'correction.\n'
             'Acceptable formats:  flat.npy or flat.npz (the first data array will be used), '
             ' flat.edf or flat.edf.gz (a single 2D array is expected), '
             '"flat.h5:/entry_1/path/to/flat" (hdf5 format with the full path to the 2D array, '
             'hdf5 is also accepted as extension), flat.mat (from a matlab file. The first '
             'array found is loaded)')
    grp_in.add_argument(
        '--dark_subtract', type=float, default=False, const=True, nargs='?',
        help='use this option to subtract the dark from the observed'
             'intensity. This is DEPRECATED, as it will mess up the Poisson statistics'
             'of the observed data (and so alters the ML optimisation). As long as '
             '--dark is used, the background is always "subtracted" dynamically '
             'during algorithms, but without altering the statistics.'
             'If given as a simple keyword, the dark is pre-subtracted. If given with a float,'
             'the dark will be multiplied by this factor before subtraction')
    grp_in.add_argument(
        '--orientation_round_robin', '--orientationroundrobin', action='store_true',
        help='use this option to test all possible combinations of xy and detector_orientation '
             'to find the correct detector configuration. There are 64 possibilities, '
             'with a redundancy of 8.')
    grp_in.add_argument(
        '--nrj', '--energy', action='store', default=p['nrj'], type=float,
        help='X-ray beam energy in keV (not needed if read from the data files)')
    grp_in.add_argument(
        '--pixel_size', '--pixelsize', action='store', default=p['pixelsize'], type=float,
        dest='pixelsize',
        help='detector pixel size in meters (not needed if read from the data files)')
    grp_in.add_argument(
        '--near_field', '--nearfield', action='store_true', default=p['near_field'],
        help="Use this for near-field ptycho optimisation. It is normally "
             "automatically triggered by the script called, but may be useful in some "
             "cases (e.g. near-field data in CXI format)")
    grp_in.add_argument(
        '--detector_distance', '--detectordistance', '--distance',
        type=float, dest='detectordistance', default=p['detectordistance'],
        help='Detector distance in meters. Mandatory unless included in the data file.')

    grp_algo = parser.add_argument_group("Algorithms")
    grp_algo.add_argument(
        '--algorithm', type=str, default=p['algorithm'],
        help='algorithm chain used for the optimization. This chain:\n\n'
             '* is divided in "steps" separated by commas, e.g: ``ML**50,DM**100``\n'
             '* is interpreted from right to left, as a mathematical operator to an '
             'object on the right-hand side\n'
             '* should not contain any space, unless the whole string is given between quotes ("")\n'
             '\nThere are two types of commands in the chain:\n\n'
             '1) **commands** which change basic parameters or perform some analysis:\n\n'
             '* ``probe=1`` or 0: activate or deactivate the probe optimisation (by default'
             '  only the object is optimised)\n'
             '* ``object=1`` or 0: activate or deactivate the object optimisation\n'
             '* ``background=1`` or 0: activate or deactivate the background optimisation.\n'
             '  When set to 1, this will initialise the background to at least 1e-2 to\n'
             '  enable the background optimisation.\n'
             '* ``background_smooth=3``: gaussian sigma for smoothing the updated background\n'
             '  [default:3, large values are possible and will use FFT convolution]\n'
             '* ``position=N`` or 0: activate or deactivate (0) position optimisation every N cycles\n'
             '  (preferably for AP or ML, some prior convergence is needed. Can also work\n'
             '  with DM but is not recommended). Recommended value is every 5 cycles\n'
             '* ``pos_mult=1``: multiplier for the calculated position shift. Can be used to accelerate\n'
             '  convergence or make it more cautious. (default:5, suitable for an object\n'
             '  with reasonable contrast)\n'
             '* ``pos_max_shift``: maximum shift of position update in pixels between iterations (default:2)\n'
             '* ``pos_min_shift``: minimum shift of position update in pixels between iterations (default:0)\n'
             '* ``pos_threshold``: if the integrated norm of the object gradient multiplied by the probe\n'
             '  lower than the average value (for all positions) multiplied by this threshold,\n'
             '  the position is not changed. This allows to avoid updating positions in areas\n'
             '  where the object is flat, and sensitivity to shifts is low. [default:0.2]\n'
             '* ``nbprobe=3``: change the number of modes for the probe (can go up or down)\n'
             '* ``regularization=1e-4``: setting the regularization parameter for the object, to penalize\n'
             '  local variations in ML runs and smooth the solution\n'
             '* ``obj_smooth=1.5`` or ``probe_smooth=1.0``: these parameters will partially smooth the object\n'
             '  and probe, softening the resulting arrays. This applies to DM and AP algorithms.\n'
             '* ``obj_inertia=0.01`` or ``obj_probe=0.001``: these parameters set the inertia of the object\n'
             '  and/or probe update, yielding more stable result. This applies to DM and AP algorithms.\n'
             '* ``ortho``: will perform orthogonalisation of the probe modes. The modes are sorted by\n'
             '  decreasing intensity.\n'
             '* ``analysis``: perform an analysis of the probe (propagation, modes). Useful combined\n'
             '  with "saveplot" to save the analysis plots\n'
             '2) **operators** which indicate the actual algorithm to apply to the object and probe:\n\n'
             '* ``AP``: alternate projections. Slow but converging algorithm\n'
             '* ``DM``: difference map. Fast early convergence, oscillating after.\n'
             '* ``ML``: maximum likelihood conjugate gradient (Poisson-noise). Robust, converging,'
             '  for final optimization.\n'
             'These operators can be combined mathematically, e.g.:\n\n'
             '* ``DM**100``: corresponds to 100 cycles of difference map\n'
             '* ``ML**40*DM**100``: 100 cycles of DM followed by 40 cycles of ML (note the order)\n'
             ' \nExample algorithms chains:\n\n'
             '* ``--algorithm ML**40,DM**100,probe=1``: activate probe optimisation, then 100 DM and 40 ML\n'
             '* ``--algorithm ML**100,DM**200,nbprobe=3,ML**40,DM**100,probe=1,DM**100``: first DM with'
             '  object update only,  then 100 DM also updating the probe, then use 3 probe modes'
             '  and do 100 DM followed by 40 ML\n'
             '* ``--algorithm ML**100*AP**200*DM**200,probe=1``: 200 DM then 200 AP then 100 ML (one step)\n'
             '* ``--algorithm "(ML**10*DM**20)**5,probe=1"``:'
             '  repeat 5 times [20 cycles of DM followed by 5 cycles of ML]'
             '  (note the quotes necessary for the parenthesis)\n'
    )
    grp_algo.add_argument('--nb_run', '--nbrun', type=int, default=1,
                          dest='nbrun', help='number of times to run the optimisation')
    grp_algo.add_argument('--run0', type=int, default=None,
                          help='number for the first run (can be used to overwrite previous run results)')
    grp_algo.add_argument('--no_rerun', '--norerun', '--no_re_run',
                          action='store_true', default=False,
                          help='if this option is used, the scan will be skipped if the result file \n'
                               'already exists. This allows to run batch jobs for multiple scans with \n'
                               'low priority (and which can be preempted) and restart them without \n'
                               'reprocessing scans. [default: save to a new run if saveprefix has a %%d\n'
                               'field for the run number, and overwrite otherwise]\n')
    grp_algo.add_argument(
        '--no_auto_center', '--noautocenter', action='store_true',
        help='Use this option to disable auto-centering of object and probe ofter'
             'each optimisation step. This is done by default to avoid drifts.')
    # TODO: deprecate --auto_center in favour of --no_auto_center ?
    grp_algo.add_argument(
        '--auto_center', '--autocenter', type=int, default=1,
        choices=[0, 1], dest='autocenter',
        help='Use this option to disable (0) auto-centering of object and probe after '
             'each optimisation step. This is done by default to avoid drifts (notably during DM).')
    grp_algo.add_argument(
        '--center_probe_n', type=int, default=5,
        help='During DM, check for a probe drift every N cycles given by this argument.'
             'Also see center_probe_max_shift. Ignored for near field ptycho.')
    grp_algo.add_argument(
        '--center_probe_max_shift', type=int, default=5,
        help='During DM, when checking for a probe drift, the object and probe positions '
             'will be corrected if the center deviates by more the given number of pixels.'
             'Ignored for near field ptycho.')
    grp_algo.add_argument(
        '--dm_loop_obj_probe', type=int, default=1,
        help='during DM, when both object and probe are updated, it may be '
             'better to loop the object and probe update for a more stable '
             'optimisation (but slower).')
    grp_algo.add_argument(
        '--use_direct_beam', '--usedirectbeam', action='store_true', default=False,
        help='if used, and a reference frame with the direct beam is given, it'
             'will also be used for the optimisation, giving an absolute reference for the probe')
    grp_algo.add_argument('--ml_obj_regularisation', '--regularization', type=float, default=0,
                          help="object regularisation factor to smooth during ML")
    grp_algo.add_argument('--obj_smooth', '--objsmooth', type=float, default=0,
                          help="object smoothing width during projection algorithms. Must be "
                               "associated with obj_inertia>0 for effect.")
    grp_algo.add_argument('--obj_inertia', '--objinertia', type=float,
                          default=p['obj_inertia'],
                          help="object inertia (0-1) during projection algorithms.")
    grp_algo.add_argument('--probe_smooth', '--probesmooth', type=float, default=0,
                          help="probe smoothing width during projection algorithms. Must be "
                               "associated with probe_inertia>0 for effect.")
    grp_algo.add_argument('--probe_inertia', '--probeinertia', type=float,
                          default=p['probe_inertia'],
                          help="probe inertia (0-1) during projection algorithms.")
    grp_algo.add_argument('--pos_mult', '--posmult', type=float, default=5,
                          help="multiplier for the calculated position shift. Can be used to accelerate "
                               "convergence or make it more cautious. (default:5, suitable "
                               "for an object with reasonable contrast)")
    grp_algo.add_argument('--pos_max_shift', '--posmaxshift', type=float, default=1,
                          help="maximum shift of position update in pixels between iterations")
    grp_algo.add_argument('--pos_min_shift', '--posminshift', type=float, default=0,
                          help="minimum shift of position update in pixels between iterations - "
                               "this can be used to ignore small shifts (due to noise)")
    grp_algo.add_argument(
        '--pos_threshold', '--posthreshold', type=float, default=0.2,
        help="if the integrated norm of the object gradient multiplied by the probe is lower "
             "than the average value (for all positions) multiplied by this threshold, "
             "the position is not changed. This allows to avoid updating positions "
             "in area where the object is flat, and sensitivity to shifts is low.")
    grp_algo.add_argument(
        '--background_smooth', '--backgroundsmooth', type=float, default=3,
        help="gaussian sigma to smooth the updated background")
    # use (bilinear) interpolation during projection ?
    grp_algo.add_argument('--interpolation', action='store_true', default=False,
                          help=argparse.SUPPRESS)
    grp_algo.add_argument('--dm_alpha', '--dmalpha', action='store', default=0.02, type=float,
                          help="alpha parameter to mix a small amount of AP with DM")
    grp_algo.add_argument('--raar_beta', '--raarbeta', action='store', default=0.9, type=float,
                          help="RAAR beta parameter")
    grp_algo.add_argument(
        '--padding', action='store', default=0, type=int,
        help=argparse.SUPPRESS)
    grp_algo.add_argument(
        '--padding_window', '--paddingwindow', action='store', default=0, type=int,
        help=argparse.SUPPRESS)

    grp = parser.add_argument_group("Display")
    grp.add_argument('--live_plot', '--liveplot', action='store_true', default=False,
                     dest='liveplot',
                     help='liveplot during optimisation\n')
    grp.add_argument('--verbose', type=int, default=p['verbose'],
                     help="print evolution of llk (and display plot if 'liveplot' is set) every N cycle")
    grp.add_argument('--fig_num', type=int, default=100,
                     help="figure number to display object and probe")

    grp = parser.add_argument_group("Output")
    grp.add_argument('--save_plot', '--saveplot', action='store', type=str,
                     default=p['saveplot'], const=True, dest='saveplot',
                     choices=['object_phase', 'object_rgba', 'object_real',
                              'object_imag', 'object_abs'], nargs='?',
                     help='will save plot at the end of the optimization (png file).\n'
                          'A string can be given to specify if only the object phase should be plotted.')
    grp.add_argument('--save_prefix', '--saveprefix', action='store', type=str,
                     default=p['saveprefix'], dest='saveprefix',
                     help='prefix to save the optimized object and probe '
                          '(as a .cxi or .npz file) and optionally image (png). '
                          'Use "saveprefix=none" to disable saving (e.g. for tests). '
                          'Several format options can be used, including the "run" number, '
                          'the "scan" number, the data directory "data_dir" or the'
                          'data prefix "data_prefix" (data filename without extension). '
                          'Examples:\n'
                          '  --save_prefix ResultsScan{scan:04d}/Run{run:04d} : the default, using'
                          '    the scan and run number to identify the result\n'
                          '  --save_prefix ResultsScan%%04d/Run%%04d : same as the previous '
                          '    format, but using old-style formatting (deprecated)\n'
                          '  --save_prefix ResultsScan{scan:04d} : only use the scan number'
                          '    in the output format (new runs will overwrite the previous ones\n'
                          '  --save_prefix {data_prefix}_run{run:04d} : combining the original'
                          '    data name and the run number\n'
                          '  --save_prefix {data_dir}/{data_prefix}_results/Run{run:04d} : '
                          '    saving along the original data file. \n'
                          '')
    grp.add_argument('--output_format', '--outputformat', action='store', type=str,
                     default='cxi', choices=['cxi', 'npz', 'nxtomo'],
                     help='choose the output format for the final object and support.'
                          'Note: nxtomo is only for ptycho-tomo, and requires the rotation angle of '
                          'each projection. This is only supported by the ID16A NFP runner.')
    grp.add_argument('--cxi_output', '--cxioutput', action='store', type=str,
                     default=p['cxi_output'],
                     choices=['object_probe', 'object_phase', 'object', 'probe'],
                     help='Choice of data included in the cxi result file.'
                          'This only affects the large object and probe arrays. '
                          'Ignored for --output_format nxtomo. Choices:\n'
                          '- object_probe: all object and probe modes are saved as complex objects.\n'
                          '- object_phase: only the first mode of the object phase is saved (as float16).\n'
                          '  This is the format which saves the most space'
                          '- object: save only the complex object\n'
                          '- probe: save only the complex probe')
    grp.add_argument('--remove_obj_phase_ramp', action='store_true',
                     default=p['remove_obj_phase_ramp'],
                     help='Use this option to save the final object after '
                          'removing the phase ramp estimated from the imperfect '
                          'centring of the diffraction data (sub-pixel shift). Calculated '
                          'diffraction patterns using such a corrected object will present '
                          'a sub-pixel shift relative to the diffraction data.')
    grp.add_argument('--save', action='store', type=str,
                     default=p['save'], choices=['final', 'all'],
                     help='When to save object and probe. Using "all" will save '
                          'after each comma-separated algorithm step.')
    grp.add_argument('--data2cxi', action='store', default=False, const=True,
                     nargs='?', choices=['crop'],
                     help='Option to save the raw data in CXI format (http://cxidb.org/cxi.html), '
                          'with all the required information for a ptychography experiment '
                          '(energy, detector distance, scan number, translation axis are '
                          'all required). if data2cxi=crop is used, the data will be saved '
                          'after centering and cropping (default is to save '
                          'the raw data). If this keyword is present, the processing stops '
                          'after exporting the data.')
    grp.add_argument('--movie', action='store_true', default=None,
                     help='create a movie of the scan with all scan positions and diffraction '
                          'frames. Requires matplotlib and ffmpeg.')

    grp = parser.add_argument_group("Initial object and probe")
    grp.add_argument(
        '--obj_max_pix', '--objmaxpix', type=int, default=p['obj_max_pix'],
        help='maximum size in pixels for the object. This is mostly to avoid mistakes when\n'
             'e.g. scan coordinates are incorrectly scaled (*1e6) and result in absurd dimensions.')
    grp.add_argument(
        '--obj_margin', '--objmargin', type=int, default=p['obj_margin'],
        help='margin (in pixels) around the calculated object area. this is useful when \n'
             'refining positions, to avoid getting outside the object area.')
    grp.add_argument('--load', action='store', type=str,
                     default=None,
                     help='load object and probe from previous optimization, giving a path'
                          'to a .cxi or .npz result file. Note that the'
                          'object and probe will be scaled if the number of pixels is '
                          'different for the probe..')
    grp.add_argument('--load_probe', '--loadprobe', action='store', type=str,
                     default=None, dest='loadprobe',
                     help='load only probe from previous optimization, giving a path'
                          'to a .cxi or .npz result file. Note that the'
                          ' probe will be scaled if the number of pixels is '
                          'different for the probe..')
    grp.add_argument('--load_pixel_size', '--loadpixelsize', '--load_pixelsize',
                     action='store', type=float, default=None, dest='loadpixelsize',
                     help='Pixel size (in meters) from a loaded probe '
                          '(and possibly object). If the pixel size is different, '
                          'the loaded arrays will be scaled to match the new pixel size. '
                          '[default: when loading previous files, object/probe pixel size is '
                          'calculated from the size of the probe array, assuming same detector '
                          'distance and pixel size]')
    grp.add_argument('--probe', action='store', type=str,
                     default=p['probe'],
                     help='String defining the initial probe, mandatory '
                          'if --load or --loadprobe is not used.\nExamples:\n\n'
                          '* focus,60e-6x200e-6,0.09: slits size (horizontal x vertical), '
                          'focal distance (all in meters)\n'
                          '* focus,200e-6,0.09: radius of the circular aperture, '
                          'focal distance (all in meters)\n'
                          '* gaussian,100e-9x200e-9: gaussian type with horizontal '
                          'x vertical FWHM, both given in meters\n'
                          '* disc,100e-9: disc-shape, with diameter given in meters')
    grp.add_argument('--defocus', action='store', type=float,
                     default=None,
                     help='defocus distance (+: towards detector). The initial probe is propagated '
                          'by this distance before being used. This is true both for calculated probes '
                          '(using probe=...)  and for probes loaded from a previous file')
    grp.add_argument('--rotate', action='store', type=float,
                     default=None,
                     help='rotate the initial probe (either simulated or loaded) by X degrees')
    grp.add_argument('--object', action='store', type=str,
                     default=p['object'],
                     help='String to define the original object. It will be initialised'
                          'over the entire area using random values. The four given values '
                          'correspond to the amplitude min and max and the phase '
                          'min and max (in radians).')
    grp.add_argument(
        '--multiscan_reuse_ptycho', '--multi_scan_reuse_ptycho',
        '--multiscanreuseptycho', const=True, nargs='?', default=None,
        help='if used as a keyword, successive scans will re-start from the previous '
             'ptycho object and probe, and skip initialisation steps. Useful for ptycho-tomo.'
             'This can also be used to supply a shorter algorithm chain which will be used after the first'
             'scan, e.g. with:\n'
             ' ``--algorithm ML**100,AP**100,DM**1000,nbprobe=2,probe=1'
             '--multiscan_reuse_ptycho ML**100,AP**100,probe=1,AP**50,probe=0``\n'
             'In the above example, the probe would be re-used so there is no need to use ``nbprobe=2`` ,'
             'and the first step (``AP**50``) would only optimise the object.\n'
             'Note that if you want to re-process the first scan with this short algorithm chain,'
             'it is possible to list the first scan twice, e.g.: ``scan=12,12,13,14,15``.'
             '[By default, every scan starts from a new object and probe]')

    grp = parser.add_argument_group("MPI options")

    class ActionMPI(argparse.Action):
        """Argparse Action to check MPI is really availabe if --mpi is used"""

        def __call__(self, parser_, namespace, values, option_string=None):
            if MPI is None:
                raise argparse.ArgumentError(self, "--mpi was given but MPI is not available")
            else:
                if MPI.COMM_WORLD.Get_size() == 1:
                    raise argparse.ArgumentError(self, '--mpi was given but MPI rank = 1')
            if values not in ['split', 'scan', 'multiscan']:
                raise argparse.ArgumentError(
                    self, "--mpi: possible choices are: 'split', 'scan' (same as 'multiscan')")
            if values == 'scan':
                values = 'multiscan'
            setattr(namespace, self.dest, values)

    grp.add_argument(
        '--mpi', action=ActionMPI, type=str,
        default=p['mpi'], choices=['split', 'scan', 'multiscan'],
        help='mpi option (only when launching the script using mpiexec): '
             'either distribute the list of scans to different processes '
             '(scan or multiscan, the default), or split a large scan '
             'in different parts, which are automatically aligned and'
             'stitched. Examples:\n\n'
             '* ``mpiexec -n 2 pynx-ptycho-cxi scan=11,12,13,14 data=scan%%02d.cxi mpi=multiscan '
             'probe=focus,120e-6x120e-6,0.1 defocus=100e-6 verbose=50 '
             'algorithm=analysis,ML**100,DM**200,nbprobe=2,probe=1``\n'
             '* ``mpiexec -n 4 pynx-ptycho-cxi  mpi=split data=scan151.cxi verbose=50 '
             'probe=focus,120e-6x120e-6,0.1 defocus=100e-6 '
             'algorithm=analysis,ML**100,DM**200,nbprobe=2,probe=1``')
    grp.add_argument('--mpi_split_nb_overlap', action='store', type=int,
                     default=1,
                     help='depth of shared neighbours between overlapping regions '
                          'when using "--mpi split"')
    grp.add_argument('--mpi_split_nb_neighbour', action='store', type=int,
                     default=20,
                     help='target number of shared neighbours for each set for overlap '
                          'when using "--mpi split"')

    grp = parser.add_argument_group("GPU")
    grp.add_argument('--gpu', action='store', type=str, default=p['gpu'],
                     help='string matching GPU name (or part of it, case-insensitive)')
    grp.add_argument('--profile', '--profiling', action='store_true', default=False,
                     dest='profiling', help=argparse.SUPPRESS)
    grp.add_argument('--stack_size', '--stacksize', action='store',
                     type=int, default=None,
                     help='number of frames processed together. This is normally '
                          'set automatically to a large value for better GPU performance. '
                          'Smaller values than the number of frames can be used to avoid '
                          'save memory during ML. An optimal value is such that '
                          '(nb_frame %% stack_size) is as small as possible.')
    grp = parser.add_argument_group("Miscellaneous")
    grp.add_argument('--instrument', action='store', type=str, default=p['instrument'],
                     help='instrument name (mostly for CXI output).')
    grp.add_argument('--share_probe', '--share-probe', nargs='?',
                     const=-1, type=int, choices=[-1, 0, 1, 2, 3, 4, 5, 6],
                     help='Use this option to process multiple scans with the same probe, '
                          'which will be shared periodically using a tailored algorithm '
                          'chain, e.g.:\n'
                          '  ``--algorithm "LoopScan(ML**50)**20,LoopScan(DM**50)**40,probe=1"``\n'
                          'will perform 2000 DM followed 1000 ML cycles, averaging the probe '
                          'every 50 cycles.\n'
                          'By default the probe will be shared while the optimisation loops '
                          'over the different scans, but by giving a integer value n, the '
                          'probe can instead be averaged at the end of each loop (value=1) '
                          'or decomposed (using truncated SVD) if n>1. '
                          '[EXPERIMENTAL]')

    return parser
