from dipy.testing.decorators import warning_for_keywords
from dipy.workflows.workflow import Workflow


class CombinedWorkflow(Workflow):
    @warning_for_keywords()
    def __init__(
        self, *, output_strategy="append", mix_names=False, force=False, skip=False
    ):
        """Workflow that combines multiple workflows.
        The workflow combined together are referred as sub flows in this class.
        """

        self._optionals = {}
        super(CombinedWorkflow, self).__init__(
            output_strategy=output_strategy, mix_names=mix_names, force=force, skip=skip
        )

    def get_sub_runs(self):
        """Returns a list of tuples
        (sub flow name, sub flow run method, sub flow short name)
        to be used in the sub flow parameters extraction.
        """
        sub_runs = []
        for flow in self._get_sub_flows():
            sub_runs.append((flow.__name__, flow.run, flow.get_short_name()))

        return sub_runs

    def _get_sub_flows(self):
        """Returns a list of sub flows used in the combined_workflow. Needs to
        be implemented in every new combined_workflow.
        """
        raise AttributeError(
            f"Error: _get_sub_flows() has to be defined for {self.__class__}"
        )

    def set_sub_flows_optionals(self, opts):
        """Sets the self._optionals variable with all sub flow arguments
        that were passed in the commandline.
        """
        self._optionals = {}
        for key, sub_dict in opts.items():
            self._optionals[key] = {k: v for k, v in sub_dict.items() if v is not None}

    def get_optionals(self, flow, **kwargs):
        """Returns the sub flow's optional arguments merged with those passed
        as params in kwargs.
        """
        opts = self._optionals[flow.__name__]
        opts.update(kwargs)

        return opts

    def run_sub_flow(self, flow, *args, **kwargs):
        """Runs the sub flow with the optional parameters passed via the
        command line. This is a convenience method to make sub flow running
        more intuitive on the concrete CombinedWorkflow side.
        """
        return flow.run(*args, **self.get_optionals(type(flow), **kwargs))
