"""
 Copyright (c) 2019-2020 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

from typing import Callable, List, Optional, Tuple, Any
import os

import networkx as nx
from copy import deepcopy
from networkx.drawing.nx_agraph import to_agraph
from torch import Tensor

from nncf.dynamic_graph.graph_matching import Expression, NodeExpression, search_all, get_edge_boundaries
from nncf.dynamic_graph.trace_tensor import TensorMeta, TracedTensor
from nncf.layers import ITERATION_MODULES

from nncf.nncf_logger import logger as nncf_logger


# pylint: disable=too-many-public-methods
class TensorMetaComparator:
    def __call__(self, lhs: TensorMeta, rhs: TensorMeta) -> bool:
        raise NotImplementedError


class DefaultTensorMetaComparator(TensorMetaComparator):
    def __call__(self, lhs: TensorMeta, rhs: TensorMeta) -> bool:
        return TensorMeta.default_comparator(lhs, rhs)


class ShapeIgnoringTensorMetaComparator(TensorMetaComparator):
    def __call__(self, lhs: TensorMeta, rhs: TensorMeta) -> bool:
        return lhs.creator_id == rhs.creator_id and lhs.index == rhs.index


class ShapeOnlyTensorMetaComparator(TensorMetaComparator):
    def __call__(self, lhs: TensorMeta, rhs: TensorMeta) -> bool:
        return lhs.shape[1:] == rhs.shape[1:]


class InputsMatcher:
    def __call__(self, node_inputs: List[TensorMeta], real_inputs: List[TensorMeta],
                 tm_comparators: List[TensorMetaComparator]) -> bool:
        raise NotImplementedError


class FirstInputsMatcher(InputsMatcher):
    def __call__(self, node_inputs: List[TensorMeta], real_inputs: List[TensorMeta],
                 tm_comparators: List[TensorMetaComparator]) -> bool:
        if not node_inputs or not real_inputs:
            return False

        if not node_inputs[0] or not real_inputs[0]:
            return False

        for tm_comparator in tm_comparators:
            if not tm_comparator(node_inputs[0], real_inputs[0]):
                return False
        return True


class DefaultInputsMatcher(InputsMatcher):
    def __call__(self, node_inputs: List[TensorMeta], real_inputs: List[TensorMeta],
                 tm_comparators: List[TensorMetaComparator]) -> bool:
        if node_inputs is None and real_inputs:
            return False

        for saved_input, actual_input in zip(node_inputs, real_inputs):
            if saved_input is None and actual_input is None:
                continue
            if (saved_input is None) != (actual_input is None):
                return False
            for tm_comparator in tm_comparators:
                if not tm_comparator(saved_input, actual_input):
                    return False
        return True


class InputAgnosticOperationExecutionContext:
    def __init__(self, operator_name: str, scope_in_model: 'Scope', call_order: int):
        self.operator_name = operator_name
        self.scope_in_model = scope_in_model
        self.call_order = call_order

    def __eq__(self, other: 'InputAgnosticOperationExecutionContext'):
        return (self.operator_name == other.operator_name) and \
               (self.scope_in_model == other.scope_in_model) and \
               (self.call_order == other.call_order)

    def __str__(self):
        return str(self.scope_in_model) + '/' + \
               self.operator_name + "_" + str(self.call_order)

    def __hash__(self):
        return hash((self.operator_name, self.scope_in_model, self.call_order))

    @staticmethod
    def from_str(s: str):
        scope_and_op, _, call_order_str = s.rpartition('_')
        scope_str, _, op_name = scope_and_op.rpartition('/')

        from nncf.dynamic_graph.context import Scope
        return InputAgnosticOperationExecutionContext(op_name,
                                                      Scope.from_str(scope_str),
                                                      int(call_order_str))


class OperationExecutionContext:
    """Information that allows to uniquely identify an operation inside the NNCF graph,
    i.e. determine whether an execution of the operator inside the module has already been
    registered as a node in the graph or not (in the latter case a new node would have to
    be created"""

    def __init__(self,
                 operator_name: str,
                 scope_in_model: 'Scope',
                 call_order: int,
                 tensor_metas: List[TensorMeta],
                 tm_comparators: List[TensorMetaComparator] = None,
                 input_matcher: InputsMatcher = None):
        self.input_agnostic = InputAgnosticOperationExecutionContext(operator_name, scope_in_model, call_order)
        # This should be a list with a length equal to the number of inputs.
        # "None" values in this list correspond to non-tensor input elements.
        self.tensor_metas = tensor_metas
        self.tm_comparators = tm_comparators if tm_comparators else [
            DefaultTensorMetaComparator()]
        self.input_matcher = input_matcher if input_matcher else DefaultInputsMatcher()

    def __eq__(self, other: 'OperationExecutionContext'):
        return (self.input_agnostic == other.input_agnostic) and \
               self.input_matcher(self.tensor_metas, other.tensor_metas, self.tm_comparators)

    def __hash__(self):
        return hash((self.operator_name, tuple(self.scope_in_model), self.call_order,
                     tuple(self.tensor_metas)))

    def __str__(self):
        input_info_str = ""
        for meta in self.tensor_metas:
            if meta is None:
                input_info_str += "N;"
            else:
                input_info_str += str(meta) + ";"

        return super().__str__() + '(' + input_info_str + ')'

    @property
    def operator_name(self):
        return self.input_agnostic.operator_name

    @property
    def scope_in_model(self) -> 'Scope':
        return self.input_agnostic.scope_in_model

    @property
    def call_order(self):
        return self.input_agnostic.call_order


class NNCFNode:
    def __init__(self, node_id: int, op_exec_context: OperationExecutionContext):
        self.node_id = node_id
        self.op_exec_context = op_exec_context

    def __str__(self):
        return str(self.node_id) + " " + str(self.op_exec_context)

    def __hash__(self):
        return hash(str(self))

    def __eq__(self, other):
        return self.node_id == other.node_id and self.op_exec_context == other.op_exec_context


class DefaultScopeNodeMatcher:
    def __init__(self, node_id_to_key_dict, nx_graph, nx_node_to_nncf_node):
        self._node_id_to_key_dict = node_id_to_key_dict
        self._nx_graph = nx_graph
        self._nx_node_to_nncf_node = nx_node_to_nncf_node
        self._inputless_nx_nodes = dict()

    def get_node_by_id(self, node_id):
        return self._nx_graph.nodes[self._node_id_to_key_dict[node_id]]

    def _find_nodes_with_matching_context_among_inputless(self, op_exec_context: OperationExecutionContext):
        node_candidates = {}
        for nx_node_key, nx_node in self._inputless_nx_nodes.items():
            if nx_node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR] == op_exec_context:
                node_candidates[nx_node_key] = nx_node
        return node_candidates

    def _find_nodes_with_matching_context_and_inputs(self, op_exec_context: OperationExecutionContext):
        node_candidates = {}
        for info in op_exec_context.tensor_metas:
            if info is None or info.creator_id is None:
                continue
            creator_id = info.creator_id
            for successor_node_key in self._nx_graph.successors(self._node_id_to_key_dict[creator_id]):
                successor_node = self._nx_graph.nodes[successor_node_key]
                if op_exec_context == successor_node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR]:
                    node_candidates[successor_node_key] = successor_node
        return node_candidates

    def add_node(self, op_exec_context: OperationExecutionContext, inputs) -> NNCFNode:
        node_id = len(self._node_id_to_key_dict)

        name_parts = (str(op_exec_context.scope_in_model), op_exec_context.operator_name)
        node_key = '{idx} {uri}'.format(uri='/'.join(name_parts), idx=node_id)

        nncf_logger.debug("New node added to NNCF graph: {}".format(node_key))

        self._node_id_to_key_dict[node_id] = node_key
        attrs = {
            NNCFGraph.ID_NODE_ATTR: node_id,
            NNCFGraph.KEY_NODE_ATTR: node_key,
            NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR: op_exec_context,
        }
        self._nx_graph.add_node(node_key, **attrs)

        has_traced_inputs = False
        for i, info in enumerate(op_exec_context.tensor_metas):
            if info is None or info.creator_id is None:
                continue
            parent = self._node_id_to_key_dict[info.creator_id]
            self._nx_graph.add_edge(parent, node_key)
            has_traced_inputs = True
            self._nx_graph.edges[parent, node_key][NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR] = info.shape
            self._nx_graph.edges[parent, node_key][NNCFGraph.IN_PORT_NAME] = i

        if not has_traced_inputs:
            self._inputless_nx_nodes[node_key] = self._nx_graph.nodes[node_key]

        return self._nx_node_to_nncf_node(self._nx_graph.nodes[node_key])

    def find_node(self, ia_op_exec_context: InputAgnosticOperationExecutionContext,
                  tensor_metas: List[TensorMeta],
                  tm_comparators: List[TensorMetaComparator]) -> NNCFNode:
        op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name,
                                                    ia_op_exec_context.scope_in_model,
                                                    ia_op_exec_context.call_order,
                                                    tensor_metas,
                                                    tm_comparators=tm_comparators)
        nncf_node_candidates = []
        node_candidates = self._find_nodes_with_matching_context_and_inputs(op_exec_context)
        if not node_candidates:
            node_candidates = self._find_nodes_with_matching_context_among_inputless(op_exec_context)

        for nx_node in node_candidates.values():
            nncf_node_candidates.append(NNCFNode(nx_node[NNCFGraph.ID_NODE_ATTR],
                                                 op_exec_context))
        result = None
        if len(nncf_node_candidates) == 1:
            result = nncf_node_candidates[0]
        if len(nncf_node_candidates) > 1:
            nncf_logger.warning("More than one node matches input")
            result = nncf_node_candidates[0]

        return result


class IterationScopeNodeMatcher(DefaultScopeNodeMatcher):
    def __init__(self, node_id_to_key_dict, nx_graph, nx_node_to_nncf_node):
        super().__init__(node_id_to_key_dict, nx_graph, nx_node_to_nncf_node)
        self._first_iteration_nodes = {}  # type: {str: {str: NNCFNode}}

    @staticmethod
    def _get_iteration_scopes(scope: 'Scope') -> List[str]:
        results = []
        scope_name = str(scope)
        for iter_scope in ITERATION_MODULES.registry_dict:
            if iter_scope in scope_name:
                results.append(iter_scope)
        return results

    def save_first_iteration_node(self, inputs, node: NNCFNode):
        """
        It finds and saves "starting" points of iteration for further matching with them on next iteration,
        instead of adding new nodes for each iteration. "Starting" points of iteration are nodes
            * that have at least one input node, which is outside of iteration scope
            * or whose all inputs are not TracedTensor
        """
        op_exec_context = node.op_exec_context
        name = node
        iter_scopes = self._get_iteration_scopes(op_exec_context.scope_in_model)
        if iter_scopes:
            for iter_scope in iter_scopes:
                if iter_scope not in self._first_iteration_nodes:
                    self._first_iteration_nodes[iter_scope] = {}
                first_nodes = self._first_iteration_nodes[iter_scope]
                has_input_outside_iteration = False
                not_traced_count = 0
                for i in inputs:
                    if isinstance(i, Tensor):
                        has_input_outside_iteration = True
                        break
                    if not isinstance(i, TracedTensor):
                        not_traced_count += 1
                        continue
                    creator_id = i.tensor_meta.creator_id
                    creator_node = self.get_node_by_id(creator_id)
                    creator_node_op_exec_ctx = creator_node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR]
                    within_scopes = self._get_iteration_scopes(creator_node_op_exec_ctx.scope_in_model)
                    if iter_scope not in within_scopes:
                        has_input_outside_iteration = True
                if not_traced_count == len(inputs):
                    has_input_outside_iteration = True
                if has_input_outside_iteration:
                    node_name = str(op_exec_context.input_agnostic)
                    first_nodes[node_name] = node
                    nncf_logger.debug('Found first iteration node: {} in scope: {}'.format(name, iter_scope))

    def add_node(self, op_exec_context: OperationExecutionContext, inputs) -> NNCFNode:
        node = super().add_node(op_exec_context, inputs)
        self.save_first_iteration_node(inputs, node)
        return node

    def find_node(self,
                  ia_op_exec_context: InputAgnosticOperationExecutionContext,
                  tensor_metas: List[TensorMeta],
                  tm_comparators: List[TensorMetaComparator]) -> NNCFNode:
        nncf_node_candidates = []
        iter_scopes = self._get_iteration_scopes(ia_op_exec_context.scope_in_model)
        # compare meta information about first input nodes during the matching. During the iteration some nodes may
        # change number of inputs, e.g. on concat of hidden outputs
        input_matcher = FirstInputsMatcher()
        op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name,
                                                    ia_op_exec_context.scope_in_model,
                                                    ia_op_exec_context.call_order,
                                                    tensor_metas,
                                                    input_matcher=input_matcher,
                                                    tm_comparators=tm_comparators)
        node_candidates = self._find_nodes_with_matching_context_and_inputs(op_exec_context)
        if not node_candidates:
            op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name,
                                                        ia_op_exec_context.scope_in_model,
                                                        ia_op_exec_context.call_order,
                                                        tensor_metas,
                                                        tm_comparators=tm_comparators)
            node_candidates = self._find_nodes_with_matching_context_among_inputless(op_exec_context)
            if not node_candidates and iter_scopes:
                # ignore information about node creator and index of input
                comparators = tm_comparators + [ShapeOnlyTensorMetaComparator()]
                op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name,
                                                            ia_op_exec_context.scope_in_model,
                                                            ia_op_exec_context.call_order,
                                                            tensor_metas,
                                                            tm_comparators=comparators)
                # match with starting points of iteration
                iter_nodes = self._match_first_iteration_nodes(op_exec_context, iter_scopes)
                for node in iter_nodes.items():
                    nncf_node_candidates.append(node[1])

        for nx_node in node_candidates.values():
            nncf_node_candidates.append(NNCFNode(nx_node[NNCFGraph.ID_NODE_ATTR],
                                                 op_exec_context))

        result = None
        if len(nncf_node_candidates) == 1:
            result = nncf_node_candidates[0]
        if len(nncf_node_candidates) > 1:
            nncf_logger.warning("More than one node matches input")
            result = nncf_node_candidates[0]

        return result

    def _match_first_iteration_nodes(self, op_exec_context: OperationExecutionContext, iter_scopes):
        node_candidates = {}
        for iter_scope in iter_scopes:
            if iter_scope in self._first_iteration_nodes:
                for name, node in self._first_iteration_nodes[iter_scope].items():
                    if op_exec_context == node.op_exec_context:
                        node_candidates[name] = node
                        break
                if node_candidates:
                    break
        return node_candidates


class NodeManager:
    def __init__(self, node_id_to_key_dict, nx_graph, nx_node_to_nncf_node):
        self.base_matcher = DefaultScopeNodeMatcher(node_id_to_key_dict, nx_graph, nx_node_to_nncf_node)
        self.iteration_matcher = IterationScopeNodeMatcher(node_id_to_key_dict, nx_graph, nx_node_to_nncf_node)

    # TODO: optimize by matching exact module type
    @staticmethod
    def _within_iteration(scope: 'Scope'):
        scope_name = str(scope)
        for iter_scope in ITERATION_MODULES.registry_dict:
            if iter_scope in scope_name:
                return True
        return False

    def choose_matcher(self, ia_op_exec_context: InputAgnosticOperationExecutionContext) -> DefaultScopeNodeMatcher:
        if self._within_iteration(ia_op_exec_context.scope_in_model):
            return self.iteration_matcher
        return self.base_matcher

    @staticmethod
    def choose_tm_comparators(ia_op_exec_context: InputAgnosticOperationExecutionContext,
                              input_comparators_per_scope:
                              List[Tuple[TensorMetaComparator, List[str]]]) -> List[TensorMetaComparator]:
        result = []
        for pairs in input_comparators_per_scope:
            comparator, scopes = pairs
            for scope in scopes:
                if scope in str(ia_op_exec_context):
                    result.append(comparator)
        return result

    def find_node(self, ia_op_exec_context: InputAgnosticOperationExecutionContext,
                  tensor_metas: List[TensorMeta],
                  input_comparators_per_scope: List[Tuple[TensorMetaComparator, List[str]]]) -> NNCFNode:
        matcher = self.choose_matcher(ia_op_exec_context)
        comparators = self.choose_tm_comparators(ia_op_exec_context, input_comparators_per_scope)
        return matcher.find_node(ia_op_exec_context, tensor_metas, comparators)

    def add_node(self, ia_op_exec_context: InputAgnosticOperationExecutionContext,
                 tensor_metas: List[TensorMeta],
                 tm_comparators_per_scope: List[Tuple[TensorMetaComparator, List[str]]],
                 inputs) -> NNCFNode:
        matcher = self.choose_matcher(ia_op_exec_context)
        tm_comparators = self.choose_tm_comparators(ia_op_exec_context, tm_comparators_per_scope)
        op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name,
                                                    ia_op_exec_context.scope_in_model,
                                                    ia_op_exec_context.call_order,
                                                    tensor_metas,
                                                    tm_comparators=tm_comparators)

        return matcher.add_node(op_exec_context, inputs)


class NNCFGraphEdge:
    def __init__(self, from_node: NNCFNode, to_node: NNCFNode, tensor_shape: Tuple):
        self.from_node = from_node
        self.to_node = to_node
        self.tensor_shape = tensor_shape

    def __str__(self):
        return str(self.from_node) + " -> " + str(self.tensor_shape) + " -> " + str(self.to_node)

    def __hash__(self):
        return hash(str(self))

    def __eq__(self, other):
        return self.from_node == other.from_node and self.to_node == other.to_node \
               and self.tensor_shape == other.tensor_shape


class NNCFGraphPatternIO:
    def __init__(self, input_edges: List[NNCFGraphEdge], output_edges: List[NNCFGraphEdge],
                 input_nodes: List[NNCFNode],
                 output_nodes: List[NNCFNode],
                 ):
        self.input_edges = input_edges
        self.output_edges = output_edges
        self.input_nodes = input_nodes
        self.output_nodes = output_nodes


class NNCFGraph:
    ID_NODE_ATTR = "id"
    KEY_NODE_ATTR = "key"
    OP_EXEC_CONTEXT_NODE_ATTR = "op_exec_context"
    ACTIVATION_SHAPE_EDGE_ATTR = "activation_shape"
    IN_PORT_NAME = "in_port"

    def __init__(self):
        self._nx_graph = nx.DiGraph()
        self._node_id_to_key_dict = dict()
        self.match_manager = NodeManager(self._node_id_to_key_dict, self._nx_graph, self._nx_node_to_nncf_node)
        self._input_nncf_nodes = []

    def find_node(self,
                  ia_op_exec_context: InputAgnosticOperationExecutionContext,
                  tensor_metas: List[TensorMeta],
                  input_comparators_per_scope: List[Tuple[TensorMetaComparator, List[str]]]) -> NNCFNode:
        return self.match_manager.find_node(ia_op_exec_context, tensor_metas, input_comparators_per_scope)

    def add_node(self, ia_op_exec_context: InputAgnosticOperationExecutionContext,
                 tensor_metas: List[TensorMeta],
                 input_comparators_per_scope: List[Tuple[TensorMetaComparator, List[str]]],
                 inputs) -> NNCFNode:
        node = self.match_manager.add_node(ia_op_exec_context, tensor_metas, input_comparators_per_scope, inputs)

        from nncf.dynamic_graph.input_wrapping import MODEL_INPUT_OP_NAME
        if node.op_exec_context.operator_name == MODEL_INPUT_OP_NAME:  # TODO: refactorable model input node name
            self._input_nncf_nodes.append(node)
        return node

    def get_nx_node_by_key(self, key: str):
        return self._nx_graph.nodes[key]

    def get_node_id_by_iap_context(self, iap_ctx: InputAgnosticOperationExecutionContext) -> str:
        for node_key, node in self._nx_graph.nodes.items():
            if node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].input_agnostic == iap_ctx:
                return node_key
        raise AttributeError('Failed to get node by context={}'.format(str(iap_ctx)))

    def get_successors(self, node_name: str):
        return self._nx_graph.successors(node_name)

    def get_all_node_keys(self):
        return self._node_id_to_key_dict.copy().values()

    def get_all_node_idxs(self):
        return self._node_id_to_key_dict.keys()

    def get_node_key_by_id(self, node_id):
        return self._node_id_to_key_dict[node_id]

    def get_matching_nncf_graph_pattern_io_list(self, expression: Expression) -> List[NNCFGraphPatternIO]:
        matched_node_key_sequences = search_all(self._nx_graph, expression)
        pattern_ios = [self._get_nncf_graph_pattern_io_list(match) for match in matched_node_key_sequences]
        return pattern_ios

    def dump_graph(self, path):
        nx.drawing.nx_pydot.write_dot(self._get_graph_for_structure_analysis(), path)

    def visualize_graph(self, path):
        out_graph = self._get_graph_for_visualization()
        nx.drawing.nx_pydot.write_dot(out_graph, path)
        try:
            A = to_agraph(out_graph)
            A.layout('dot')
            png_path = os.path.splitext(path)[0]+'.png'
            A.draw(png_path)
        except ImportError:
            nncf_logger.warning("Graphviz is not installed - only the .dot model visualization format will be used. "
                                "Install pygraphviz into your Python environment and graphviz system-wide to enable "
                                "PNG rendering.")

    def is_output_node(self, node: NNCFNode) -> bool:
        return not list(self._nx_graph.successors(self._node_id_to_key_dict[node.node_id]))

    def get_nx_graph_copy(self) -> nx.DiGraph:
        return deepcopy(self._nx_graph)

    def get_input_nodes(self) -> List[NNCFNode]:
        return self._input_nncf_nodes

    def get_graph_outputs(self) -> List[NNCFNode]:
        outputs = []
        for nx_node_key, deg in self._nx_graph.out_degree():
            if deg == 0:
                outputs.append(self._nx_node_to_nncf_node(self._nx_graph.nodes[nx_node_key]))
        return outputs

    # pylint:disable=protected-access
    def get_nx_edge(self, node_u: NNCFNode, node_v: NNCFNode):
        nx_node_u = self._nx_graph._node[self._node_id_to_key_dict[node_u.node_id]]
        nx_node_v = self._nx_graph._node[self._node_id_to_key_dict[node_v.node_id]]
        return self._nx_graph.edges[nx_node_u['key'], nx_node_v['key']]

    def get_next_nodes(self, node: NNCFNode) -> Optional[List[NNCFNode]]:
        nx_node_keys = self._nx_graph.succ[self._node_id_to_key_dict[node.node_id]]
        return [self._nx_node_to_nncf_node(self._nx_graph.nodes[key]) for key in nx_node_keys]

    def get_previous_nodes(self, node: NNCFNode) -> Optional[List[NNCFNode]]:
        nx_node_keys = self._nx_graph.pred[self._node_id_to_key_dict[node.node_id]]
        return [self._nx_node_to_nncf_node(self._nx_graph.nodes[key]) for key in nx_node_keys]

    def get_inputs_count(self, node: NNCFNode) -> int:
        return self._nx_graph.in_degree()[self._node_id_to_key_dict[node.node_id]]

    def traverse_graph(self, curr_node: NNCFNode, traverse_function: Callable[[NNCFNode], Tuple[bool, List[Any]]],
                       traverse_forward: bool = True):
        output = []
        return self._traverse_graph_recursive_helper(curr_node, traverse_function, output, traverse_forward)

    def _traverse_graph_recursive_helper(self, curr_node: NNCFNode,
                                         traverse_function: Callable[[NNCFNode, List[Any]], Tuple[bool, List[Any]]],
                                         output: List[Any], traverse_forward: bool):
        is_finished, output = traverse_function(curr_node, output)
        get_nodes_fn = self.get_next_nodes if traverse_forward else self.get_previous_nodes
        if not is_finished:
            for node in get_nodes_fn(curr_node):
                self._traverse_graph_recursive_helper(node, traverse_function, output, traverse_forward)
        return output

    def get_nodes_count(self):
        return self._nx_graph.number_of_nodes()

    @staticmethod
    def node_type_fn(node: dict) -> str:
        return node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].operator_name

    def get_output_shapes_for_ia_op_exec_context(self,
                                                 ia_op_exec_context: InputAgnosticOperationExecutionContext)\
                                                 -> List[Tuple]:
        node_key = self.get_node_id_by_iap_context(ia_op_exec_context)
        succs = list(self._nx_graph.successors(node_key))
        edge_list = [self._nx_graph.edges[node_key, to_node_key] for to_node_key in succs]
        return [edge[NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR] for edge in edge_list]

    def _get_graph_for_structure_analysis(self, extended=False) -> nx.DiGraph:
        """The graph to dump has certain node attributes omitted, compared to the graph stored
         inside NNCFGraph."""
        out_graph = nx.DiGraph()
        for node_name, node in self._nx_graph.nodes.items():
            op_exec_context = node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR]
            scope_str = str(op_exec_context.scope_in_model)
            attrs_node = {
                'type': op_exec_context.operator_name,
                'id': node[NNCFGraph.ID_NODE_ATTR],
                'scope': scope_str
            }
            if 'color' in node:
                attrs_node['color'] = node['color']
            if 'label' in node:
                attrs_node['label'] = node['label']
            if 'style' in node:
                attrs_node['style'] = node['style']

            out_graph.add_node(node_name, **attrs_node)
        if extended:
            for u, v in self._nx_graph.edges:
                out_graph.add_edge(u, v, label=self._nx_graph.edges[u, v][NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR])
        else:
            for u, v in self._nx_graph.edges:
                out_graph.add_edge(u, v)

        return out_graph

    def _get_graph_for_visualization(self) -> nx.DiGraph:
        """A user-friendly graph .dot file, making it easier to debug the network and setup
        ignored/target scopes."""
        out_graph = nx.DiGraph()
        for node_name, node in self._nx_graph.nodes.items():
            op_exec_context = node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR]

            attrs_node = {}
            attrs_node['label'] = str(node[NNCFGraph.ID_NODE_ATTR]) + ' ' + str(op_exec_context.input_agnostic)

            out_graph.add_node(node_name, **attrs_node)

        for u, v in self._nx_graph.edges:
            out_graph.add_edge(u, v, label=self._nx_graph.edges[u, v][NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR])

        mapping = {k: v["label"] for k, v in out_graph.nodes.items()}
        out_graph = nx.relabel_nodes(out_graph, mapping)
        for node in out_graph.nodes.values():
            node.pop("label")

        return out_graph

    def _get_topologically_last_nodes(self, matches: List[List[str]]) -> List[str]:
        topological_order = {node: k for k, node in enumerate(nx.topological_sort(self._nx_graph))}
        insertion_points = {max(match, key=topological_order.__getitem__) for match in matches}
        for match in matches:
            for node in match:
                if len(list(self._nx_graph.successors(node))) > 1:
                    insertion_points.add(node)

        return list(insertion_points)

    def _get_nncf_graph_pattern_input_output(self, match: List[str]) -> NNCFGraphPatternIO:
        out_edge_boundary = list(nx.edge_boundary(self._nx_graph, match, data=True))
        complement = list(filter(lambda x: x not in match, self._nx_graph.nodes.keys()))
        in_edge_boundary = list(nx.edge_boundary(self._nx_graph, complement, data=True))
        boundary = in_edge_boundary + out_edge_boundary
        input_nncf_edges = []
        output_nncf_edges = []
        input_nncf_nodes = []
        output_nncf_nodes = []
        for key in match:
            # Currently we treat the nodes without incoming edges as "input" and the nodes without
            # outcoming edges as "output".
            # A proper way to find the input nodes would be to mark the tensors arriving at NNCFNetwork's
            # "forward" as input, then drop the marking once the first operation with an input tensor
            # has been done; the node corresponding to this operation would be "input" by definition.
            # Same with output nodes - should check the model output for TracedTensors and mark the
            # nodes from which such tensors originated as "output".
            # TODO: implement the functionality above.
            if not list(self._nx_graph.successors(key)):
                output_nncf_nodes.append(self._nx_node_to_nncf_node(self._nx_graph.nodes[key]))
            if not list(self._nx_graph.predecessors(key)):
                input_nncf_nodes.append(self._nx_node_to_nncf_node(self._nx_graph.nodes[key]))

        for nx_edge in boundary:
            from_node_key = nx_edge[0]
            to_node_key = nx_edge[1]
            data = nx_edge[2]
            nncf_edge = NNCFGraphEdge(self._nx_node_to_nncf_node(self._nx_graph.nodes[from_node_key]),
                                      self._nx_node_to_nncf_node(self._nx_graph.nodes[to_node_key]),
                                      data[NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR])
            if from_node_key in match:
                output_nncf_edges.append(nncf_edge)
            elif to_node_key in match:
                input_nncf_edges.append(nncf_edge)
            else:
                raise RuntimeError("Invalid graph expression supplied!")

        return NNCFGraphPatternIO(input_nncf_edges, output_nncf_edges,
                                  input_nncf_nodes, output_nncf_nodes)

    def _get_nncf_graph_pattern_io_list(self, match: List[str]) -> NNCFGraphPatternIO:
        in_edge_boundary, out_edge_boundary = get_edge_boundaries(match, self._nx_graph)
        boundary = in_edge_boundary + out_edge_boundary
        input_nncf_edges = []
        output_nncf_edges = []
        input_nncf_nodes = []
        output_nncf_nodes = []
        for key in match:
            # Currently we treat the nodes without incoming edges as "input" and the nodes without
            # outcoming edges as "output".
            # A proper way to find the input nodes would be to mark the tensors arriving at NNCFNetwork's
            # "forward" as input, then drop the marking once the first operation with an input tensor
            # has been done; the node corresponding to this operation would be "input" by definition.
            # Same with output nodes - should check the model output for TracedTensors and mark the
            # nodes from which such tensors originated as "output".
            # TODO: implement the functionality above.
            if not list(self._nx_graph.successors(key)):
                output_nncf_nodes.append(self._nx_node_to_nncf_node(self._nx_graph.nodes[key]))
            if not list(self._nx_graph.predecessors(key)):
                input_nncf_nodes.append(self._nx_node_to_nncf_node(self._nx_graph.nodes[key]))

        for nx_edge in boundary:
            from_node_key = nx_edge[0]
            to_node_key = nx_edge[1]
            data = nx_edge[2]
            nncf_edge = NNCFGraphEdge(self._nx_node_to_nncf_node(self._nx_graph.nodes[from_node_key]),
                                      self._nx_node_to_nncf_node(self._nx_graph.nodes[to_node_key]),
                                      data[NNCFGraph.ACTIVATION_SHAPE_EDGE_ATTR])
            if from_node_key in match:
                output_nncf_edges.append(nncf_edge)
            elif to_node_key in match:
                input_nncf_edges.append(nncf_edge)
            else:
                raise RuntimeError("Invalid graph expression supplied!")

        return NNCFGraphPatternIO(input_nncf_edges, output_nncf_edges,
                                  input_nncf_nodes, output_nncf_nodes)

    @staticmethod
    def _nx_node_to_nncf_node(nx_node) -> 'NNCFNode':
        return NNCFNode(nx_node[NNCFGraph.ID_NODE_ATTR],
                        nx_node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR])

    def find_node_in_nx_graph_by_scope(self, scope: 'Scope') -> Optional[dict]:
        """
        Looking for node with scope == scope in networkx graph.
        :param self: graphs to work on
        :param scope: Scope to find in graph
        :return: node from networkx graph for graph (or None if such scope is not found)
        """
        nodes = self._nx_graph.nodes
        for node_key in nodes:
            if nodes[node_key][NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].scope_in_model == scope:
                return nodes[node_key]
        return None

    def get_op_nodes_in_scope(self, scope: 'Scope') -> List:
        matching_graph_op_nodes = []
        for _, node in self._nx_graph.nodes.items():
            op_scope = node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].input_agnostic.scope_in_model
            if op_scope in scope:
                matching_graph_op_nodes.append(node)
        return matching_graph_op_nodes

class NNCFNodeExpression(NodeExpression):
    def __init__(self, node_type: str = None, filter_fn=None):
        super().__init__(node_type, filter_fn, node_type_fn=NNCFGraph.node_type_fn)
