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

from collections import deque
from typing import List, Set

from extensions.middle.InsertLayoutPropagationTransposes import is_output_data_in_correct_layout, \
    InsertLayoutPropagationTranspose, mark_input_as_in_correct_layout, mark_output_as_in_correct_layout
from extensions.ops.gather import Gather
from extensions.ops.transpose import Transpose
from mo.front.common.partial_infer.utils import int64_array
from mo.front.tf.graph_utils import create_op_with_const_inputs
from mo.graph.graph import Graph
from mo.graph.port import Port
from mo.middle.replacement import MiddleReplacementPattern


class LayoutChangeForConstantShapePaths(MiddleReplacementPattern):
    enabled = True
    graph_condition = [lambda graph: graph.graph['layout'] == 'NHWC',
                       lambda graph: not graph.graph['cmd_params'].static_shape]
    force_clean_up = True

    def run_after(self):
        return [InsertLayoutPropagationTranspose]

    def run_before(self):
        return []

    @staticmethod
    def get_next_in_ports(in_port: Port) -> Set[Port]:
        next_in_ports = set()
        for out_port in in_port.node.out_ports().values():
            next_in_ports.update(out_port.get_destinations())
        return next_in_ports

    def find_shape_subgraph_endpoints(self, out_ports: List[Port], visited: set = None,
                                      action: callable = None) -> Set[Port]:
        """
        Searches for input ports of data dependent operations starting from output ports passed to the function.
        Condition for data dependent operations is absence of node output value.

        :param out_ports: list of output ports to start search from
        :param visited: set of input ports that were visited to avoid visiting them more than once
        :param action: function to call on the each input port of shape sub-graph
        :return: set of input ports of data dependent operations
        """
        if visited is None:
            visited = set()

        deque_of_in_ports = deque()
        for out_port in out_ports:
            deque_of_in_ports.extend(out_port.get_destinations())

        end_points_in_ports = set()
        while len(deque_of_in_ports):
            in_port = deque_of_in_ports.popleft()
            if in_port in visited:
                continue
            next_in_ports = self.get_next_in_ports(in_port)
            if any([port.data.get_value() is None for port in next_in_ports]):
                end_points_in_ports.add(in_port)
            else:
                deque_of_in_ports.extend(next_in_ports)
                if action is not None:
                    action(in_port)
            visited.add(in_port)
        return end_points_in_ports

    def find_and_replace_pattern(self, graph: Graph):
        shape_ops = graph.get_op_nodes(op='ShapeOf')

        # 1. Inserting Gather to N*C format on constant shape paths
        for shape in shape_ops:
            source_port = shape.in_port(0).get_source()
            if is_output_data_in_correct_layout(source_port.node, source_port.idx):
                continue  # data is already in N*C format

            name = shape.soft_get('name', shape.id)
            rank = source_port.data.get_shape().size

            if rank in [4, 5]:
                index = int64_array([0, *list(range(2, rank)), 1])
            else:
                continue  # data is layout independent

            gather = create_op_with_const_inputs(graph, op=Gather, port_value_dict={1: index, 2: int64_array(0)},
                                                 op_attrs={'name': name + '/GatherNCHWtoNHWC'})
            shape.out_port(0).get_connection().insert_node(gather)

        # 2. Inserting Gather/Transpose to NC* format
        shape_sub_graph_end_points = self.find_shape_subgraph_endpoints([shape.out_port(0) for shape in shape_ops])
        for in_port in shape_sub_graph_end_points:
            name = in_port.node.soft_get('name', in_port.node.id)
            shape = in_port.data.get_shape()

            should_switch_layout = not any([is_output_data_in_correct_layout(port.node, port.idx)
                                            for port in in_port.node.out_ports().values() if not port.disconnected()])
            should_insert_gather = should_switch_layout and len(shape) == 1 and shape.item(0) in [4, 5]
            should_insert_transpose = should_switch_layout and len(shape) in [4, 5]

            if should_insert_gather:
                # we should turn input permutation off to perform it with the following gather insertion
                in_port.__setattr__('input_permutation', None)
                index = int64_array([0, shape.item(0) - 1, *list(range(1, shape.item(0) - 1))])
                gather = create_op_with_const_inputs(graph, op=Gather,
                                                     port_value_dict={1: index, 2: int64_array(0)},
                                                     op_attrs={'name': name + '/GatherNHWCtoNCHW'})
                in_port.get_connection().insert_node(gather)
            elif should_insert_transpose:
                # we should turn input permutation off to perform it with the following transpose insertion
                in_port.__setattr__('input_permutation', None)
                order = int64_array([0, len(shape) - 1, *list(range(1, len(shape) - 1))])
                transpose = create_op_with_const_inputs(graph, op=Transpose, port_value_dict={1: order},
                                                        op_attrs={'name': name + '/TransposeNHWCtoNCHW',
                                                                  'override_output_shape': True})
                mark_input_as_in_correct_layout(transpose, 0)
                mark_output_as_in_correct_layout(transpose, 0)
                in_port.get_connection().insert_node(transpose)
            else:
                continue  # data is layout independent
