fedbiomed.researcher.experiment

@exp_exceptions
def aggregated_params(self) -> dict:
    """Retrieves all aggregated parameters of each round of training

    Returns:
        Dictionary of aggregated parameters keys stand for each round of training
    """

    return self._aggregated_params

@exp_exceptions
def aggregator(self) -> Aggregator:
    """ Retrieves aggregator class that will be used for aggregating model parameters.

    To set or update aggregator: [`set_aggregator`][fedbiomed.researcher.experiment.Experiment.set_aggregator].

    Returns:
        A class or an object that is an instance of [Aggregator][fedbiomed.researcher.aggregators.Aggregator]

    """
    return self._aggregator

@exp_exceptions
def breakpoint(self) -> None:
    """
    Saves breakpoint with the state of the training at a current round. The following Experiment attributes will
    be saved:
      - round_current
      - round_limit
      - tags
      - experimentation_folder
      - aggregator
      - node_selection_strategy
      - training_data
      - training_args
      - model_args
      - training_plan_path
      - training_plan_class
      - aggregated_params
      - job (attributes returned by the Job, aka job state)
      - secagg

    Raises:
        FedbiomedExperimentError: experiment not fully defined, experiment did not run any round yet, or error when
            saving breakpoint
    """
    # at this point, we run the constructor so all object variables are defined

    # check pre-requisistes for saving a breakpoint
    #
    # need to have run at least 1 round to save a breakpoint
    if self._round_current < 1:
        msg = ErrorNumbers.FB413.value + \
            ' - need to run at least 1 before saving a breakpoint'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    elif self._fds is None:
        msg = ErrorNumbers.FB413.value + \
            ' - need to define `training_data` for saving a breakpoint'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    elif self._node_selection_strategy is None:
        msg = ErrorNumbers.FB413.value + \
            ' - need to define `strategy` for saving a breakpoint'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    elif self._job is None:
        msg = ErrorNumbers.FB413.value + \
            ' - need to define `job` for saving a breakpoint'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

        # conditions are met, save breakpoint
    breakpoint_path, breakpoint_file_name = \
        choose_bkpt_file(self._experimentation_folder, self._round_current - 1)

    state = {
        'training_data': self._fds.data(),
        'training_args': self._training_args.dict(),
        'model_args': self._model_args,
        'training_plan_path': self._job.training_plan_file,  # only in Job we always model saved to a file
        # with current version
        'training_plan_class': self._job.training_plan_name,  # not always available properly
        # formatted in Experiment with current version
        'round_current': self._round_current,
        'round_limit': self._round_limit,
        'experimentation_folder': self._experimentation_folder,
        'aggregator': self._aggregator.save_state(breakpoint_path, global_model=self._global_model),  # aggregator state
        'node_selection_strategy': self._node_selection_strategy.save_state(),
        # strategy state
        'tags': self._tags,
        'aggregated_params': self._save_aggregated_params(
            self._aggregated_params, breakpoint_path),
        'job': self._job.save_state(breakpoint_path),  # job state
        'secagg': self._secagg.save_state()
    }

    # rewrite paths in breakpoint : use the links in breakpoint directory
    state['training_plan_path'] = create_unique_link(
        breakpoint_path,
        # - Need a file with a restricted characters set in name to be able to import as module
        'model_' + str("{:04d}".format(self._round_current - 1)), '.py',
        # - Prefer relative path, eg for using experiment result after
        # experiment in a different tree
        os.path.join('..', os.path.basename(state["training_plan_path"]))
    )

    # save state into a json file.

    breakpoint_file_path = os.path.join(breakpoint_path, breakpoint_file_name)
    try:
        with open(breakpoint_file_path, 'w') as bkpt:
            json.dump(state, bkpt)
        logger.info(f"breakpoint for round {self._round_current - 1} saved at " +
                    os.path.dirname(breakpoint_file_path))
    except (OSError, ValueError, TypeError, RecursionError) as e:
        # - OSError: heuristic for catching open() and write() errors
        # - see json.dump() documentation for documented errors for this call
        msg = ErrorNumbers.FB413.value + f' - save failed with message {str(e)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

@exp_exceptions
def check_training_plan_status(self) -> Responses:
    """ Method for checking training plan status, ie whether it is approved or not by the nodes

    Returns:
        Training plan status for answering nodes

    Raises:
        FedbiomedExperimentError: bad argument type
    """
    # at this point, self._job exists (initialized in constructor)
    if self._job is None:
        # cannot check training plan status if job not defined
        msg = ErrorNumbers.FB412.value + \
              ', in method `check_training_plan_status` : no `job` defined for experiment'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # always returns a `Responses()` object
    responses = self._job.check_training_plan_is_approved_by_nodes()

    return responses

@exp_exceptions
def experimentation_folder(self) -> str:
    """Retrieves the folder name where experiment data/result are saved.

    Please see also[`set_experimentation_folder`]
    [fedbiomed.researcher.experiment.Experiment.set_experimentation_folder]

    Returns:
        File name where experiment related files are saved
    """

    return self._experimentation_folder

@exp_exceptions
def experimentation_path(self) -> str:
    """Retrieves the file path where experimentation folder is located and experiment related files are saved.

    Returns:
        Experiment directory where all experiment related files are saved
    """

    return os.path.join(environ['EXPERIMENTS_DIR'], self._experimentation_folder)

@exp_exceptions
def info(self) -> None:
    """Prints out the information about the current status of the experiment.

    Lists  all the parameters/arguments of the experiment and informs whether the experiment can be run.

    Raises:
        FedbiomedExperimentError: Inconsistent experiment due to missing variables
    """

    # at this point all attributes are initialized (in constructor)
    info = {
        'Arguments': [
            'Tags',
            'Nodes filter',
            'Training Data',
            'Aggregator',
            'Strategy',
            'Job',
            'Training Plan Path',
            'Training Plan Class',
            'Model Arguments',
            'Training Arguments',
            'Rounds already run',
            'Rounds total',
            'Experiment folder',
            'Experiment Path',
            'Breakpoint State',
            'Secure Aggregation'
        ],
        # max 60 characters per column for values - can we do that with tabulate() ?
        'Values': ['\n'.join(findall('.{1,60}',
                                     str(e))) for e in [
            self._tags,
            self._nodes,
            self._fds,
            self._aggregator.aggregator_name if self._aggregator is not None else None,
            self._node_selection_strategy,
            self._job,
            self._training_plan_path,
            self._training_plan_class,
            self._model_args,
            self._training_args,
            self._round_current,
            self._round_limit,
            self._experimentation_folder,
            self.experimentation_path(),
            self._save_breakpoints,
            f'- Using: {self._secagg}\n- Active: {self._secagg.active}'
        ]
        ]
    }
    print(tabulate(info, headers='keys'))

    # definitions that may be missing for running the experiment
    # (value None == not defined yet for _fds et _job,
    # False == no valid model for _training_plan_is_defined )
    may_be_missing = {
        '_fds': 'Training Data',
        '_node_selection_strategy': 'Strategy',
        '_training_plan_is_defined': 'Training Plan',
        '_job': 'Job'
    }
    # definitions found missing
    missing = ''

    for key, value in may_be_missing.items():
        try:
            if eval('self.' + key) is None or eval('self.' + key) is False:
                missing += f'- {value}\n'
        except Exception:
            # should not happen, all eval variables should be defined
            msg = ErrorNumbers.FB400.value + \
                f', in method `info` : self.{key} not defined for experiment'
            logger.critical(msg)
            raise FedbiomedExperimentError(msg)
    if missing:
        print(f'\nExperiment cannot be run (not fully defined), missing :\n{missing}')
    else:
        print('\nExperiment can be run now (fully defined)')

@exp_exceptions
def job(self) -> Union[Job, None]:
    """Retrieves the [`Job`][fedbiomed.researcher.job] that manages training rounds.

    Returns:
        Initialized `Job` object. None, if it isn't declared yet or not information to set to job. Please see
            [`set_job`][fedbiomed.researcher.experiment.Experiment.set_job].
    """

    return self._job

@classmethod
@exp_exceptions
def load_breakpoint(cls: Type[TExperiment],
                    breakpoint_folder_path: Union[str, None] = None) -> TExperiment:
    """
    Loads breakpoint (provided a breakpoint has been saved)
    so experience can be resumed. Useful if training has crashed
    researcher side or if user wants to resume experiment.

    Args:
      cls: Experiment class
      breakpoint_folder_path: path of the breakpoint folder. Path can be absolute or relative eg:
        "var/experiments/Experiment_xxxx/breakpoints_xxxx". If None, loads latest breakpoint of the latest
        experiment. Defaults to None.

    Returns:
        Reinitialized experiment object. With given object-0.2119,  0.0796, -0.0759, user can then use `.run()`
            method to pursue model training.

    Raises:
        FedbiomedExperimentError: bad argument type, error when reading breakpoint or bad loaded breakpoint
            content (corrupted)
    """
    # check parameters type
    if not isinstance(breakpoint_folder_path, str) and breakpoint_folder_path is not None:
        msg = (
            f"{ErrorNumbers.FB413.value}: load failed, `breakpoint_folder_path`"
            f" has bad type {type(breakpoint_folder_path)}"
        )
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # get breakpoint folder path (if it is None) and state file
    breakpoint_folder_path, state_file = find_breakpoint_path(breakpoint_folder_path)
    breakpoint_folder_path = os.path.abspath(breakpoint_folder_path)

    try:
        path = os.path.join(breakpoint_folder_path, state_file)
        with open(path, "r", encoding="utf-8") as file:
            saved_state = json.load(file)
    except (json.JSONDecodeError, OSError) as exc:
        # OSError: heuristic for catching file access issues
        msg = (
            f"{ErrorNumbers.FB413.value}: load failed,"
            f" reading breakpoint file failed with message {exc}"
        )
        logger.critical(msg)
        raise FedbiomedExperimentError(msg) from exc
    if not isinstance(saved_state, dict):
        msg = (
            f"{ErrorNumbers.FB413.value}: load failed, breakpoint file seems"
            f" corrupted. Type should be `dict` not {type(saved_state)}"
        )
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # retrieve breakpoint training data
    bkpt_fds = saved_state.get('training_data')
    bkpt_fds = FederatedDataSet(bkpt_fds)
    # retrieve breakpoint sampling strategy
    bkpt_sampling_strategy_args = saved_state.get("node_selection_strategy")

    bkpt_sampling_strategy = cls._create_object(bkpt_sampling_strategy_args, data=bkpt_fds)

    # initializing experiment
    loaded_exp = cls(tags=saved_state.get('tags'),
                     nodes=None,  # list of previous nodes is contained in training_data
                     training_data=bkpt_fds,
                     node_selection_strategy=bkpt_sampling_strategy,
                     round_limit=saved_state.get("round_limit"),
                     training_plan_class=saved_state.get("training_plan_class"),
                     training_plan_path=saved_state.get("training_plan_path"),
                     model_args=saved_state.get("model_args"),
                     training_args=saved_state.get("training_args"),
                     save_breakpoints=True,
                     experimentation_folder=saved_state.get('experimentation_folder'),
                     secagg=SecureAggregation.load_state(saved_state.get('secagg')))

    # nota: we are initializing experiment with no aggregator: hence, by default,
    # `loaded_exp` will be loaded with FedAverage.

    # changing `Experiment` attributes
    loaded_exp._set_round_current(saved_state.get('round_current'))

    # TODO: checks when loading parameters
    training_plan = loaded_exp.training_plan()
    if training_plan is None:
        msg = ErrorNumbers.FB413.value + ' - load failed, ' + \
            'breakpoint file seems corrupted, `training_plan` is None'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    else:
        loaded_exp._aggregated_params = loaded_exp._load_aggregated_params(
            saved_state.get('aggregated_params')
        )

    # retrieve and change federator
    bkpt_aggregator_args = saved_state.get("aggregator")

    bkpt_aggregator = loaded_exp._create_object(bkpt_aggregator_args, training_plan=training_plan)
    loaded_exp.set_aggregator(bkpt_aggregator)

    # changing `Job` attributes
    loaded_exp._job.load_state(saved_state.get('job'))

    logger.info(f"Experimentation reload from {breakpoint_folder_path} successful!")
    return loaded_exp

@exp_exceptions
def model_args(self) -> dict:
    """Retrieves model arguments.

    Please see also [`set_model_args`][fedbiomed.researcher.experiment.Experiment.set_model_args]

    Returns:
        The arguments that are going to be passed to [`training_plans`][fedbiomed.common.training_plans]
            classes in built time on the node side.
    """
    return self._model_args

@exp_exceptions
def monitor(self) -> Monitor:
    """Retrieves the monitor object

    Monitor is responsible for receiving and parsing real-time training and validation feed-back from each node
    participate to federated training. See [`Monitor`][fedbiomed.researcher.monitor.Monitor]

    Returns:
        Monitor object that will always exist with experiment to retrieve feed-back from the nodes.
    """
    return self._monitor

@exp_exceptions
def nodes(self) -> Union[List[str], None]:
    """Retrieves the `nodes` that are chosen for federated training.

    Please see [`set_nodes`][fedbiomed.researcher.experiment.Experiment.set_nodes] to set `nodes`.

    Returns:
        Object that contains meta-data for the datasets of each node. `None` if nodes are not set.
    """
    return self._nodes

@exp_exceptions
def round_current(self) -> int:
    """Retrieves the round where the experiment is at.

    Returns:
        Indicates the round number that the experiment will perform next.
    """
    return self._round_current

@exp_exceptions
def round_limit(self) -> Union[int, None]:
    """Retrieves the round limit from the experiment object.

    Please see  also [`set_round_limit`][fedbiomed.researcher.experiment.Experiment.set_training_data] to change
    or set round limit.

    Returns:
        Round limit that shows maximum number of rounds that can be performed. `None` if it isn't declared yet.
    """
    return self._round_limit

@exp_exceptions
def run(self, rounds: Union[int, None] = None, increase: bool = False) -> int:
    """Run one or more rounds of an experiment, continuing from the point the
    experiment had reached.

    Args:
        rounds: Number of experiment rounds to run in this call.
            * `None` means "run all the rounds remaining in the experiment" computed as
                maximum rounds (`round_limit` for this experiment) minus the number of
                rounds already run rounds (`round_current` for this experiment).
                It does nothing and issues a warning if `round_limit` is `None` (no
                round limit defined for the experiment)
            * `int` >= 1 means "run at most `rounds` rounds".
                If `round_limit` is `None` for the experiment, run exactly `rounds` rounds.
                If a `round_limit` is set for the experiment and the number or rounds would
            increase beyond the `round_limit` of the experiment:
            - if `increase` is True, increase the `round_limit` to
              (`round_current` + `rounds`) and run `rounds` rounds
            - if `increase` is False, run (`round_limit` - `round_current`)
              rounds, don't modify the maximum `round_limit` of the experiment
              and issue a warning.
        increase: automatically increase the `round_limit`
            of the experiment for executing the specified number of `rounds`.
            Does nothing if `round_limit` is `None` or `rounds` is None.
            Defaults to False

    Returns:
        Number of rounds have been run

    Raises:
        FedbiomedExperimentError: bad argument type or value
    """
    # check rounds is a >=1 integer or None
    if rounds is None:
        pass
    elif isinstance(rounds, int):
        if rounds < 1:
            msg = ErrorNumbers.FB410.value + \
                f', in method `run` param `rounds` : value {rounds}'
            logger.critical(msg)
            raise FedbiomedExperimentError(msg)
    else:
        # bad type
        msg = ErrorNumbers.FB410.value + \
            f', in method `run` param `rounds` : type {type(rounds)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
        # check increase is a boolean
    if not isinstance(increase, bool):
        msg = ErrorNumbers.FB410.value + \
            f', in method `run` param `increase` : type {type(increase)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # compute number of rounds to run + updated rounds limit
    if rounds is None:
        if isinstance(self._round_limit, int):
            # run all remaining rounds in the experiment
            rounds = self._round_limit - self._round_current
            if rounds == 0:
                # limit already reached
                logger.warning(f'Round limit of {self._round_limit} already reached '
                               'for this experiment, do nothing.')
                return 0
        else:
            # cannot run if no number of rounds given and no round limit exists
            logger.warning('Cannot run, please specify a number of `rounds` to run or '
                           'set a `round_limit` to the experiment')
            return 0

    else:
        # at this point, rounds is an int >= 1
        if isinstance(self._round_limit, int):
            if (self._round_current + rounds) > self._round_limit:
                if increase:
                    # dont change rounds, but extend self._round_limit as necessary
                    logger.debug(f'Auto increasing total rounds for experiment from {self._round_limit} '
                                 f'to {self._round_current + rounds}')
                    self._round_limit = self._round_current + rounds
                else:
                    new_rounds = self._round_limit - self._round_current
                    if new_rounds == 0:
                        # limit already reached
                        logger.warning(f'Round limit of {self._round_limit} already reached '
                                       'for this experiment, do nothing.')
                        return 0
                    else:
                        # reduce the number of rounds to run in the experiment
                        logger.warning(f'Limit of {self._round_limit} rounds for the experiment '
                                       f'will be reached, reducing the number of rounds for this '
                                       f'run from {rounds} to {new_rounds}')
                        rounds = new_rounds

    # At this point `rounds` is an int > 0 (not None)

    # run the rounds
    for _ in range(rounds):
        if isinstance(self._round_limit, int) and self._round_current == (self._round_limit - 1) \
                and self._training_args['test_on_global_updates'] is True:
            # Do "validation after a round" only if this a round limit is defined and we reached it
            # and validation is active on global params
            # When this condition is met, it also means we are running the last of
            # the `rounds` rounds in this function
            test_after = True
        else:
            test_after = False

        increment = self.run_once(increase=False, test_after=test_after)

        if increment == 0:
            # should not happen
            msg = ErrorNumbers.FB400.value + \
                f', in method `run` method `run_once` returns {increment}'
            logger.critical(msg)
            raise FedbiomedExperimentError(msg)

    return rounds

@exp_exceptions
def run_once(self, increase: bool = False, test_after: bool = False) -> int:
    """Run at most one round of an experiment, continuing from the point the
    experiment had reached.

    If `round_limit` is `None` for the experiment (no round limit defined), run one round.
    If `round_limit` is not `None` and the `round_limit` of the experiment is already reached:
    * if `increase` is False, do nothing and issue a warning
    * if `increase` is True, increment total number of round `round_limit` and run one round

    Args:
        increase: automatically increase the `round_limit` of the experiment if needed. Does nothing if
            `round_limit` is `None`. Defaults to False
        test_after: if True, do a second request to the nodes after the round, only for validation on aggregated
            params. Intended to be used after the last training round of an experiment. Defaults to False.

    Returns:
        Number of rounds really run

    Raises:
        FedbiomedExperimentError: bad argument type or value
    """
    # check increase is a boolean
    if not isinstance(increase, bool):
        msg = ErrorNumbers.FB410.value + \
              f', in method `run_once` param `increase` : type {type(increase)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # nota:  we should never have self._round_current > self._round_limit, only ==
    if self._round_limit is not None and self._round_current >= self._round_limit:
        if increase is True:
            logger.debug(f'Auto increasing total rounds for experiment from {self._round_limit} '
                         f'to {self._round_current + 1}')
            self._round_limit = self._round_current + 1
        else:
            logger.warning(f'Round limit of {self._round_limit} was reached, do nothing')
            return 0

    # at this point, self._aggregator always exists and is not None
    # self.{_node_selection_strategy,_job} exist but may be None

    # check pre-requisites are met for running a round
    # for component in (self._node_selection_strategy, self._job):
    if self._node_selection_strategy is None:
        msg = ErrorNumbers.FB411.value + ', missing `node_selection_strategy`'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    elif self._job is None:
        msg = ErrorNumbers.FB411.value + ', missing `job`'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # Ready to execute a training round using the job, strategy and aggregator
    if self._global_model is None:
        self._global_model = self._job.training_plan.get_model_params()
        # initial server state, before optimization/aggregation

    self._aggregator.set_training_plan_type(self._job.training_plan.type())
    # Sample nodes using strategy (if given)
    self._job.nodes = self._node_selection_strategy.sample_nodes(self._round_current)

    # If secure aggregation is activated ---------------------------------------------------------------------
    secagg_arguments = None
    if self._secagg.active:
        self._secagg.setup(parties=[environ["ID"]] + self._job.nodes,
                           job_id=self._job.id)
        secagg_arguments = self._secagg.train_arguments()
    # --------------------------------------------------------------------------------------------------------

    # Check aggregator parameter(s) before starting a round
    self._aggregator.check_values(n_updates=self._training_args.get('num_updates'),
                                  training_plan=self._job.training_plan)
    logger.info('Sampled nodes in round ' + str(self._round_current) + ' ' + str(self._job.nodes))

    aggr_args_thr_msg, aggr_args_thr_file = self._aggregator.create_aggregator_args(self._global_model,
                                                                                    self._job.nodes)

    # Trigger training round on sampled nodes
    _ = self._job.start_nodes_training_round(round_=self._round_current,
                                             aggregator_args_thr_msg=aggr_args_thr_msg,
                                             aggregator_args_thr_files=aggr_args_thr_file,
                                             do_training=True,
                                             secagg_arguments=secagg_arguments)

    # refining/normalizing model weights received from nodes
    model_params, weights, total_sample_size, encryption_factors = self._node_selection_strategy.refine(
        self._job.training_replies[self._round_current], self._round_current)

    self._aggregator.set_fds(self._fds)

    if self._secagg.active:
        flatten_params = self._secagg.aggregate(
            round_=self._round_current,
            encryption_factors=encryption_factors,
            total_sample_size=total_sample_size,
            model_params=model_params
        )
        # FIXME: Access TorchModel through non-private getter once it is implemented
        aggregated_params: Dict[str, Union['torch.tensor', 'nd.ndarray']] = \
            self._job.training_plan._model.unflatten(flatten_params)

    else:
        # aggregate models from nodes to a global model
        aggregated_params = self._aggregator.aggregate(model_params,
                                                       weights,
                                                       global_model=self._global_model,
                                                       training_plan=self._job.training_plan,
                                                       training_replies=self._job.training_replies,
                                                       node_ids=self._job.nodes,
                                                       n_updates=self._training_args.get('num_updates'),
                                                       n_round=self._round_current)

    # write results of the aggregated model in a temp file

    # Export aggregated parameters to a local file and upload it.
    # Also assign the new values to the job's training plan's model.
    self._global_model = aggregated_params  # update global model
    aggregated_params_path, _ = self._job.update_parameters(aggregated_params)
    logger.info(f'Saved aggregated params for round {self._round_current} '
                f'in {aggregated_params_path}')

    self._aggregated_params[self._round_current] = {'params': aggregated_params,
                                                    'params_path': aggregated_params_path}

    self._round_current += 1

    # Update round in monitor for the next round
    self._monitor.set_round(round_=self._round_current + 1)

    if self._save_breakpoints:
        self.breakpoint()

    # do final validation after saving breakpoint :
    # not saved in breakpoint for current round, but more simple
    if test_after:
        # FIXME: should we sample nodes here too?
        aggr_args_thr_msg, aggr_args_thr_file = self._aggregator.create_aggregator_args(self._global_model,
                                                                                        self._job.nodes)
        self._job.start_nodes_training_round(round_=self._round_current,
                                             aggregator_args_thr_msg=aggr_args_thr_msg,
                                             aggregator_args_thr_files=aggr_args_thr_file,
                                             do_training=False)

    return 1

@exp_exceptions
def save_breakpoints(self) -> bool:
    """Retrieves the status of saving breakpoint after each round of training.

    Returns:
        `True`, If saving breakpoint is active. `False`, vice versa.
    """

    return self._save_breakpoints

@property
def secagg(self) -> SecureAggregation:
    """Gets secagg object `SecureAggregation`

    Returns:
        Secure aggregation object.
    """
    return self._secagg

@exp_exceptions
def set_aggregator(self, aggregator: Union[Aggregator, Type[Aggregator], None]) -> \
        Aggregator:
    """Sets aggregator + verification on arguments type

    Args:
        aggregator: Object or class defining the method for aggregating local updates. Default to None
            (use `FedAverage` for aggregation)

    Returns:
        aggregator (Aggregator)

    Raises:
        FedbiomedExperimentError : bad aggregator type
    """

    if aggregator is None:
        # default aggregator
        self._aggregator = FedAverage()
    elif inspect.isclass(aggregator):
        # a class is provided, need to instantiate an object
        if issubclass(aggregator, Aggregator):
            self._aggregator = aggregator()
        else:
            # bad argument
            msg = ErrorNumbers.FB410.value + f' `aggregator` : {aggregator} class'
            logger.critical(msg)
            raise FedbiomedExperimentError(msg)
    elif isinstance(aggregator, Aggregator):
        # an object of a proper class is provided, nothing to do
        self._aggregator = aggregator
    else:
        # other bad type or object
        msg = ErrorNumbers.FB410.value + f' `aggregator` : {type(aggregator)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    # at this point self._aggregator is (non-None) aggregator object
    self.aggregator_args["aggregator_name"] = self._aggregator.aggregator_name
    if self._fds is not None:
        self._aggregator.set_fds(self._fds)

    return self._aggregator

@exp_exceptions
def set_experimentation_folder(self, experimentation_folder: Union[str, None]) -> str:
    """Sets `experimentation_folder`, the folder name where experiment data/result are saved.

    Args:
        experimentation_folder: File name where experiment related files are saved

    Returns:
        experimentation_folder (str)

    Raises:
        FedbiomedExperimentError : bad `experimentation_folder` type
    """
    if experimentation_folder is None:
        self._experimentation_folder = create_exp_folder()
    elif isinstance(experimentation_folder, str):
        sanitized_folder = sanitize_filename(experimentation_folder, platform='auto')
        self._experimentation_folder = create_exp_folder(sanitized_folder)
        if (sanitized_folder != experimentation_folder):
            logger.warning(f'`experimentation_folder` was sanitized from '
                           f'{experimentation_folder} to {sanitized_folder}')
    else:
        msg = ErrorNumbers.FB410.value + \
            f' `experimentation_folder` : {type(experimentation_folder)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

        # at this point self._experimentation_folder is a str valid for a foldername

    # _job doesn't always exist at this point
    if self._job is not None:
        logger.debug('Experimentation folder changed, you may need to update `job`')

    return self._experimentation_folder

@exp_exceptions
def set_job(self) -> Union[Job, None]:
    """Setter for job, it verifies pre-requisites are met for creating a job
    attached to this experiment. If yes, instantiate a job ; if no, return None.

    Returns:
        The object that is initialized for creating round jobs.
    """
    # at this point all are defined among:
    # self.{_reqs,_fds,_training_plan_is_defined,_training_plan,_training_plan_path,_model_args,_training_args}
    # self._experimentation_folder => self.experimentation_path()
    # self._round_current

    if self._job is not None:
        # a job is already defined, and it may also have run some rounds
        logger.debug('Experimentation `job` changed after running '
                     '{self._round_current} rounds, may give inconsistent results')

    if self._training_plan_is_defined is not True:
        # training plan not properly defined yet
        self._job = None
        logger.debug('Experiment not fully configured yet: no job. Missing proper training plan '
                     f'definition (training_plan={self._training_plan_class} '
                     f'training_plan_path={self._training_plan_path})')
    elif self._fds is None:
        # not training data yet
        self._job = None
        logger.debug('Experiment not fully configured yet: no job. Missing training data')
    else:
        # meeting requisites for instantiating a job
        self._job = Job(reqs=self._reqs,
                        training_plan_class=self._training_plan_class,
                        training_plan_path=self._training_plan_path,
                        model_args=self._model_args,
                        training_args=self._training_args,
                        data=self._fds,
                        keep_files_dir=self.experimentation_path())

    return self._job

@exp_exceptions
def set_model_args(self, model_args: dict) -> dict:
    """Sets `model_args` + verification on arguments type

    Args:
        model_args (dict): contains model arguments passed to the constructor
            of the training plan when instantiating it : output and input feature
            dimension, etc.

    Returns:
        Model arguments that have been set.

    Raises:
        FedbiomedExperimentError : bad model_args type
    """
    if isinstance(model_args, dict):
        self._model_args = model_args
    else:
        # bad type
        msg = ErrorNumbers.FB410.value + f' `model_args` : {type(model_args)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    # self._model_args always exist at this point

    if self._job is not None:
        logger.debug('Experimentation model_args changed, you may need to update `job`')
    return self._model_args

@exp_exceptions
def set_nodes(self, nodes: Union[List[str], None]) -> Union[List[str], None]:
    """Sets for nodes + verifications on argument type

    Args:
        nodes: List of node_ids to filter the nodes to be involved in the experiment.

    Returns:
        List of tags that are set. None, if the argument `nodes` is None.

    Raises:
        FedbiomedExperimentError : Bad nodes type
    """

    if isinstance(nodes, list):
        for node in nodes:
            if not isinstance(node, str):
                msg = ErrorNumbers.FB410.value + f' `nodes` : list of {type(node)}'
                logger.critical(msg)
                raise FedbiomedExperimentError(msg)
        self._nodes = nodes
    elif nodes is None:
        self._nodes = nodes
    else:
        msg = ErrorNumbers.FB410.value + f' `nodes` : {type(nodes)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    # self._nodes always exist at this point

    if self._fds is not None:
        logger.debug('Experimentation nodes filter changed, you may need to update `training_data`')

    return self._nodes

@exp_exceptions
def set_round_limit(self, round_limit: Union[int, None]) -> Union[int, None]:
    """Sets `round_limit` + verification on arguments type

    Args:
        round_limit: the maximum number of training rounds (nodes <-> central server) that should be executed
            for the experiment. `None` means that no limit is defined.

    Returns:
        Round limit for experiment of federated learning

    Raises:
        FedbiomedExperimentError : bad rounds type or value
    """
    # at this point round_current exists and is an int >= 0

    if round_limit is None:
        # no limit for training rounds
        self._round_limit = None
    elif isinstance(round_limit, int):
        # at this point round_limit is an int
        if round_limit < 0:
            msg = ErrorNumbers.FB410.value + f' `round_limit` can not be negative: {round_limit}'
            logger.critical(msg)
            raise FedbiomedExperimentError(msg)
        elif round_limit < self._round_current:
            # self._round_limit can't be less than current round
            logger.error(f'cannot set `round_limit` to less than the number of already run rounds '
                         f'({self._round_current})')
        else:
            self._round_limit = round_limit
    else:
        msg = ErrorNumbers.FB410.value + f' `round_limit` : {type(round_limit)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    # at this point self._round_limit is a Union[int, None]
    return self._round_limit

@exp_exceptions
def set_save_breakpoints(self, save_breakpoints: bool) -> bool:
    """ Setter for save_breakpoints + verification on arguments type

    Args:
        save_breakpoints (bool): whether to save breakpoints or
            not after each training round. Breakpoints can be used for resuming
            a crashed experiment.

    Returns:
        Status of saving breakpoints

    Raises:
        FedbiomedExperimentError: bad save_breakpoints type
    """
    if isinstance(save_breakpoints, bool):
        self._save_breakpoints = save_breakpoints
        # no warning if done during experiment, we may change breakpoint policy at any time
    else:
        msg = ErrorNumbers.FB410.value + f' `save_breakpoints` : {type(save_breakpoints)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    return self._save_breakpoints

@exp_exceptions
def set_secagg(self, secagg: Union[bool, SecureAggregation]):

    if isinstance(secagg, bool):
        self._secagg = SecureAggregation(active=secagg, timeout=10)
    elif isinstance(secagg, SecureAggregation):
        self._secagg = secagg
    else:
        msg = f"{ErrorNumbers.FB410.value}: Expected `secagg` argument bool or `SecureAggregation`, " \
              f"but got {type(secagg)}"
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    return self._secagg

@exp_exceptions
def set_strategy(self, node_selection_strategy: Union[Strategy, Type[Strategy], None]) -> \
        Union[Strategy, None]:
    """Sets for `node_selection_strategy` + verification on arguments type

    Args:
        node_selection_strategy: object or class defining how nodes are sampled at each round for training, and
            how non-responding nodes are managed. Defaults to None:
            - use `DefaultStrategy` if training_data is initialized
            - else strategy is None (cannot be initialized), experiment cannot
              be launched yet

    Returns:
        node selection strategy class

    Raises:
        FedbiomedExperimentError : bad strategy type
    """
    if self._fds is not None:
        if node_selection_strategy is None:
            # default node_selection_strategy
            self._node_selection_strategy = DefaultStrategy(self._fds)
        elif inspect.isclass(node_selection_strategy):
            # a class is provided, need to instantiate an object
            if issubclass(node_selection_strategy, Strategy):
                self._node_selection_strategy = node_selection_strategy(self._fds)
            else:
                # bad argument
                msg = ErrorNumbers.FB410.value + \
                    f' `node_selection_strategy` : {node_selection_strategy} class'
                logger.critical(msg)
                raise FedbiomedExperimentError(msg)
        elif isinstance(node_selection_strategy, Strategy):
            # an object of a proper class is provided, nothing to do
            self._node_selection_strategy = node_selection_strategy
        else:
            # other bad type or object
            msg = ErrorNumbers.FB410.value + \
                f' `node_selection_strategy` : {type(node_selection_strategy)}'
            logger.critical(msg)
            raise FedbiomedExperimentError(msg)
    else:
        # cannot initialize strategy if not FederatedDataSet yet
        self._node_selection_strategy = None
        logger.debug('Experiment not fully configured yet: no node selection strategy')

    # at this point self._node_selection_strategy is a Union[Strategy, None]
    return self._node_selection_strategy

@exp_exceptions
def set_tags(self, tags: Union[List[str], str, None]) -> Union[List[str], None]:
    """Sets tags + verifications on argument type

    Args:
        tags: List of string with data tags or string with one data tag. Empty list
            of tags ([]) means any dataset is accepted, it is different from None (tags not set, cannot search
            for training_data yet).

    Returns:
        List of tags that are set. None, if the argument `tags` is None.

    Raises:
        FedbiomedExperimentError : Bad tags type
    """

    if isinstance(tags, list):
        for tag in tags:
            if not isinstance(tag, str):
                msg = ErrorNumbers.FB410.value + f' `tags` : list of {type(tag)}'
                logger.critical(msg)
                raise FedbiomedExperimentError(msg)
        self._tags = tags
    elif isinstance(tags, str):
        self._tags = [tags]
    elif tags is None:
        self._tags = tags
    else:
        msg = ErrorNumbers.FB410.value + f' `tags` : {type(tags)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)
    # self._tags always exist at this point

    if self._fds is not None:
        logger.debug('Experimentation tags changed, you may need to update `training_data`')

    return self._tags

@exp_exceptions
def set_tensorboard(self, tensorboard: bool) -> bool:
    """
    Sets the tensorboard flag

    Args:
        tensorboard: If `True` tensorboard log files will be writen after receiving training feedbacks

    Returns:
        Status of tensorboard
    """

    if isinstance(tensorboard, bool):
        self._tensorboard = tensorboard
        self._monitor.set_tensorboard(tensorboard)
    else:
        msg = ErrorNumbers.FB410.value + f' `tensorboard` : {type(tensorboard)}'
        logger.critical(msg)
        raise FedbiomedExperimentError(msg)

    return self._tensorboard

@exp_exceptions
def set_test_metric(self, metric: Union[MetricTypes, str, None], **metric_args: dict) -> \
        Tuple[Union[str, None], Dict[str, Any]]:
    """ Sets a metric for federated model validation

    Args:
        metric: A class as an instance of [`MetricTypes`][fedbiomed.common.metrics.MetricTypes]. [`str`][str] for
            referring one of  metric which provided as attributes in [`MetricTypes`]
            [fedbiomed.common.metrics.MetricTypes]. None, if it isn't declared yet.
        **metric_args: A dictionary that contains arguments for metric function. Arguments
            should be compatible with corresponding metrics in [`sklearn.metrics`][sklearn.metrics].

    Returns:
        Metric and  metric args as tuple

    Raises:
        FedbiomedExperimentError: Invalid type for `metric` argument
    """
    self._training_args['test_metric'] = metric

    # using **metric_args, we know `test_metric_args` is a Dict[str, Any]
    self._training_args['test_metric_args'] = metric_args

    if self._job is not None:
        # job setter function exists, use it
        self._job.training_args = self._training_args
        logger.debug('Experimentation training_args updated for `job`')

    return metric, metric_args