# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2024, ETH Zurich
# All rights reserved.
#
# Please, refer to the LICENSE file in the root directory.
# SPDX-License-Identifier: BSD-3-Clause

import copy
import functools
from typing import Any, Optional

import dace
from dace import (
    data as dace_data,
    dtypes as dace_dtypes,
    properties as dace_properties,
    subsets as dace_sbs,
    transformation as dace_transformation,
)
from dace.sdfg import (
    nodes as dace_nodes,
    propagation as dace_propagation,
    type_inference as dace_type_inference,
    utils as dace_sutils,
)

from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations


@dace_properties.make_properties
class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
    """The transformation moves dataflow into the if branches.

    Essentially transforms code from this
    ```python
    __arg1 = foo(...)
    __arg2 = bar(...)
    if __cond:
        __output = __arg1
    else:
        __output = __arg2
    ```
    into this
    ```python
    if __cond:
        __output = foo(...)
    else:
        __output = bar(...)
    ```
    I.e. computation that is needed in only one branch is relocated or inlined into
    that specific branch, which might reduce computation.

    Despite its name the transformation is not only able to handle `if` expressions,
    but more general `switch` like expressions. The requirements are:
    - The expression must be inside a nested SDFG containing only a `ConditionalBlock`.
    - The matched nested SDFG must be inside a Map (might be dropped).
    - For every incoming connector there must be exactly one AccessNode in the
        entire nested SDFG (might be dropped).
    - Every branch must write to every output of the nested SDFG.
    - The only dataflow allowed inside the branches is `(i) -> (o)` i.e. AccessNode
        to AccessNode connections.

    Furthermore, the transformation should be applied as long as possible, i.e.
    it should be applied by passing it to `SDFG.apply_transformations_repeated()`.

    Args:
        ignore_upstream_blocks: If `True` do not require that upstream `if_block`s
            have to be processed first.

    Note:
        - If there is a chain of suitable `if` expression, i.e. an `if` expression is
            in the upstream dataflow of the one that is currently matched, then the
            transformation will not apply. This is done to ensure that everything is
            properly inlined into the branches. This behaviour can be disabled by
            setting `ignore_upstream_blocks` to `True`.
            The reason for this behaviour is that the current implementation can only
            handle the case where the `if` expression is _directly_ inside a Map.

    Todo:
        - Allow that an inconnector can be used multiple times, as long as it is in
            different branches.
        - Extend the implementation that it is also able to handle `if_blocks` that
            are not inside a Map. This would allow to drop the need to process
            upstream `if` expressions first.
    """

    if_block = dace_transformation.PatternNode(dace_nodes.NestedSDFG)

    ignore_upstream_blocks = dace_properties.Property(
        dtype=bool,
        default=False,
        desc="If 'True' ignore 'if_block's that are upstream.",
    )

    @classmethod
    def expressions(cls) -> Any:
        return [dace.sdfg.utils.node_path_graph(cls.if_block)]

    def __init__(
        self,
        ignore_upstream_blocks: Optional[bool] = None,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        if ignore_upstream_blocks is not None:
            self.ignore_upstream_blocks = ignore_upstream_blocks
        super().__init__(*args, **kwargs)

    def can_be_applied(
        self,
        graph: dace.SDFGState,
        expr_index: int,
        sdfg: dace.SDFG,
        permissive: bool = False,
    ) -> bool:
        scope_dict = graph.scope_dict()
        if_block: dace_nodes.NestedSDFG = self.if_block
        assert isinstance(if_block, dace_nodes.NestedSDFG)

        # NOTE: The main benefit of requiring that the `if_block` must be inside a
        #   Map is that we know that every transient we encounter is single use data.
        enclosing_map = scope_dict[if_block]
        if not isinstance(enclosing_map, dace_nodes.MapEntry):
            return False

        # Test if the `if_block` is valid. This will also give us the names.
        if_block_spec = self._partition_if_block(if_block)
        if if_block_spec is None:
            return False

        # Compute the dataflow that is relocated.
        raw_relocatable_dataflow, non_relocatable_dataflow = (
            {
                conn_name: gtx_transformations.utils.find_upstream_nodes(
                    start=if_block,
                    state=graph,
                    start_connector=conn_name,
                    limit_node=enclosing_map,
                )
                for conn_name in conn_names
            }
            for conn_names in if_block_spec
        )
        relocatable_dataflow = self._filter_relocatable_dataflow(
            sdfg=sdfg,
            state=graph,
            if_block=if_block,
            raw_relocatable_dataflow=raw_relocatable_dataflow,
            non_relocatable_dataflow=non_relocatable_dataflow,
            enclosing_map=enclosing_map,
        )

        # If no branch has something to inline then we are done.
        if all(len(rel_df) == 0 for rel_df in relocatable_dataflow.values()):
            return False

        # Check if relatability is possible.
        if not self._check_for_data_and_symbol_conflicts(
            sdfg=sdfg,
            state=graph,
            relocatable_dataflow=relocatable_dataflow,
            enclosing_map=enclosing_map,
            if_block=if_block,
        ):
            return False

        # Because the transformation can only handle `if` expressions that
        #  are _directly_ inside a Map, we must check if the upstream contains
        #  suitable `if` expressions that must be processed first. The simplest way
        #  is to return `False` and not apply. However, this implies that the
        #  transformation is applied in a loop until it applies nowhere anymore.
        # NOTE: This is a restriction due to the current implementation.
        if not self.ignore_upstream_blocks:
            for reloc_dataflow in relocatable_dataflow.values():
                if any(
                    self._has_if_block_relocatable_dataflow(
                        sdfg=sdfg,
                        state=graph,
                        upstream_if_block=upstream_if_block,
                        enclosing_map=enclosing_map,
                    )
                    for upstream_if_block in reloc_dataflow
                    if isinstance(upstream_if_block, dace_nodes.NestedSDFG)
                ):
                    return False

        return True

    def apply(
        self,
        graph: dace.SDFGState,
        sdfg: dace.SDFG,
    ) -> None:
        if_block: dace_nodes.NestedSDFG = self.if_block
        if_block_spec = self._partition_if_block(if_block)
        assert if_block_spec is not None
        enclosing_map = graph.scope_dict()[if_block]

        # Find the dataflow that should be relocated.
        raw_relocatable_dataflow, non_relocatable_dataflow = (
            {
                conn_name: gtx_transformations.utils.find_upstream_nodes(
                    start=if_block,
                    state=graph,
                    start_connector=conn_name,
                    limit_node=enclosing_map,
                )
                for conn_name in conn_names
            }
            for conn_names in if_block_spec
        )
        relocatable_dataflow = self._filter_relocatable_dataflow(
            sdfg=sdfg,
            state=graph,
            if_block=if_block,
            raw_relocatable_dataflow=raw_relocatable_dataflow,
            non_relocatable_dataflow=non_relocatable_dataflow,
            enclosing_map=enclosing_map,
        )

        # Create a mapping from connector names to the corresponding AccessNode inside the branch and the branch state.
        # This is necessary to properly patch the dataflow inside the branch, i.e. to connect it to the global AccessNode
        # corresponding to the connector if necessary. We need to gather this information for all the connectors because
        # a node in the dataflow of one connector might be the global AccessNode of another connector and we have to handle
        # it properly in `_replicate_dataflow_into_branch`.
        conn_name_to_access_node_map: dict[str, tuple[dace.SDFGState, dace_nodes.AccessNode]] = {}
        for conn_name in relocatable_dataflow.keys():
            conn_name_to_access_node_map[conn_name] = self._find_branch_for(
                if_block=if_block,
                connector=conn_name,
            )

        # Create a map of the old to the new nodes to keep track of the old nodes copied
        # and their corresponding new nodes. The map should be per branch of the ConditionalBlock
        # because there could be a node that has to be copied in all branches. Since the
        # `relocatable_dataflow` nodes are not disjoint we need this mapping to avoid copying the same
        # node multiple times and to properly connect the copied nodes.
        old_to_new_nodes_map: dict[tuple[dace_nodes.Node, dace.SDFGState], dace_nodes.Node] = dict()
        # Finally relocate the dataflow
        for conn_name, nodes_to_move in relocatable_dataflow.items():
            self._replicate_dataflow_into_branch(
                state=graph,
                sdfg=sdfg,
                if_block=if_block,
                enclosing_map=enclosing_map,
                nodes_to_move=nodes_to_move,
                connector=conn_name,
                conn_name_to_access_node_map=conn_name_to_access_node_map,
                old_to_new_nodes_map=old_to_new_nodes_map,
            )

        self._update_symbol_mapping(if_block, sdfg)

        self._remove_outside_dataflow(
            sdfg=sdfg,
            state=graph,
            relocatable_dataflow=relocatable_dataflow,
        )

        # Because we relocate some node it seems that DaCe gets a bit confused.
        #  So we have to reset the list. Without it the `test_if_mover_chain`
        #  test would fail.
        sdfg.reset_cfg_list()

        # Readjust the Subsets.
        #  As a side effect this call will also properly propagate the nested SDFG `if_block`.
        dace_propagation.propagate_memlets_map_scope(sdfg, graph, enclosing_map)

    def _replicate_dataflow_into_branch(
        self,
        sdfg: dace.SDFG,
        state: dace.SDFGState,
        if_block: dace_nodes.NestedSDFG,
        enclosing_map: dace_nodes.MapEntry,
        nodes_to_move: set[dace_nodes.Node],
        connector: str,
        conn_name_to_access_node_map: dict[str, tuple[dace.SDFGState, dace_nodes.AccessNode]],
        old_to_new_nodes_map: dict[tuple[dace_nodes.Node, dace.SDFGState], dace_nodes.Node],
    ) -> None:
        """Replicate the dataflow in `nodes_to_move` from `state` into `if_block`.

        First the function will determine into which branch, inside `if_block`,
        the dataflow has to be replicated. It will then copy the dataflow, nodes
        listed in `nodes_to_move` and insert them into that state.
        The function will then create the edges to connect them in the same way
        as they where outside. If there is an outer data dependency, for example
        a read to a global memory, then the function will patch that inside the
        `if_block`.
        At the end the function will remove the `connector`, but it will not remove
        the original dataflow.
        In case of nodes existing in multiple `nodes_to_move` sets coming from multiple
        connectors, the function will only copy the necessary nodes and edges only once
        based on the `old_to_new_nodes_map` keys and values. In case a connector to the
        NestedSDFG is exists in the dataflow of another connector, the function
        will remove the connection of the original connector and replace the global
        AccessNode of the NestedSDFG with a temporary one and remove the original connector.

        Args:
            sdfg: The sdfg that we process, the one that contains `state`.
            state: The state we operate on, the one that contains `if_block`.
            if_block: The `if_block` into which we inline.
            enclosing_map: The enclosing map.
            nodes_to_move: The list of nodes that should be removed.
            connector: The connector that should be inlined.
            conn_name_to_access_node_map: A mapping from connector names to the
                corresponding AccessNode inside the branch and the branch state.
            old_to_new_nodes_map: A mapping from the old nodes to the new nodes.
                The keys of the mapping are tuples of the old node and the branch
                state for which the new node was created. The values are the new nodes.
        """
        # Nothing to relocate nothing to do.
        if len(nodes_to_move) == 0:
            return

        inner_sdfg: dace.SDFG = if_block.sdfg
        branch_state, connector_node = conn_name_to_access_node_map[connector]

        # Replicate the nodes and store them in the `old_to_new_nodes_map` mapping.
        # Add the SDFGState to the key of the dictionary because we have to create
        # new node for the different branches.
        unique_old_nodes: list[dace_nodes.Node] = []
        for old_node in nodes_to_move:
            if (old_node, branch_state) in old_to_new_nodes_map:
                continue
            unique_old_nodes.append(old_node)
            copy_of_old_node = copy.deepcopy(old_node)
            old_to_new_nodes_map[(old_node, branch_state)] = copy_of_old_node
            branch_state.add_node(copy_of_old_node)

            # There might be AccessNodes inside `nodes_to_move`, we now have to make sure
            #  that they are present inside the nested ones. By our base assumption they
            #  are transients and single use, because they are only used in one place. Make
            #  sure that we don't add new nodes if the nodes are already in the arrays of
            #  the inner SDFG.
            if not isinstance(old_node, dace_nodes.AccessNode):
                continue
            if old_node.data in inner_sdfg.arrays:
                continue
            assert sdfg.arrays[old_node.data].transient
            # TODO(phimuell): Handle the case we need to rename something.
            inner_sdfg.add_datadesc(
                old_node.data,
                sdfg.arrays[old_node.data].clone(),
                find_new_name=False,
            )

        # Now add the edges between the edges that have been replicated inside the
        #  branch state, these are the outgoing edges.
        #  Now add the outgoing edges between the replicated nodes inside the branch state.
        #  The data dependencies (incomming edges) of the nodes, i.e. the not relocated dataflow,
        #  are still missing.
        for node in unique_old_nodes:
            for oedge in state.out_edges(node):
                if oedge.dst is if_block:
                    if oedge.dst_conn == connector:
                        # This connection maps the outside data into the nested SDFG, thus its destination is technically the same data.
                        # TODO(phimuell): Make subsets complete.
                        # TODO(phimuell): Check if this Memlet is always correct, especially in case of slicing.
                        branch_state.add_edge(
                            old_to_new_nodes_map[(oedge.src, branch_state)],
                            oedge.src_conn,
                            connector_node,
                            None,
                            dace.Memlet.from_memlet(oedge.data),
                        )
                    else:
                        assert oedge.dst_conn in conn_name_to_access_node_map
                        # Some of the nodes_to_move are also in the dataflow of some other connector whose dataflow
                        # we should also move. Since we already moved the dataflow of the other connector we connect
                        # the dataflow to the global AccessNode corresponding to the other connector
                        branch_state.add_edge(
                            old_to_new_nodes_map[(oedge.src, branch_state)],
                            oedge.src_conn,
                            conn_name_to_access_node_map[oedge.dst_conn][1],
                            None,
                            dace.Memlet.from_memlet(oedge.data),
                        )
                        # Handle the other connector AccessNode and connector
                        inner_access_node_of_connector_name = conn_name_to_access_node_map[
                            oedge.dst_conn
                        ][1].data
                        if not inner_sdfg.arrays[inner_access_node_of_connector_name].transient:
                            inner_sdfg.arrays[inner_access_node_of_connector_name].transient = True
                            if_block.remove_in_connector(inner_access_node_of_connector_name)
                else:
                    # Restore a connection between two nodes that were relocated.
                    assert oedge.dst in nodes_to_move
                    branch_state.add_edge(
                        old_to_new_nodes_map[(oedge.src, branch_state)],
                        oedge.src_conn,
                        old_to_new_nodes_map[(oedge.dst, branch_state)],
                        oedge.dst_conn,
                        dace.Memlet.from_memlet(oedge.data),
                    )

        # Now we have to satisfy the data dependencies, i.e. forward all nodes that
        #  could not have been moved inside `if_block` but are still needed to compute
        #  the final result. We find them by scanning the input edges of the nodes
        #  that have been relocated.
        for node in unique_old_nodes:
            for iedge in state.in_edges(node):
                if iedge.src in nodes_to_move:
                    # Inner data dependency, there is nothing to do and the edge was
                    #  created above.
                    continue
                if iedge.data.is_empty():
                    # Empty Memlets are there to maintain some order relation, "happens
                    #  before". Depending on the situation we can remove or have to
                    #  recreate them. The case where the connection comes from a
                    #  node within the relocated dataflow is handled above.
                    assert iedge.src is enclosing_map
                    continue

                # Now we have to figuring out where the data is coming from, since
                #  the data was not relocated, we must patch it into `if_block`.
                if iedge.src is enclosing_map:
                    # The data is coming from outside the Map scope, i.e. not defined
                    #  inside the Map scope, so we have to trace it back.
                    memlet_path = state.memlet_path(iedge)
                    outer_data = memlet_path[0].src
                else:
                    # The data is defined somewhere in the Map scope itself.
                    outer_data = iedge.src
                # TODO(phimuell): It is possible that this does not lead to an
                #   AccessNode on the outside, but to something inside the Map scope
                #   such as the MapExit of an inner map. To handle such a case we need
                #   to construct the set of nodes to move differently, i.e.
                #   considering this case already there.
                if not isinstance(outer_data, dace_nodes.AccessNode):
                    raise NotImplementedError()
                assert not gtx_transformations.utils.is_view(outer_data, sdfg)

                # If the data is not yet available in the inner SDFG made
                #  patch it through.
                if outer_data.data not in inner_sdfg.arrays:
                    inner_desc = sdfg.arrays[outer_data.data].clone()
                    inner_desc.transient = False
                    # TODO(phimuell): Handle the case we need to rename something.
                    inner_sdfg.add_datadesc(outer_data.data, inner_desc, False)
                    # TODO(phimeull): We pass the whole data inside the SDFG.
                    #   Find out if there are cases where this is wrong.
                    state.add_edge(
                        iedge.src,
                        iedge.src_conn,
                        if_block,
                        outer_data.data,
                        dace.Memlet(
                            data=outer_data.data, subset=dace_sbs.Range.from_array(inner_desc)
                        ),
                    )
                    if_block.add_in_connector(outer_data.data)
                else:
                    # This is the case that we found a node, that refers to data that
                    #  was already patched into the `if_block`. We would have to remove
                    #  this, but since this function just replicates the dataflow,
                    #  it will not do that. Instead we postpone this to the cleanup
                    #  phase, see `_remove_outside_dataflow()`.
                    pass

                if (outer_data, branch_state) not in old_to_new_nodes_map:
                    assert all(
                        outer_data.data != mapped_node.data
                        for mapped_node, mapped_branch_state in old_to_new_nodes_map.keys()
                        if isinstance(mapped_node, dace_nodes.AccessNode)
                        and mapped_branch_state == branch_state
                    )
                    assert outer_data.data in inner_sdfg.arrays
                    assert not inner_sdfg.arrays[outer_data.data].transient
                    old_to_new_nodes_map[(outer_data, branch_state)] = branch_state.add_access(
                        outer_data.data, copy.copy(outer_data.debuginfo)
                    )

                # Now create the edge in the inner state.
                branch_state.add_edge(
                    old_to_new_nodes_map[(outer_data, branch_state)],
                    None,
                    old_to_new_nodes_map[(iedge.dst, branch_state)],
                    iedge.dst_conn,
                    copy.deepcopy(iedge.data),
                )

        # The old connector name is no longer valid.
        inner_sdfg.arrays[connector].transient = True
        if_block.remove_in_connector(connector)

    def _remove_outside_dataflow(
        self,
        sdfg: dace.SDFG,
        state: dace.SDFGState,
        relocatable_dataflow: dict[str, set[dace_nodes.Node]],
    ) -> None:
        """Removes the original dataflow, that has been relocated.

        The function will also remove data containers that are no longer in use.
        """
        all_relocatable_dataflow: set[dace_nodes.Node] = functools.reduce(
            lambda s1, s2: s1.union(s2), relocatable_dataflow.values(), set()
        )

        # Before we can clean the original nodes, we must clean the dataflow. If a
        #  node, that was relocated, has incoming connections we must remove them
        #  and the parent dataflow.
        for node_to_remove in all_relocatable_dataflow:
            for iedge in list(state.in_edges(node_to_remove)):
                if iedge.src in all_relocatable_dataflow:
                    continue
                dace_sutils.remove_edge_and_dangling_path(state, iedge)

            if isinstance(node_to_remove, dace_nodes.AccessNode):
                assert node_to_remove.desc(sdfg).transient
                sdfg.remove_data(node_to_remove.data, validate=False)

        # Remove the original nodes (data descriptors were deleted in the loop above).
        state.remove_nodes_from(all_relocatable_dataflow)

    def _update_symbol_mapping(
        self,
        if_block: dace_nodes.NestedSDFG,
        parent: dace.SDFG,
    ) -> None:
        """Updates the symbol mapping of the nested SDFG.

        The function assumes that the symbols that are missing in the nested SDFG
        are available in the parent SDFG.
        """
        symbol_mapping = if_block.symbol_mapping
        missing_symbols = [ms for ms in if_block.sdfg.free_symbols if ms not in symbol_mapping]
        symbol_mapping.update({s: s for s in missing_symbols})
        if_block.symbol_mapping = symbol_mapping  # Performs conversion.

        # Add new global symbols to nested SDFG.
        #  The code is based on `SDFGState.add_nested_sdfg()`.
        if_block_symbols = if_block.sdfg.symbols
        parent_symbols = parent.symbols
        for new_sym in missing_symbols:
            if new_sym in if_block_symbols:
                # The symbol is already known, so we check that it is the same type as in the
                #  parent SDFG.
                assert if_block_symbols[new_sym] == parent_symbols[new_sym]

            elif new_sym in parent_symbols:
                # The symbol is known to the parent SDFG, so take the type from there.
                if_block.sdfg.add_symbol(new_sym, parent_symbols[new_sym])

            else:
                # Figuring out the type of the symbol based on the computation we do.
                # TODO(phimuell): Maybe switch to `symbols_defined_at()` as it is indicated
                #   in the `SDFGState.add_nested_sdfg()` function.
                if_block.sdfg.add_symbol(
                    new_sym,
                    dace_type_inference.infer_expr_type(new_sym, parent_symbols)
                    or dace_dtypes.typeclass(int),
                )

    def _check_for_data_and_symbol_conflicts(
        self,
        sdfg: dace.SDFG,
        state: dace.SDFGState,
        relocatable_dataflow: dict[str, set[dace_nodes.Node]],
        if_block: dace_nodes.NestedSDFG,
        enclosing_map: dace_nodes.MapEntry,
    ) -> bool:
        """Check if the relocation would cause any conflict, such as a symbol clash."""

        # TODO(phimuell): There is an obscure case where the nested SDFG, on its own,
        #   defines a symbol that is also mapped, for example a dynamic Map range.
        #   It is probably not a problem, because of the scopes DaCe adds when
        #   generating the C++ code.

        # Create a subgraph to compute the free symbols, i.e. the symbols that
        #  need to be supplied from the outside. However, this are not all.
        #  Note, just adding some "well chosen" nodes to the set will not work.
        all_relocated_dataflow: set[dace_nodes.Node] = functools.reduce(
            lambda s1, s2: s1.union(s2), relocatable_dataflow.values(), set()
        )
        requiered_symbols: set[str] = dace.sdfg.state.StateSubgraphView(
            state, all_relocated_dataflow
        ).free_symbols

        inner_data_names = if_block.sdfg.arrays.keys()
        for node_to_check in all_relocated_dataflow:
            if (
                isinstance(node_to_check, dace_nodes.AccessNode)
                and node_to_check.data in inner_data_names
            ):
                # There is already a data descriptor that is used on the inside as on
                #  the outside. Thus we would have to perform some renaming, which we
                #  currently do not.
                # TODO(phimell): Handle this case.
                return False

            for iedge in state.in_edges(node_to_check):
                src_node = iedge.src
                if src_node not in all_relocated_dataflow:
                    # This means that `src_node` is not relocated but mapped into the
                    #  `if` block. This means that `edge` is replicated as well.
                    # NOTE: This code is based on the one found in `DataflowGraphView`.
                    # TODO(phimuell): Do we have to inspect the full Memlet path here?
                    assert isinstance(src_node, dace_nodes.AccessNode) or src_node is enclosing_map
                    requiered_symbols |= iedge.data.used_symbols(True, edge=iedge)

        # The (beyond the enclosing Map) data is also mapped into the `if` block, so we
        #  have to consider that as well.
        for iedge in state.in_edges(if_block):
            if iedge.src is enclosing_map and (not iedge.data.is_empty()):
                outside_desc = sdfg.arrays[iedge.data.data]
                if isinstance(outside_desc, dace_data.View):
                    return False  # Handle this case.
                requiered_symbols |= outside_desc.used_symbols(True)

        # A conflicting symbol is a free symbol of the relocatable dataflow, that is not a
        #  direct mapping. For example if there is a symbol `n` on the inside and outside
        #  then everything is okay if the symbol mapping is `{n: n}` i.e. the symbol has the
        #  same meaning inside and outside. Everything else is not okay.
        symbol_mapping = if_block.symbol_mapping
        conflicting_symbols = requiered_symbols.intersection((str(k) for k in symbol_mapping))
        for conflicting_symbol in conflicting_symbols:
            if conflicting_symbol != str(symbol_mapping[conflicting_symbol]):
                return False

        return True

    def _find_branch_for(
        self,
        if_block: dace_nodes.NestedSDFG,
        connector: str,
    ) -> tuple[dace.SDFGState, dace_nodes.AccessNode]:
        """
        Locates the branch and the AccessNode to where the dataflow should be relocated.
        """
        inner_sdfg: dace.SDFG = if_block.sdfg
        conditional_block: dace.sdfg.state.ConditionalBlock = next(iter(inner_sdfg.nodes()))

        # This will locate the state where the first AccessNode that refers to
        #  `connector` is found. Since `_partition_if_block()` makes sure that
        #  there is only one match this is okay. But it must be changed, if we
        #  lift this restriction.
        for inner_state in conditional_block.all_states():
            connector_nodes: list[dace_nodes.AccessNode] = [
                dnode for dnode in inner_state.data_nodes() if dnode.data == connector
            ]
            if len(connector_nodes) == 0:
                continue
            break
        else:
            raise ValueError(f"Did not find a branch associated to '{connector}'.")

        assert isinstance(inner_state, dace.SDFGState)
        assert inner_state.in_degree(connector_nodes[0]) == 0
        assert inner_state.out_degree(connector_nodes[0]) > 0
        return inner_state, connector_nodes[0]

    def _has_if_block_relocatable_dataflow(
        self,
        sdfg: dace.SDFG,
        state: dace.SDFGState,
        upstream_if_block: dace_nodes.NestedSDFG,
        enclosing_map: dace_nodes.MapEntry,
    ) -> bool:
        """Check if `upstream_if_block` has relocatable dataflow.

        The function is used to enforce the rule that no `if` expression should be
        processed before all suitable `if` expressions in its upstream dataflow of
        the matched `if_block` has been processed first.

        Args:
            sdfg: The SDFG on which we operate.
            state: The state on which we operate on.
            upstream_if_block: The `if` expression that was found in the relocatable
                dataflow of the matched `if_block`.
            enclosing_map: The limiting node, i.e. the MapEntry of the Map `if_block`
                is located in.
        """
        if_block_spec = self._partition_if_block(upstream_if_block)
        if if_block_spec is None:
            return False

        raw_relocatable_dataflow, non_relocatable_dataflow = (
            {
                conn_name: gtx_transformations.utils.find_upstream_nodes(
                    start=upstream_if_block,
                    state=state,
                    start_connector=conn_name,
                    limit_node=enclosing_map,
                )
                for conn_name in conn_names
            }
            for conn_names in if_block_spec
        )
        filtered_relocatable_dataflow = self._filter_relocatable_dataflow(
            sdfg=sdfg,
            state=state,
            if_block=upstream_if_block,
            raw_relocatable_dataflow=raw_relocatable_dataflow,
            non_relocatable_dataflow=non_relocatable_dataflow,
            enclosing_map=enclosing_map,
        )
        if all(len(rel_df) == 0 for rel_df in filtered_relocatable_dataflow.values()):
            return False

        return True

    def _filter_relocatable_dataflow(
        self,
        sdfg: dace.SDFG,
        state: dace.SDFGState,
        if_block: dace_nodes.NestedSDFG,
        raw_relocatable_dataflow: dict[str, set[dace_nodes.Node]],
        non_relocatable_dataflow: dict[str, set[dace_nodes.Node]],
        enclosing_map: dace_nodes.MapEntry,
    ) -> dict[str, set[dace_nodes.Node]]:
        """Partition the dependencies.

        The function expects the dataflow that is upstream of every connector
        of the `if_block`. The function will then scan the dataflow and compute
        the parts that actually can be relocated and returns a `dict` mapping
        every relocatable input connector to the set of nodes that can be relocated.

        The returned sets can include duplicate nodes, i.e. a node can be in the
        dataflow of multiple connectors. The function that performs the actual
        relocation (`_replicate_dataflow_into_branch`) will take care of that and
        make sure that such nodes are only copied once.

        Args:
            state: The state on which we operate.
            if_block: The `if_block` that is processed.
            raw_relocatable_dataflow: The connectors and their associated dataflow
                that can be relocated, not yet filtered.
            non_relocatable_dataflow: The connectors and their associated dataflow
                that can not be relocated.
            enclosing_map: The limiting node, i.e. the MapEntry of the Map where
                `if_block` is located in.
        """

        # Remove the parts of the dataflow that is unrelocatable.
        all_non_relocatable_dataflow: set[dace_nodes.Node] = functools.reduce(
            lambda s1, s2: s1.union(s2), non_relocatable_dataflow.values(), set()
        )
        relocatable_dataflow = {
            conn_name: rel_df.difference(all_non_relocatable_dataflow)
            for conn_name, rel_df in raw_relocatable_dataflow.items()
        }

        # Find the known_nodes for each branch
        known_nodes: dict[dace.SDFGState, set[dace_nodes.Node]] = dict()
        for conn_name, rel_df in relocatable_dataflow.items():
            branch_state, _ = self._find_branch_for(if_block=if_block, connector=conn_name)
            if branch_state not in known_nodes:
                known_nodes[branch_state] = set()
            known_nodes[branch_state].update(rel_df)

        multiple_df_nodes: set[dace_nodes.Node] = set()
        # Find intersect of all known_nodes sets which are the nodes that are in the dataflow
        # of multiple branches and thus doesn't make sense to relocate
        for branch_state, known_nodes_set in known_nodes.items():
            for other_branch_state, other_known_nodes_set in known_nodes.items():
                if branch_state != other_branch_state:
                    multiple_df_nodes.update(known_nodes_set.intersection(other_known_nodes_set))

        if multiple_df_nodes:
            # Remove from the relocatable dataflow the nodes that appear in multiple branches
            # as it doesn't make sense to relocate them and duplicate them in both branches.
            relocatable_dataflow = {
                conn_name: rel_df.difference(multiple_df_nodes)
                for conn_name, rel_df in relocatable_dataflow.items()
            }

        # TODO(phimuell): If we operate outside of a Map we also have to make sure that
        #   the data is single use data, is not an AccessNode that refers to global
        #   memory nor is a source AccessNode.
        def filter_nodes(
            nodes_proposed_for_reloc: set[dace_nodes.Node],
        ) -> set[dace_nodes.Node]:
            has_been_updated = True
            while has_been_updated:
                has_been_updated = False

                for reloc_node in list(nodes_proposed_for_reloc):
                    # The node was already handled in a previous iteration.
                    if reloc_node not in nodes_proposed_for_reloc:
                        continue

                    assert (
                        state.in_degree(reloc_node) > 0
                    )  # Because we are currently always inside a Map

                    # If the node is needed by anything that is not also moved
                    #  into the `if` body, then it has to remain outside. For that we
                    #  have to pretend that `if_block` is also relocated.
                    if any(
                        oedge.dst not in nodes_proposed_for_reloc
                        for oedge in state.out_edges(reloc_node)
                        if oedge.dst is not if_block
                    ):
                        nodes_proposed_for_reloc.remove(reloc_node)
                        has_been_updated = True
                        continue

                    # We do not look at all incoming nodes, but have to ignore some of them.
                    #  We ignore `enclosed_map` because it acts as boundary, and the node
                    #  on the other side of it is mapped into the `if` body anyway. We
                    #  ignore the AccessNodes because they will either be relocated into
                    #  the `if` body or be mapped (remain outside but made accessible
                    #  inside), thus their relocation state is of no concern for
                    #  `reloc_node`.
                    non_mappable_incoming_nodes: set[dace_nodes.Node] = {
                        iedge.src
                        for iedge in state.in_edges(reloc_node)
                        if not (
                            (iedge.src is enclosing_map)
                            or isinstance(iedge.src, dace_nodes.AccessNode)
                        )
                    }
                    if non_mappable_incoming_nodes.issubset(nodes_proposed_for_reloc):
                        # All nodes that can not be mapped into the `if` body are
                        #  currently scheduled to be relocated, thus there is not
                        #  problem.
                        pass

                    else:
                        # Only some of the non mappable nodes are selected to be
                        #  moved inside the `if` body. This means that `reloc_node`
                        #  can also not be moved because of its input dependencies.
                        #  Since we can not relocate `reloc_node` this also implies
                        #  that none of its input can. Thus we remove them from
                        #  `nodes_proposed_for_reloc`.
                        nodes_proposed_for_reloc.difference_update(non_mappable_incoming_nodes)
                        nodes_proposed_for_reloc.remove(reloc_node)
                        has_been_updated = True

            return nodes_proposed_for_reloc

        return {
            conn_name: filter_nodes(rel_df) for conn_name, rel_df in relocatable_dataflow.items()
        }

    def _partition_if_block(
        self,
        if_block: dace_nodes.NestedSDFG,
    ) -> Optional[tuple[set[str], set[str]]]:
        """Check if `if_block` can be processed and partition the input connectors.

        The function will check if `if_block` has the right structure, i.e. if it is
        roughly equivalent to the Python expression `a if cond else b`.
        It will also identify which input connectors refer to dataflow that can
        be inlined into the `if_block` and which can not.

        Returns:
            If `if_block` is unsuitable the function will return `None`.
            If `if_block` meets the structural requirements the function will return
            two sets of strings. The first set contains the connectors that can be
            relocated and the second one of the conditions that can not be relocated.
        """
        # TODO(phimuell): Change the return type to `tuple[list[str], list[str]]` and sort the connectors, such that the operation is deterministic.
        # There shall only be one output and three inputs with given names.
        if len(if_block.out_connectors.keys()) == 0:
            return None

        # These are all the output names.
        output_names: set[str] = set(if_block.out_connectors.keys())

        # We require that the nested SDFG contains a single node, which is a
        #  `ConditionalBlock` containing two branches.
        inner_sdfg: dace.SDFG = if_block.sdfg
        if inner_sdfg.number_of_nodes() != 1:
            return None
        inner_if_block = next(iter(inner_sdfg.nodes()))
        if not isinstance(inner_if_block, dace.sdfg.state.ConditionalBlock):
            return None

        # Defining it outside will ensure that there is only one AccessNode for every
        #  inconnector, which is something `_find_branch_for()` relies on.
        reference_count: dict[str, int] = {conn_name: 0 for conn_name in if_block.in_connectors}

        for _, branch in inner_if_block.branches:
            output_count: dict[str, int] = {conn_name: 0 for conn_name in output_names}
            for inner_state in branch.all_states():
                assert isinstance(inner_state, dace.SDFGState)
                for node in inner_state.nodes():
                    if not isinstance(node, dace_nodes.AccessNode):
                        return None
                    if node.data in reference_count:
                        reference_count[node.data] += 1
                        exp_in_deg, exp_out_deg = 0, 1
                    elif node.data in output_count:
                        output_count[node.data] += 1
                        exp_in_deg, exp_out_deg = 1, 0
                    else:
                        return None
                    if inner_state.in_degree(node) != exp_in_deg:
                        return None
                    if inner_state.out_degree(node) != exp_out_deg:
                        return None

        # The connectors that can be pulled inside must appear exactly once.
        #  In theory they could appear more, but then we would have to replicate
        #  the dataflow to different locations which is not supported.
        #  So the ones that can be relocated were found exactly once. Zero would
        #  mean they can not be relocated and more than one means that we do not
        #  support it yet.
        relocatable_connectors = {
            conn_name for conn_name, conn_count in reference_count.items() if conn_count == 1
        }
        non_relocatable_connectors = {
            conn_name
            for conn_name in reference_count.keys()
            if conn_name not in relocatable_connectors
        }

        if len(non_relocatable_connectors) == 0:
            return None
        if len(relocatable_connectors) == 0:
            return None
        return relocatable_connectors, non_relocatable_connectors
