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

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

"""
This package includes tests for the CDI command-line scripts.
"""

import os
import platform
import sys
import subprocess
import unittest
import tempfile
import warnings

has_mpi = False
try:
    import mpi4py
    import shutil

    if shutil.which('mpiexec') is not None:
        has_mpi = True
except ImportError:
    pass

from pynx.ptycho.test.test_ptycho import make_ptycho_data, make_ptycho_data_cxi
from pynx.processing_unit import has_cuda, has_opencl

exclude_cuda = False
exclude_opencl = False
if 'PYNX_PU' in os.environ:
    if 'opencl' in os.environ['PYNX_PU'].lower():
        exclude_cuda = True
    elif 'cuda' in os.environ['PYNX_PU'].lower():
        exclude_opencl = True

if 'opencl' in sys.argv or '--opencl' in sys.argv or not has_cuda:
    exclude_cuda = True
if 'cuda' in sys.argv or '--cuda' in sys.argv or not has_opencl:
    exclude_opencl = True

oldparse = True if 'oldparse' in sys.argv or '--oldparse' in sys.argv else False
liveplot = True if 'liveplot' in sys.argv or '--liveplot' in sys.argv \
                   or 'live_plot' in sys.argv or '--live_plot' in sys.argv else False


class TestPtychoRunner(unittest.TestCase):
    """
    Class for tests of the Ptycho runner scripts
    """

    @classmethod
    def setUpClass(cls):
        # Directory contents will automatically get cleaned up on deletion
        cls.tmp_dir_obj = tempfile.TemporaryDirectory()
        cls.tmp_dir = cls.tmp_dir_obj.name

    @unittest.skipUnless(oldparse, "old parsing tests not selected")
    def test_ptycho_runner_old(self):
        # Old parser (deprecated)
        path = make_ptycho_data_cxi(dsize=160, nb_frame=40, nb_photons=1e9, dir_name=self.tmp_dir)
        options = {'cxi': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                           'algorithm=ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                           'verbose=10', 'save=all', 'saveplot'],
                   'cxi_stack16': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                   'algorithm=ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                   'stack_size=16'],
                   'cxi_stack20': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                   'algorithm=ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                   'stack_size=20'],
                   'cxi_saveprefix1': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                       'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                       'verbose=10', 'save=all', 'saveplot', 'saveprefix=test%02d_%02d'],
                   'cxi_saveprefix2': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                       'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                       'verbose=10', 'save=all', 'saveplot', 'saveprefix=test%02d',
                                       'defocus=200e-6'],
                   'cxi_saveprefix3': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                       'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                       'verbose=10', 'save=all', 'saveplot', 'saveprefix=test',
                                       'defocus=-200e-6'],
                   'cxi_2scan': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                 'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                 'verbose=10', 'save=all', 'scan=1,1'],
                   'cxi_no_rerun': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                    'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                    'verbose=10', 'save=all', 'no_rerun'],
                   'cxi_2scan_no_rerun': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                          'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                          'verbose=10', 'save=all', 'no_rerun', 'scan=1,1'],
                   'cxiobj': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                              'algorithm=DM**20,nbprobe=1,probe=1', 'cxi_output=object'],
                   'cxiobjph': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                'algorithm=DM**20,nbprobe=1,probe=1', 'cxi_output=object_phase'],
                   'cxidm_alpha': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                   'algorithm=DM**10,dm_alpha=0.1,DM**10,nbprobe=1,probe=1',
                                   'dm_alpha=0'],
                   'cxipr': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                             'algorithm=DM**20,nbprobe=1,probe=1', 'cxi_output=probe'],
                   'cxi_mask_flat': ['pynx-ptycho-cxi', 'data=%s' % path, 'probe=gaussian,400e-9x400e-9',
                                     'algorithm=ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                     'verbose=10', 'save=all', 'saveplot', 'mask=mask.npz', 'flatfield=flatfield.npz'],
                   'cxi_mpi': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                               'data=%s' % path, 'probe=gaussian,400e-9x400e-9', 'mpi=split',
                               'algorithm=analysis,ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                               'verbose=10', 'save=all', 'saveplot'],
                   'cxi_mpi_scan': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                                    'data=%s' % path, 'probe=gaussian,400e-9x400e-9', 'mpi=split',
                                    'algorithm=analysis,ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                    'verbose=10', 'save=all', 'saveplot', 'scan=1,2'],
                   'cxi_mpi_no_rerun': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                                        'data=%s' % path, 'probe=gaussian,400e-9x400e-9', 'mpi=split',
                                        'algorithm=analysis,ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                        'verbose=10', 'save=all', 'saveplot', 'no_rerun'],
                   'cxi_mpi_scan_no_rerun': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                                             'data=%s' % path, 'probe=gaussian,400e-9x400e-9', 'mpi=split',
                                             'algorithm=analysis,ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                             'verbose=10', 'save=all', 'saveplot', 'scan=1,2', 'no_rerun'],
                   'simulation': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                  'algorithm=analysis,ML**20,AP**20,RAAR**10,'
                                  'DM**20,background=1,DM**20,nbprobe=1,probe=1',
                                  'verbose=10', 'save=all', 'saveplot', 'simul_background=0.05'],
                   'simulation_maxframe': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                           'maxframe=60'],
                   'simulation_modulo': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                         'moduloframe=2'],
                   'simulation_modulo1': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                          'moduloframe=2,1'],
                   'simulation_nf': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                     'algorithm=analysis,ML**20,AP**20,RAAR**10,'
                                     'DM**20,background=1,DM**20,nbprobe=1,probe=1',
                                     'verbose=10', 'save=all', 'saveplot', 'near_field'],
                   'simulation_nf_direct': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                            'algorithm=analysis,ML**20,AP**20,RAAR**10,DM**20,'
                                            'DM**20,nbprobe=1,probe=1', 'verbose=10', 'save=all', 'saveplot',
                                            'near_field', 'direct_beam'],
                   'simulation_nf_use_direct_pos': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                                    'algorithm=analysis,ML**20,AP**20,RAAR**10,DM**20,position=2,'
                                                    'DM**20,probe=1', 'verbose=10', 'save=all', 'saveplot',
                                                    'near_field', 'direct_beam', 'use_direct_beam'],
                   'simulation_pos': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                      'algorithm=analysis,ML**20,AP**20,RAAR**10,'
                                      'DM**20,position=2,DM**20,nbprobe=1,probe=1',
                                      'verbose=10', 'save=all', 'saveplot'],
                   'simulation_pos_direct': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                             'algorithm=analysis,ML**20,AP**20,RAAR**10,DM**20,position=2,'
                                             'DM**20,nbprobe=1,probe=1',
                                             'verbose=10', 'save=all', 'saveplot', 'direct_beam'],
                   'simulation_multi': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                        'algorithm=analysis,ML**20,AP**20,RAAR**10,'
                                        'DM**20,position=2,DM**20,nbprobe=1,probe=1',
                                        'verbose=10', 'multiscan_reuse_ptycho=ML**20,DM**20,probe=1',
                                        'scan=1,2'],
                   'simulation_multi_nf': ['pynx-ptycho-simulation', 'frame_nb=64', 'frame_size=128',
                                           'near_field', 'algorithm=analysis,ML**20,AP**20,RAAR**10,'
                                                         'DM**20,position=2,DM**20,nbprobe=1,probe=1',
                                           'verbose=10', 'multiscan_reuse_ptycho=ML**20,DM**20,probe=1',
                                           'scan=1,2'],
                   }
        # # Also test all options as saved in a *.ini file
        # for k in list(options.keys()):
        #     options[f"{k}_as_parameters_file"] = options[k]

        options.update({
            # We need to run those expected failures at the end
            # the failure triggers the deletion of the TemporaryDirectory (?)
            'cxi_overwrite_FAIL': ['pynx-ptycho-cxi', 'data=%s' % path,
                                   'probe=gaussian,400e-9x400e-9',
                                   'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                   'verbose=10', 'saveprefix=%s' % (path[:-4])],
            'cxi_FAIL': ['pynx-ptycho-cxi', 'data=this_file_does_not_exist.cxi',
                         'probe=gaussian,400e-9x400e-9',
                         'algorithm=ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                         'verbose=10', 'save=all', 'saveplot']
        })

        for k, args in options.items():
            # Only run one language for the old command-line argument style
            langs = ['opencl'] if exclude_cuda else ['cuda']
            args0 = args
            for lang in langs:
                # test as_parameters_file
                if "as_parameters_file" in k:
                    with tempfile.NamedTemporaryFile(mode='w', suffix=".config",
                                                     dir=self.tmp_dir, delete=False) as partxt:
                        partxt.write(f"# Converted parameters file from: {' '.join(args[1:])}")
                        for v in args[1:]:
                            if '=' in v:
                                # We don't really need the spaces in ' = ' but this is more realistic formatting
                                if 'algorithm' in v or 'multiscan' in v or 'probe=' in v or 'data=' in v:
                                    # parsing ignores quotes, but again more realistic test
                                    tmp = v.split('=')
                                    partxt.write(f"{tmp[0]} = '{tmp[1]}'\n")
                                else:
                                    partxt.write(f"{v.replace('=', ' = ')}\n")
                            else:
                                partxt.write(f"{v} = True\n")

                        args = args[:1] + [partxt.name]

                opt_liveplot = [[]]
                if liveplot:
                    opt_liveplot.append(['liveplot'])
                for opt_live in opt_liveplot:
                    if 'mpi' in k and (opt_live == 'liveplot' or not has_mpi
                                       or platform.system() in ['Windows', 'Darwin']):
                        continue
                    # print(k, lang, opt_live)
                    with self.subTest(k, config=' '.join(args0 + [lang] + opt_live),
                                      command=" ".join(args + opt_live)):
                        my_env = os.environ.copy()
                        my_env["PYNX_PU"] = lang
                        with subprocess.Popen(args + opt_live, stderr=subprocess.PIPE, stdout=subprocess.PIPE,
                                              cwd=self.tmp_dir, env=my_env) as p:
                            stdout, stderr = p.communicate(timeout=200)
                            res = p.returncode
                            # print("#" * 100)
                            # print(k, res)
                            # print(stdout.decode())
                            # print(stderr.decode())
                            if 'FAIL' not in k:
                                self.assertTrue(res == 0, msg=stderr.decode())
                            else:
                                # Expected failure
                                self.assertFalse(res == 0, msg=stderr.decode())

    def test_ptycho_runner(self):
        # Using argparse
        path = make_ptycho_data_cxi(dsize=160, nb_frame=40, nb_photons=1e9, dir_name=self.tmp_dir)
        options = {
            'help': ['pynx-ptycho-simulation', '--help'],
            'cxi': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                    '--algorithm', 'ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                    '--verbose', '10', '--save', 'all', '--saveplot'],
            'cxi_stack16': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                            '--algorithm', 'ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                            '--stack_size', '16'],
            'cxi_stack20': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                            '--algorithm', 'ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                            '--stack_size', '20'],
            'cxi_saveprefix1': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                                '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                '--verbose', '10', '--save', 'all', '--saveplot', '--saveprefix', 'test%02d_%02d'],
            'cxi_saveprefix2': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                                '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                '--verbose', '10', '--save', 'all', '--saveplot', '--saveprefix', 'test%02d',
                                '--defocus', '200e-6'],
            'cxi_saveprefix3': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                                '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                '--verbose', '10', '--save', 'all', '--saveplot', '--saveprefix', 'test',
                                '--defocus=-200e-6'],
            'cxi_2scan': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                          '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                          '--verbose', '10', '--save', 'all', '--scan', '1,1'],
            'cxi_no_rerun': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                             '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                             '--verbose', '10', '--save', 'all', '--no_rerun'],
            'cxi_2scan_no_rerun': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                                   '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                   '--verbose', '10', '--save', 'all', '--no_rerun', '--scan', '1,1'],
            'cxiobj': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                       '--algorithm', 'DM**20,nbprobe=1,probe=1', '--cxi_output', 'object'],
            'cxiobjph': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                         '--algorithm', 'DM**20,nbprobe=1,probe=1', '--cxi_output', 'object_phase'],
            'cxidm_alpha': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                            '--algorithm', 'DM**10,dm_alpha=0.1,DM**10,nbprobe=1,probe=1',
                            '--dm_alpha', '0'],
            'cxipr': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                      '--algorithm', 'DM**20,nbprobe=1,probe=1', '--cxi_output', 'probe'],
            'cxi_mask_flat': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                              '--algorithm', 'ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                              '--verbose', '10', '--save', 'all', '--saveplot', '--mask', 'mask.npz',
                              '--flatfield', 'flatfield.npz'],
            'cxi_mpi': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                        '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9', '--mpi', 'split',
                        '--algorithm', 'analysis,ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                        '--verbose', '10', '--save', 'all', '--saveplot'],
            'cxi_mpi_scan': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                             '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9', '--mpi', 'split',
                             '--algorithm', 'analysis,ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                             '--verbose', '10', '--save', 'all', '--saveplot', '--scan', '1,2'],
            'cxi_mpi_no_rerun': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                                 '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9', '--mpi', 'split',
                                 '--algorithm', 'analysis,ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                 '--verbose', '10', '--save', 'all', '--saveplot', '--no_rerun'],
            'cxi_mpi_scan_no_rerun': ['mpiexec', '-n', '2', 'pynx-ptycho-cxi',
                                      '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9', '--mpi', 'split',
                                      '--algorithm', 'analysis,ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                      '--verbose', '10', '--save', 'all', '--saveplot', '--scan', '1,2', '--no_rerun'],
            'simulation': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                           '--algorithm',
                           'analysis,ML**20,AP**20,RAAR**10,DM**20,background=1,DM**20,nbprobe=1,probe=1',
                           '--verbose', '10', '--save', 'all', '--saveplot', '--simul_background', '0.05'],
            'simulation_maxframe': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                    '--maxframe', '60'],
            'simulation_modulo': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                  '--moduloframe', '2'],
            'simulation_modulo1': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                   '--moduloframe', '2', '1'],
            'simulation_nf': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                              '--algorithm',
                              'analysis,ML**20,AP**20,RAAR**10,DM**20,background=1,DM**20,nbprobe=1,probe=1',
                              '--verbose', '10', '--save', 'all', '--saveplot', '--near_field'],
            'simulation_nf_direct': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                     '--algorithm',
                                     'analysis,ML**20,AP**20,RAAR**10,DM**20,DM**20,nbprobe=1,probe=1',
                                     '--verbose', '10', '--save', 'all', '--saveplot',
                                     '--near_field', '--direct_beam'],
            'simulation_nf_use_direct_pos': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                             '--algorithm',
                                             'analysis,ML**20,AP**20,RAAR**10,DM**20,position=2, DM**20,probe=1',
                                             '--verbose', '10', '--save', 'all', '--saveplot',
                                             '--near_field', '--direct_beam', '--use_direct_beam'],
            'simulation_pos': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                               '--algorithm',
                               'analysis,ML**20,AP**20,RAAR**10,DM**20,position=2,DM**20,nbprobe=1,probe=1',
                               '--verbose', '10', '--save', 'all', '--saveplot'],
            'simulation_pos_direct': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                      '--algorithm',
                                      'analysis,ML**20,AP**20,RAAR**10,DM**20,position=2,DM**20,nbprobe=1,probe=1',
                                      '--verbose', '10', '--save', 'all', '--saveplot', '--direct_beam'],
            'simulation_multi': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                 '--algorithm', 'analysis,ML**20,AP**20,RAAR**10,'
                                                'DM**20,position=2,DM**20,nbprobe=1,probe=1',
                                 '--verbose', '10', '--multiscan_reuse_ptycho', 'ML**20,DM**20,probe=1',
                                 '--scan', '1', '2'],
            'simulation_multi_nf': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                    '--near_field', '--algorithm',
                                    'analysis,ML**20,AP**20,RAAR**10,DM**20,position=2,DM**20,nbprobe=1,probe=1',
                                    '--verbose', '10', '--multiscan_reuse_ptycho', 'ML**20,DM**20,probe=1',
                                    '--scan', '1', '2'],
            'simulation_2scan': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                 '--scan', '1,2', '--algorithm',
                                 'ML**20,DM**20,probe=1',
                                 '--verbose', '10', '--save', 'all', '--saveplot'],
            'simulation_2scan_shareprobe': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                            '--scan', '1,2', '--share_probe', '--algorithm',
                                            'LoopScan(ML**20),LoopScan(DM**20),probe=1',
                                            '--verbose', '10', '--save', 'all', '--saveplot'],
            'simulation_2scan_shareprobe1': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                             '--scan', '1,2', '--share_probe', '1', '--algorithm',
                                             'LoopScan(ML**20),LoopScan(DM**20),probe=1',
                                             '--verbose', '10', '--save', 'all', '--saveplot'],
            'simulation_4scan_shareprobe2': ['pynx-ptycho-simulation', '--frame_nb', '64', '--frame_size', '128',
                                             '--scan', '1,2,3,4', '--share_probe', '2', '--algorithm',
                                             'LoopScan(ML**20),LoopScan(DM**20),probe=1',
                                             '--verbose', '10', '--save', 'all', '--saveplot'],
            'dark': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                              '--algorithm', 'ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                              '--verbose', '10', '--save', 'all', '--dark', 'dark.npz'],
            # dark_subtract is deprecated, still should test for it
            'dark_subtract': ['pynx-ptycho-cxi', '--data', '%s' % path, '--probe', 'gaussian,400e-9x400e-9',
                              '--algorithm', 'ML**20,RAAR**10,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                              '--verbose', '10', '--save', 'all', '--dark', 'dark.npz', '--dark_subtract'],
        }
        # Also test all options as saved in a *.ini file
        for k in list(options.keys()):
            options[f"{k}_as_parameters_file"] = options[k]

        options.update({
            # We need to run those expected failures at the end
            # the failure triggers the deletion of the TemporaryDirectory (?)
            'cxi_overwrite_FAIL': ['pynx-ptycho-cxi', '--data', '%s' % path,
                                   '--probe', 'gaussian,400e-9x400e-9',
                                   '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                                   '--verbose', '10', '--saveprefix', '%s' % (path[:-4])],
            'cxi_FAIL': ['pynx-ptycho-cxi', '--data', 'this_file_does_not_exist.cxi',
                         '--probe', 'gaussian,400e-9x400e-9',
                         '--algorithm', 'ML**20,DM**20,nbprobe=2,DM**20,nbprobe=1,probe=1',
                         '--verbose', '10', '--save', 'all', '--saveplot'],
        })

        for k, args in options.items():
            langs = []
            if not exclude_cuda:
                langs += ['cuda']
            if not exclude_opencl:
                langs += ['opencl']
            for lang in langs:
                my_env = os.environ.copy()
                my_env["PYNX_PU"] = lang
                args0 = args
                if "as_parameters_file" in k and (lang == 'cuda' or 'cuda' not in langs) and args[0] != 'mpiexec':
                    with tempfile.NamedTemporaryFile(mode='w', suffix=".config",
                                                     dir=self.tmp_dir, delete=False) as partxt:
                        partxt.write(f"# Converted parameters file from: {' '.join(args[1:])}")
                        nb = len(args)
                        for i in range(nb - 1):
                            v = args[i]
                            if v.startswith('--'):
                                # New argument ? Finish previous line first
                                if args[i - 1].startswith('--') and '=' not in args[i - 1]:
                                    partxt.write(f" = True")
                                partxt.write(f"\n")

                                # Handle new argument
                                v = v[2:]
                                if '=' in v:
                                    # We don't really need the spaces in ' = ' but this is more realistic formatting
                                    if 'algorithm' in v or 'multiscan' in v or 'probe' in v or 'data' in v:
                                        # parsing ignores quotes, but again more realistic test
                                        tmp = v.split('=')
                                        partxt.write(f"{tmp[0]} = '{tmp[1]}'")
                                    else:
                                        partxt.write(f"{v.replace('=', ' = ')}")
                                else:
                                    partxt.write(f"{v}")
                            else:
                                # value/option for an argument
                                if args[i - 1].startswith('--'):
                                    partxt.write(' = ')
                                else:
                                    partxt.write(',')
                                partxt.write(v)

                        partxt.write('\n')
                        args = args[:1] + [partxt.name]

                opt_liveplot = [[]]
                if liveplot:
                    opt_liveplot.append(['--liveplot'])
                for opt_live in opt_liveplot:
                    if 'mpi' in k and ('--liveplot' in opt_live or not has_mpi
                                       or platform.system() in ['Windows', 'Darwin']):
                        continue
                    # print(lang, " ".join(v + opt_live))
                    with self.subTest(k, config=' '.join(args0 + opt_live + [lang]), command=" ".join(args + opt_live)):
                        with subprocess.Popen(args + opt_live, stderr=subprocess.PIPE, stdout=subprocess.PIPE,
                                              cwd=self.tmp_dir, env=my_env) as p:
                            stdout, stderr = p.communicate(timeout=200)
                            res = p.returncode
                            # print("#" * 100)
                            # print(k, res)
                            # print(stdout.decode())
                            # print(stderr.decode())
                            if 'FAIL' not in k:
                                self.assertTrue(res == 0, msg=stderr.decode())
                            else:
                                # Expected failure
                                self.assertFalse(res == 0, msg=stderr.decode())


def suite():
    load_tests = unittest.defaultTestLoader.loadTestsFromTestCase
    test_suite = unittest.TestSuite([load_tests(TestPtychoRunner)])
    return test_suite


if __name__ == '__main__':
    # sys.stdout = io.StringIO()
    warnings.simplefilter('ignore')
    res = unittest.TextTestRunner(verbosity=2, descriptions=False).run(suite())
