# Copyright (C) 2018-2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from mo.front.common.replacement import FrontReplacementPattern
from mo.front.extractor import add_output_ops
from mo.graph.graph import Graph, get_edge_attribute_between_nodes, set_edge_attribute_between_nodes


class OutputCut(FrontReplacementPattern):
    enabled = True
    run_not_recursively = True
    force_clean_up = True

    def run_after(self):
        from extensions.front.user_data_repack import UserDataRepack
        return [UserDataRepack]

    def run_before(self):
        return []

    def find_and_replace_pattern(self, graph: Graph):
        add_output_ops(graph, graph.graph['packed_outputs'], inputs=graph.graph['user_shapes'])

        # For keeping tensor names information for output nodes fake outputs are added
        # to graph during the model loading. In the following code fake outputs are removed
        # and tensor names information is moved to output->Result edge.
        for node in graph.get_op_nodes(needs_removal=True):
            fw_info = None
            in_node = None
            for in_port_idx in node.in_edges():
                node_idx = node.in_edge(in_port_idx)['in']
                if node_idx in node.in_nodes():
                    in_node = node.in_node(node_idx)
                    fw_info_value = get_edge_attribute_between_nodes(in_node, node, 'fw_tensor_debug_info')
                    if fw_info_value:
                        fw_info = fw_info_value
                        break
            graph.erase_node(node)

            if fw_info is not None and in_node is not None:
                for out_idx in in_node.out_nodes():
                    set_edge_attribute_between_nodes(in_node, in_node.out_node(out_idx),
                                                     'fw_tensor_debug_info', fw_info)
