#    Copyright 2023 ONERA - contact luis.bernardos@onera.fr
#
#    This file is part of MOLA.
#
#    MOLA is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    MOLA is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public License
#    along with MOLA.  If not, see <http://www.gnu.org/licenses/>.

import os
from typing import List
import copy
from fnmatch import fnmatch
from pathlib import Path
import numpy as np

from treelab import cgns
from mola import __MOLA_PATH__
import mola.naming_conventions as names
from mola.logging import mola_logger, MolaAssertionError, MolaException, CYAN, GREEN, RED, YELLOW, ENDC
from mola.cfd.preprocess.run_manager import run_manager
from mola.cfd.preprocess.write_cfd_files import write_cfd_files
from mola import server as SV
from . import Workflow

class WorkflowManager():
    '''
    This class allows to submit in sequential or parallel mode, maybe on a remote machine, 
    a bunch of workflows by applying variations from a base workflow.
    It is particularly useful to perform parametric studies.

    Parameters
    ----------

    workflow : :py:class:`Workflow` object
        Base workflow, used to generate variations around it. See the documentation 
        of :py:class:`WorkflowDispatcher` object for details.

    root_directory : str
        Target root directory from which the multiple workflows will be deployed.

    data_directory : str, optional
        Directory where the data files common to the workflows will be placed
        (refer to the documentation of :py:class:`WorkflowSender` for more information).
        It is relative to **root_directory**.

        By default it is set to ``'SHARED_DATA'``, so shared data files will be in 
        ``<root_directory>/SHARED_DATA/``.

    skip_if_exists: bool, optional
        If True, it will simply pass over directories that were already created in a 
        previous usage of ``WorkflowManager``. 
        
        .. hint:: It is useful to add new operating points to a parametric study previously submitted.
        
        If False, an error is raise is a directory already exists.

        Default value is False.

    Example
    -------

    .. code-block:: python

        from mola.workflow import Workflow
        from mola.workflow import manager as WM 

        workflow = Workflow(...)
        manager = WM.WorkflowManager(workflow, '/the/path/to/my/parametric/study/')

        for BCWall in ['WallViscous', 'WallInviscid']:
            manager.new_job(BCWall)  # There will be two job submitted, one for 'WallViscous' and one for 'WallInviscid'
            for velocity in [50., 20., 80.]:
                manager.add_variations(
                    [
                        ('RunManagement|JobName', BCWall),
                        ('RunManagement|RunDirectory', f'Velocity_{velocity}'),
                        ('Flow|Velocity', velocity),
                        ('BoundaryConditions|Family=Wall|Type', BCWall),
                    ], 
                )
        
        manager.prepare()  # at the end of this command, all the directories are deployed and job files are ready to be submitted
        manager.submit()  # submit all the jobs '{names.FILE_JOB_SEQUENCE}' generated by the previous line

    ``add_variations`` takes a list of parameters to make vary compared to the base workflow.
    Each variation is a tuple of two elements:

    #. a path (:py:class:`str`) to the parameter to make vary in the base workflow. 
       The separator is a pipe '|'. In the example given above:

       * ``'RunManagement|JobName'`` selects ``workflow.RunManagement['JobName']``.

       * ``'BoundaryConditions|Family=Wall|Type'`` selects ``workflow.BoundaryConditions[<index>]['Type']``.
         Because ``workflow.BoundaryConditions`` is a :py:class:`list`, ``Family=Wall`` is a mean to select 
         the boundary condition `BC` in this list (remember `BC` is a :py:class:`dict`) such as ``BC['Family']=='Wall'``. 
         ``workflow.BoundaryConditions[<index>]`` is that ``BC``, and finally the value corresponding to the key 
         ``'Type'`` is selected.

    #. the new value to apply to this selected parameter.

    '''

    def __init__(self, 
                 arg, 
                 root_directory='.', 
                 data_directory='SHARED_DATA', 
                 skip_if_exists=True,
                 manager_file_path=names.FILE_WORKFLOW_MANAGER,
                 ):
        
        if isinstance(arg, str) and arg.endswith('.cgns'):
            # init reading a file previously written by WorkflowManager.write
            self.read(arg)
        else:
            # arg is the base workflow of the manager object
            assert isinstance(arg, Workflow)
            self.root_directory = root_directory
            self.data_directory = data_directory
            self.skip_if_exists = skip_if_exists
            self.base_workflow = arg
            self.dispatcher = WorkflowDispatcher(self.base_workflow)
            self.machine = None
            self.sequential_managers = None
            self.manager_file_path = manager_file_path
    
    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        else:
            return (
                self.root_directory == other.root_directory \
                and self.data_directory == other.data_directory \
                and self.skip_if_exists == other.skip_if_exists \
                and self.machine == other.machine \
                and self.base_workflow == other.base_workflow \
                and self.dispatcher == other.dispatcher \
            )

    def new_job(self, directory):
        self.dispatcher.new_job(directory)
    
    def add_variations(self, variations, initialize_from_previous=True):
        self.dispatcher.add_variations(variations, initialize_from_previous)

    def prepare(self):
        if self.sequential_managers is None:
            self._init_sequential_managers()
        SV.makedirs_remote(self.root_directory, machine=self.machine)
        for sequence_of_workflows in self.sequential_managers:
            sequence_of_workflows.prepare()
        
        self.write()

    def submit(self):
        for sequence_of_workflows in self.sequential_managers:
            sequence_of_workflows.submit()

    def _init_sequential_managers(self):
        if self.sequential_managers is not None:
            raise MolaAssertionError('Attribute sequential_managers was already initialized !')
        self.sequential_managers = []
        for i, sequence in enumerate(self.dispatcher.table_of_workflows):
            root_absolute_directory = os.path.join(self.root_directory, self.dispatcher.root_directories[i])
            data_absolute_directory = os.path.join(self.root_directory, self.data_directory)

            sequential_manager = WorkflowSequentialManager(
                sequence, 
                root_directory=root_absolute_directory, 
                data_directory=data_absolute_directory, 
                skip_if_exists=self.skip_if_exists
                )
            self.sequential_managers.append(sequential_manager)

        self.machine = sequential_manager.machine
    
    def get_run_directories(self):
        return [workflow.RunManagement['RunDirectory'] 
                for sequential_manager in self.sequential_managers 
                for workflow in sequential_manager.workflows]
    
    def write(self):
        if self.sequential_managers is None:
            self._init_sequential_managers()

        tree = cgns.Tree()
        d = self.dispatcher.convert_to_dict()
        tree.setParameters(
            ContainerName = 'WorkflowManager',
            Machine = self.machine,
            RootDirectory = self.root_directory,
            DataDirectory = self.data_directory,
            SkipIfExists = self.skip_if_exists, 
            BaseWorkflow = self.base_workflow.convert_to_dict(), 
            RunDirectories = self.get_run_directories(),
            WorkflowDispatcher = d,
            )
        tree.save(self.manager_file_path)

    def read(self, manager_filepath):
        tree = cgns.load(manager_filepath)
        WorkflowDispatcher = cgns.load_from_path(manager_filepath, 'WorkflowManager/WorkflowDispatcher')
        parameters = tree.getParameters('WorkflowManager', transform_numpy_scalars=True)
        self.root_directory = parameters['RootDirectory']
        self.data_directory = parameters['DataDirectory']
        self.skip_if_exists = parameters['SkipIfExists']
        self.dispatcher = read_workflow_dispatcher_from_tree(WorkflowDispatcher)
        self.base_workflow = self.dispatcher.base_workflow
        self.machine = parameters['Machine']
        self.run_directories = parameters['RunDirectories']

    def get_synchronized_directories(self, filenames: list=[], update_from_remote_machine: bool=True): 
        if not hasattr(self, 'run_directories'):
            self.run_directories = self.get_run_directories()

        def _get_status(directory):
            for status in names.STATUS_FILES:
                status_file = Path(directory) / names.FILE_JOB_COMPLETED
                if status_file.is_file():
                    return status
            return 'not ended'

        def _get_completed_run_directories(local_directories):
            completed_run_directories = []
            for local_dir in local_directories:

                status = _get_status(local_dir)
                if status == names.FILE_JOB_COMPLETED:
                    mola_logger.info(f'{local_dir} -> {GREEN}{names.FILE_JOB_COMPLETED}{ENDC}')
                    completed_run_directories.append(local_dir)

                elif status == names.FILE_JOB_FAILED:
                    mola_logger.info(f'{local_dir} -> {RED}{names.FILE_JOB_FAILED}{ENDC}')
                
                else:
                    mola_logger.info(f'{local_dir} -> {YELLOW}{status}{ENDC}')
                    
            return completed_run_directories

        if not SV.run_on_localhost(machine=self.machine, run_directory=self.run_directories[0]):
            local_directories = [
                os.path.relpath(run_directory, self.root_directory) 
                for run_directory in self.run_directories
                ]
            if update_from_remote_machine:
                # copy locally all the required files from each directory
                included_files = set(
                    [names.DIRECTORY_OUTPUT, f'*/{names.FILE_OUTPUT_1D}']
                    + names.STATUS_FILES 
                    + filenames
                )
                SV.get_all_cases_from_workflow_manager(
                    included_files=included_files, 
                    excluded_files=set('*')
                )

        else:
            local_directories = [
                os.path.relpath(run_directory, '.') 
                for run_directory in self.run_directories
                ]
        
        completed_run_directories = _get_completed_run_directories(local_directories)
        
        return completed_run_directories
        

    def gather_signals(self, queries, filename=None, keep_last_point=False, update_from_remote_machine=True):
        """
        Need for an extraction method from queries (path, metadata)
        get scalars from signals.cgns and plot a curve
        plot several curves (one per case) on the same plot
        get a surface from each instant
        """
        if not filename:
            filename = os.path.join(names.DIRECTORY_OUTPUT, names.FILE_OUTPUT_1D)
        if not isinstance(queries, list):
            queries = [queries]

        def _extract(signals: cgns.Tree, queries: list) -> dict:
            # This extraction function must be developped
            # The query should be a path, a part of a path, or parameters to check metadata of extractions in names.CGNS_NODE_EXTRACTION_LOG
            results = dict()
            for query in queries:
                path = query
                try:
                    node = signals.getAtPath(path)
                except:
                    raise MolaException(f'Cannot find node at {path} in {filename}')
                if node:
                    results[node.name()] = node.value()
            return results
        
        local_run_directories = self.get_synchronized_directories([filename], update_from_remote_machine=update_from_remote_machine)

        all_data = dict()
        user_dir = os.getcwd()
        for local_dir in local_run_directories:
            os.chdir(local_dir)
            signals = cgns.load(filename)
            data = _extract(signals, queries)
            if keep_last_point:
                for v in data: 
                    data[v] = data[v][-1]
            all_data[local_dir] = data
            os.chdir(user_dir)
        
        return all_data

    @staticmethod
    def rearange_signals(signals):
        '''
        Input: 
        
        >>> signals = {'<job_dir>/<case_dir1>': {'Thrust': 10., 'Power': 100.}, '<job_dir>/<case_dir2>': {'Thrust': 12., 'Power': 110.}}

        Output:

        >>> {'<job_dir>': {'Thrust': np.array([10., 12.]), 'Power': np.array([100., 110.]), 'case': ['<case_dir1>', '<case_dir2>']}}

        '''
        ordered_signals = dict()
        for path, data in signals.items():
            path = Path(path)
            dir_multi = path.parent.name
            dir_case = path.name

            if dir_multi not in ordered_signals:
                ordered_signals[dir_multi] = dict((key, [value]) for key, value in data.items())
                ordered_signals[dir_multi]['case'] = [dir_case]
            else:
                for key, value in data.items():
                    ordered_signals[dir_multi][key].append(value)
                ordered_signals[dir_multi]['case'].append(dir_case)

        for dir_multi, perfo_on_iso in ordered_signals.items():
            for key, a_list in perfo_on_iso.items():
                if key == 'case': 
                    # case list is a list of str
                    continue
                ordered_signals[dir_multi][key] = np.array(a_list)

        # TODO need to sort cases ? 

        return ordered_signals


class WorkflowDispatcher():
    '''
    This class allows creating "variations" from a base workflow. 
    It may be used to prepare a parametric study on one or several parameters.

    To create a WorkflowDispatcher object, just give a base workflow:
     
    .. code-block:: python

        from mola.workflow import Workflow
        from mola.workflow import manager as WM 

        workflow = Workflow(...)
        dispatcher = WM.WorkflowDispatcher(workflow)

    Once instanciated, the user may use the method ``new_job`` and 
    ``add_variations`` to generate workflows derivated from the base workflow.
    (see the example below).

    .. hint::

        This class is intended to work with :py:class:`WorkflowManager`,
        in order to submit the workflows generated by :py:class:`WorkflowDispatcher`.

    '''

    def __init__(self, workflow):
        self.base_workflow = workflow
        self.table_of_workflows = []
        self.workflows_in_current_job = None
        self.root_directories = []
    
    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        else:
            same_table_of_workflows = True
            for workflows1, workflows2 in zip(self.table_of_workflows, other.table_of_workflows):
                for w, w2 in zip(workflows1, workflows2):
                    if w != w2:
                        same_table_of_workflows = False
                        break
            return (
                self.root_directories == other.root_directories \
                and self.base_workflow == other.base_workflow
                and same_table_of_workflows
            )

    def new_job(self, directory):
        self.workflows_in_current_job = []
        self.table_of_workflows.append(self.workflows_in_current_job)
        self.root_directories.append(directory)
    
    def add_variations(self, variations, initialize_from_previous=True):
        self._check_new_job_is_declared()
        new_workflow = copy.deepcopy(self.base_workflow)

        if initialize_from_previous and len(self.workflows_in_current_job) > 0:
            self._add_variations_to_initialize_from_previous(variations)

        for request, value in variations:
            path_in_workflow = self.request_to_paths(request)
            set_value_on_leaf(new_workflow, path_in_workflow, value)

        self.workflows_in_current_job.append(new_workflow)
    
    def _check_new_job_is_declared(self):
        if self.workflows_in_current_job is None:
            raise MolaAssertionError('Before calling `add_variations`, `new_job` must be called first to declare directory.')
    
    def _add_variations_to_initialize_from_previous(self, variations):
        try:
            previous_workflow = self.workflows_in_current_job[-1]
            init_variations = [
                ('Initialization|Method', 'from_previous'), 
            ]
            variations += init_variations
        except IndexError:
            mola_logger.warning(f'Cannot initialize the first case of a sequence ({self.root_directories[-1]}) from a previous case')

    @staticmethod
    def request_to_paths(request):
        if isinstance(request, str):
            request = request.split('|')
        return request

    def get_directories_in_current_job(self, worflows_sequence=None):
        if worflows_sequence is None:
            worflows_sequence = self.workflows_in_current_job
        return [workflow.RunManagement['RunDirectory'] for workflow in worflows_sequence]

    def get_local_directories(self):
        directories = []
        for worflows_sequence in self.table_of_workflows:
            job_directories = self.get_directories_in_current_job(worflows_sequence)
            directories.append(job_directories)
        return directories

    def get_directories(self):
        directories = []
        for root, worflows_sequence in zip(self.root_directories, self.table_of_workflows):
            job_directories = self.get_directories_in_current_job(worflows_sequence)
            directories.extend([os.path.join(root, d) for d in job_directories])
        return directories

    def convert_to_dict(self):
        TableOfWorkflows = dict()
        for i, workflows_in_job in enumerate(self.table_of_workflows):
            TableOfWorkflows[f'_list_.{i}'] = [w.convert_to_dict() for w in workflows_in_job]

        params = dict(
            RootDirectories = self.root_directories,
            BaseWorkflow = self.base_workflow.convert_to_dict(),
            TableOfWorkflows = TableOfWorkflows,
        )
        return params
    

def read_workflow_dispatcher_from_tree(tree):
    from . import AVAILABLE_WORKFLOWS

    def _build_workflow_tree(t):
        workflow_tree = cgns.Tree()
        t = t.copy()
        t.setName(names.CONTAINER_WORKFLOW_PARAMETERS)
        workflow_tree.addChild(t)
        return workflow_tree
        
    BaseWorkflow = tree.get(Name='BaseWorkflow', Depth=1)
    workflow_name = BaseWorkflow.get(Name='Name', Depth=1).value()
    PreviouslyUsedWorkflow = AVAILABLE_WORKFLOWS.get(workflow_name)
    base_workflow = PreviouslyUsedWorkflow(tree=_build_workflow_tree(BaseWorkflow))

    dispatcher = WorkflowDispatcher(base_workflow)

    dispatcher.workflows_in_current_job = None
    dispatcher.root_directories = tree.get(Name='RootDirectories', Depth=1).value()
    if isinstance(dispatcher.root_directories, str):
        dispatcher.root_directories = [dispatcher.root_directories]

    dispatcher.table_of_workflows = []
    TableOfWorkflows = tree.get(Name='TableOfWorkflows', Depth=1)
    for workflows_in_job in TableOfWorkflows.children():
        workflows = []
        for workflow in workflows_in_job.children():
            workflows.append(PreviouslyUsedWorkflow(tree=_build_workflow_tree(workflow)))
        dispatcher.table_of_workflows.append(workflows)

    return dispatcher

class WorkflowSequentialManager():
    '''
    This class is instanciated by :py:meth:`WorkflowManager._init_sequential_managers`.
    It is **not intented to be called directly by user**.

    It handles the creation of one "sequence" of cases (generated from **workflows**), 
    run sequentially with the job file :mola_name:`FILE_JOB_SEQUENCE`.
    '''

    _sequential_job_filename = names.FILE_JOB_SEQUENCE

    def __init__(self, workflows: List[Workflow], root_directory, data_directory, skip_if_exists=False):

        self.workflows = workflows
        self.root_directory = root_directory
        self.data_directory = data_directory
        self.skip_if_exists = skip_if_exists

        # Following parameter will be set to True if at least one new case is created.
        # The only case when it will keep value False is if all cases have already been 
        # created by a previous WorkflowManager.
        self.at_least_one_new_case = False  


        self._check_structure_of_workflows()
        self._check_and_set_local_paths()
        self._check_all_workflows_have_different_working_directories()
        self._set_machine()

    def _check_structure_of_workflows(self):
        assert isinstance(self.workflows, list)
        for workflow in self.workflows:
            assert isinstance(workflow, Workflow)

    def _check_all_workflows_have_different_working_directories(self):
        assert len(set(self.cases_local_paths)) == len(self.workflows), 'Several workflows have the same working directory'

    def _check_and_set_local_paths(self):
        self.cases_local_paths = []
        for workflow in self.workflows:
            path = workflow.RunManagement['RunDirectory']
            if not path.startswith(os.path.sep):
                # local path
                self.cases_local_paths.append(path)
                workflow.RunManagement['RunDirectory'] = os.path.join(self.root_directory, path)
            else:
                # absolute path
                if not path.startswith(self.root_directory):
                    raise MolaAssertionError(
                        f'RunDirectory {path} is not consistent with the root path {self.root_directory}.'
                        )
                self.cases_local_paths.append(path.replace(self.root_directory, ''))

    def _set_machine(self):
        first_workflow = self.workflows[0] 
        RunManagement = copy.deepcopy(first_workflow.RunManagement)
        scheduler_options = run_manager.set_default(RunManagement, check_run_dir=not self.skip_if_exists)
        self.machine = RunManagement['Machine']
        self.manager = RunManagement['Scheduler']
        self.job_text = write_cfd_files.get_job_text(first_workflow.Solver, RunManagement, scheduler_options)

    def prepare(self):
        SV.makedirs_remote(self.root_directory, machine=self.machine)
        previous_workflow = None

        for workflow in self.workflows:
            mola_logger.info(f"\n{CYAN}  > preparing {workflow.RunManagement['RunDirectory']}...{ENDC}")
            if self.skip_if_exists and SV.is_directory(workflow.RunManagement['RunDirectory'], self.machine):
                mola_logger.user_warning(f"Skip directory {workflow.RunManagement['RunDirectory']} that already exists")
                previous_workflow = workflow
                continue
            else:
                if not self.at_least_one_new_case:
                    # Stop process if a previous job is still running
                    self.check_job_is_not_submitted_or_running()
                self.at_least_one_new_case = True

            # If needed, update initialization field from previous case
            if workflow.Initialization.get('Method') == 'from_previous':
                workflow.Initialization['Method'] = 'copy'
                try:
                    previous_case_path = Path(previous_workflow.RunManagement['RunDirectory'])
                except IndexError:
                    mola_logger.warning(f'Cannot initialize the first case of a sequence ({self.root_directories[-1]}) from a previous case')
                workflow.Initialization['Source'] = f'../{previous_case_path.name}/{names.FILE_INPUT_SOLVER}'
                copy_options = dict(excluded_attributes=['Extractions', 'Initialization'])
            else:
                copy_options = None

            workflow.write_tree_remote(data_directory=self.data_directory, copy_options=copy_options)
            previous_workflow = workflow

        if self.at_least_one_new_case:
            self.write_sequence_job()  

    def check_job_is_not_submitted_or_running(self):
        user = None # TODO allow different user
        job_path = os.path.join(self.root_directory, self._sequential_job_filename)
        if SV.is_file(job_path, machine=self.machine, user=user):
            command = f"grep 'job-name' {job_path} | awk -F'=' '{{print $2}}'"
            output = SV.submit_command(command, self.machine, user=user)
            job_name = output.split('\n')[-2]
            assert len(job_name) > 0

            if SV.job_is_submitted_or_running(job_name, self.machine):
                raise MolaAssertionError((
                    f'The job at path {job_path} is already running or submitted, '
                    'hence you cannot modify add new subdirectories or submit this job again. '
                    'Wait for this job to finish or cancel it.'
                ))

    def write_sequence_job(self):
        paths_in_bash = '"{}"'.format(' '.join(self.cases_local_paths))
        loop_on_cases = self._build_loop_on_cases(paths_in_bash, self._sequential_job_filename)
        SV.save_file_maybe_remote(self._sequential_job_filename, self.job_text + loop_on_cases, 
                                  self.root_directory, machine=self.machine, force_copy=True)

    def submit(self):
        if self.at_least_one_new_case:
            mola_logger.info(f'  > submission of job sequence in {self.root_directory}')
            if self.manager == 'SLURM':
                command = f"cd {self.root_directory}; sbatch {self._sequential_job_filename}"
            else:
                command = f"cd {self.root_directory}; ./{self._sequential_job_filename}"

            SV.submit_command(command, self.machine)
        else:
            mola_logger.user_warning(f'  > job sequence in {self.root_directory} was not submitted, because all cases were already existing.')

    @staticmethod
    def _build_loop_on_cases(sequence_of_paths, sequential_job_filename):

        loop_on_cases = f'''
SECONDS=0
SEQUENCE_OF_PATHS={sequence_of_paths}
MAX_NUMBER_OF_ATTEMPTS=10

for case in $SEQUENCE_OF_PATHS; do

    echo "entering $case"
    cd $case

    if [ -f "{names.FILE_ERROR_PREPARING_WORKFLOW}" ] || [ -f "{names.FILE_JOB_FAILED}" ]; then
        echo "this run has failed, end job."
        exit 0
    elif [ ! -f "{names.FILE_INPUT_SOLVER}" ]; then
        mola_prepare {names.FILE_INPUT_WORKFLOW}
    fi

    NUMBER_OF_ATTEMPTS=0

    while [ "$NUMBER_OF_ATTEMPTS" -lt "$MAX_NUMBER_OF_ATTEMPTS" ]; do

        NUMBER_OF_ATTEMPTS=$((NUMBER_OF_ATTEMPTS+1))

        export ELAPSED_TIME_MOLA_MANAGER=$SECONDS
        echo "compute case $case at $ELAPSED_TIME_MOLA_MANAGER s"
        ./{names.FILE_JOB}

        if [ -f "{names.FILE_NEWJOB_REQUIRED}" ]; then
            rm {names.FILE_NEWJOB_REQUIRED}
            echo "LAUNCHING THIS JOB AGAIN"
            cd ..
            sbatch {sequential_job_filename} --dependency=singleton
            exit 0
        elif [ -f "{names.FILE_JOB_COMPLETED}" ]; then
            # exit the while loop
            break
        elif [ -f "{names.FILE_ERROR_PREPARING_WORKFLOW}" ] || [ -f "{names.FILE_JOB_FAILED}" ]; then
            echo "this run has failed, end job."
            exit 0
        else
            # there must be an error that was not raised -> exit the while loop
            echo "unraised error, end job."
            exit 0
        fi
    done

    cd ..

done
'''
        return loop_on_cases


class WorkflowSender():
    '''
    This class allow to save in a file :mola_name:`FILE_INPUT_WORKFLOW` 
    the parametrization of a workflow that has been just instanciated, 
    without calling its method ``prepare``, 
    and to send this file to ``workflow.RunManagement['RunDirectory']``, 
    eventually on a remote machine designated by ``workflow.RunManagement['RunDirectory']``.

    Notice than the workflow may have paths to files in its parameters, 
    for instance ``workflow.RawMeshComponents[<one_component>]['Source'] = 'mesh.cgns'``.
    In that case, these files are also copied to a distant repository given by **data_directory**. 
    If not given, the same directory ``self.workflow.RunManagement['RunDirectory']`` will be taken.
    By default, files that are copied are found following the pattern ``*.cgns``.

    Paths in the workflow parameters are then automatically updated to be consistent with
    the new localization of theses files.

    .. note:: 

        The most direct application of this class is its wrapping in the method 
        :meth:`~mola.workflow.Workflow.write_tree_remote` of
        the class :py:class:`~mola.workflow.Workflow`. 

    .. note:: 

        If you need to change the patterns used to found files to copy;
        the argument ``patterns_to_copy_files`` may be pass to :py:meth:`WorkflowSender.__init__`
        Remember that by default, ``patterns_to_copy_files=['*.cgns']``.

    Example
    -------

    .. code-block::python

        from mola.workflow import Workflow
        from mola.workflow import manager as WM 

        workflow = Workflow(...)
        sender = WM.WorkflowSender(workflow, data_directory=data_directory)
        sender.apply()

    '''

    _workflow_filename = names.FILE_INPUT_WORKFLOW

    def __init__(self, workflow, data_directory=None, copy_options=None): 
        self.workflow = copy.deepcopy(workflow)
        self.run_directory = copy.deepcopy(self.workflow.RunManagement['RunDirectory'])

        if data_directory is None:
            self.data_directory = self.run_directory
        else:
            self.data_directory = data_directory

        self.copy_options = dict(
            patterns = ['*.cgns'],
            operation = self._get_adapted_path,
            excluded_attributes = ['Extractions'],
        )
        if isinstance(copy_options, dict):
            self.copy_options.update(copy_options)
        
        run_manager.set_default_machine(self.workflow.RunManagement)
        self.machine = self.workflow.RunManagement['Machine']
        
    def apply(self):
        
        destination = os.path.join(self.run_directory, self._workflow_filename)
        self.workflow.RunManagement['RunDirectory'] = '.'
        
        run_on_localhost = SV.run_on_localhost(self.machine, self.run_directory)
        if run_on_localhost:    
            self.workflow.set_workflow_parameters_in_tree()    
            SV.makedirs_remote(self.run_directory, machine='localhost')    
            self.workflow.write_tree(filename=destination)
            
        else:
            self._copy_files_and_update_paths_in_workflow(**self.copy_options)

            self.workflow.set_workflow_parameters_in_tree()
            self.workflow.write_tree(filename=self._workflow_filename)
            SV.move_remote(
                source_path=self._workflow_filename, 
                source_machine='localhost', 
                destination_path=destination, 
                destination_machine=self.machine,
                )

    def _copy_files_and_update_paths_in_workflow(self, patterns, operation, excluded_attributes):
        files2copy = find_matching_leaves(self.workflow, patterns, operation, excluded_attributes)
        files2copy = [filename for filename in files2copy if not filename.endswith(names.FILE_INPUT_SOLVER)]
        
        for filename in files2copy:
            self._copy_file_to_data_directory(filename)
        
    def _copy_file_to_data_directory(self, path):
        machine = self.machine
        filename = path.split(os.path.sep)[-1]
        new_filename = os.path.join(self.data_directory, filename)

        if not SV.is_existing_path(new_filename, machine=machine):
            mola_logger.info(f'copy {path} to {self.data_directory}')
            SV.makedirs_remote(self.data_directory, machine=machine)
            SV.copy_remote(
                source_path=path,
                destination_path=new_filename,
                destination_machine=machine
            )
    
    def _get_adapted_path(self, path):
        filename = path.split(os.path.sep)[-1]
        dir_path = os.path.relpath(self.data_directory, start=self.run_directory)
        return os.path.join(dir_path, filename)
    

def set_value_on_leaf(tree, path: str, value)-> None:
    '''
    Set a new value for the leaf in Workflow parameters selected 
    by its path.

    Parameters
    ----------
    tree : Workflow or dict or list
        Workflow or a sub-tree of Workflow parameters.
    path : str
        Path of the leaf to modify in tree, using the pipe '|' as separator.
        Examples: 
        * `'RunManagement|JobName'`
        * `'BoundaryConditions|Family=Wall|Type'``

    value : 
        the new value to set for the selected leaf.

    '''
    current_path = path[0]

    if isinstance(tree, Workflow):
        if len(path) == 1:
            setattr(tree, current_path, value)
        else:
            attr = getattr(tree, current_path)
            set_value_on_leaf(attr, path[1:], value)

    elif isinstance(tree, dict):
        if len(path) == 1:
            tree[current_path] = value
        else:
            set_value_on_leaf(tree[current_path], path[1:], value)
    
    elif isinstance(tree, (list, tuple)):
        if not len(path) > 1:
            raise MolaAssertionError(
                f'Path cannot end with a list (current path is {path}). '
                'You must add the key of the selected dictionary to set.'
                )
        assert current_path.count('=') == 1
        # search the good element 
        name, search_value = current_path.split('=')
        found = False
        for dico in tree:
            if name in dico and str(dico[name]) == search_value:
                found = True
                break
        assert found
        set_value_on_leaf(dico, path[1:], value)
    
    else:
        raise MolaException(f'Unknown type: {type(tree)}')
          
def get_value_on_leaf(tree, path: str):
    # for now, it used only in tests
    try:
        current_path = path[0]
    except IndexError:
        return tree
    
    if isinstance(tree, Workflow):
        return get_value_on_leaf(getattr(tree, current_path), path[1:])

    elif isinstance(tree, dict):
        return get_value_on_leaf(tree[current_path], path[1:])
    
    elif isinstance(tree, (list, tuple)):
        assert current_path.count('=') == 1
        # search the good element 
        name, search_value = current_path.split('=')
        found = False
        for dico in tree:
            if name in dico and str(dico[name]) == search_value:
                found = True
                break
        assert found
        return get_value_on_leaf(dico, path[1:])
    
    else:
        raise MolaException(f'Unknown type: {type(tree)}')

def find_matching_leaves(tree, patterns: list, operation=None, excluded_attributes=[]) -> list:
    '''
    Find leaves in Workflow parameters matching one of the given patterns, 
    optionaly apply (in-place) an operation on them,
    and return the list of values of matching leaves.

    Parameters
    ----------
    tree : Workflow or dict or list
        Workflow or a sub-tree of Workflow parameters.

    patterns : :py:class:`list` of :py:class:`str`
        patterns used to find leaves in **tree**. 
        For example, to select CGNS files used by the Workflow:

        >>> patterns = ['*.cgns']

    operation : function, optional
        If not None, this function is applied to all matching leaves.
        It must return a variable (same type as the input).
    
    excluded_attributes : :py:class:`list` of :py:class:`str`
        Names of attributes of Workflow to exclude of the search.

    Returns
    -------
    list
        values of leaves matching patterns.
    '''
    matches = []

    if isinstance(tree, Workflow):
        for attr_name, attr_value in tree.__dict__.items():
            if attr_name.startswith('_') or attr_name in excluded_attributes:
                continue
            matches.extend(find_matching_leaves(attr_value, patterns, operation))

    elif isinstance(tree, dict):
        for key, value in tree.items():
            if isinstance(value, list):
                for item in value:
                    matches.extend(find_matching_leaves(item, patterns, operation))
            elif isinstance(value, dict):
                matches.extend(find_matching_leaves(value, patterns, operation))
            else:
                for pattern in patterns:
                    if fnmatch(str(value), pattern):
                        matches.append(value)
                        if operation:
                            tree[key] = operation(value)

    elif isinstance(tree, list):
        for item in tree:
            matches.extend(find_matching_leaves(item, patterns, operation))

    return matches


