"""
 Copyright (C) 2018-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.
"""

import logging as log
import math

import numpy as np

from extensions.back.FuseTransposesSequence import FuseTransposesSequence
from extensions.middle.FuseReshapesSequence import FuseReshapesSequence
from extensions.middle.RemoveRedundantReshapes import RemoveRedundantReshapes
from mo.back.replacement import BackReplacementPattern
from mo.front.common.partial_infer.utils import int64_array
from mo.front.tf.graph_utils import create_op_node_with_second_input
from mo.graph.graph import Graph, Node
from mo.middle.passes.fusing.helpers import get_next_operation
from mo.ops.op import PermuteAttrs
from mo.ops.reshape import Reshape


def split_input_permute_dimension(dim: int, permute_order: np.array):
    """
    Creates updated permutation for a given permutation order and the *input* dimension index to be split into two.
    :param dim: the input tensor dimension axis to split
    :param permute_order: the permutation order
    :return: the new permutation order after split of the specified dimension into two
    """
    assert dim < len(permute_order)
    new_permute_order = list()
    for permute_index in permute_order:
        if permute_index < dim:
            new_permute_order.append(permute_index)
        elif permute_index > dim:
            new_permute_order.append(permute_index + 1)
        else:
            new_permute_order.append(permute_index)
            new_permute_order.append(permute_index + 1)
    return int64_array(new_permute_order)


def split_output_permute_dimension(dim: int, permute_order: np.array):
    """
    Creates updated permutation for a given permutation order and the *output* dimension index to be split into two.
    :param dim: the output tensor dimension axis to split
    :param permute_order: the permutation order
    :return: the new permutation order after split of the specified dimension into two
    """
    assert dim < len(permute_order)
    return split_input_permute_dimension(permute_order[dim], permute_order)


def match_shapes(input_shape: np.array, output_shape: np.array):
    """
    Calculates "match" shape for the given input and output shape of the Reshape layer. The function splits some of the
    input/output dimensions into several ones to make new input and output shapes equal. For example,
    input_shape=(1,32,64,60)->Reshape->output_shape=(8,4,64,20,3) is converted to
    input_shape=(1,8,4,64,20,3)->Reshape->output_shape=(1,8,4,64,20,3).

    :param input_shape: input shape of the Reshape
    :param output_shape: output shape of the Reshape
    :return: "match" shape or None if it is not possible to calculate match shape
    """
    matched_shape = list()

    in_ind = 0
    out_ind = 0
    in_left = input_shape[0]
    out_left = output_shape[0]
    while in_ind < len(input_shape) or out_ind < len(output_shape):
        if in_ind < len(input_shape) and out_ind < len(output_shape):
            g = math.gcd(in_left, out_left)
            matched_shape.append(g)
            if g == 1 and in_left != 1 and out_left != 1:  # shapes cannot be matched
                return None
            in_left //= g
            out_left //= g
            if in_left == 1:
                in_ind += 1
                if in_ind < len(input_shape):
                    in_left *= input_shape[in_ind]
            if out_left == 1:
                out_ind += 1
                if out_ind < len(output_shape):
                    out_left *= output_shape[out_ind]
        else:
            matched_shape.append(1)
            if out_ind != len(output_shape):
                out_ind += 1
            else:
                in_ind += 1
    return int64_array(matched_shape)


def split_dims_indices(input_shape: np.array, match_shape: np.array):
    """
    Returns list of indices of the input shape to be split to match the match_shape shape
    :param input_shape: input shape
    :param match_shape: match shape
    :return: list of indices (indices may be repetitive)
    """
    result = list()
    in_ind = 0
    match_ind = 0
    in_left = input_shape[0]
    while match_ind < len(match_shape):
        if in_ind >= len(input_shape):
            assert match_shape[match_ind] == 1, 'Total number of elements in input shape and output shape are not equal'
            match_ind += 1
            result.append(in_ind - 1)
        elif match_shape[match_ind] == input_shape[in_ind] and match_shape[match_ind] == 1:
            match_ind += 1
            in_ind += 1
            if in_ind < len(input_shape):
                in_left *= input_shape[in_ind]
        elif in_left > match_shape[match_ind]:
            if in_left > match_shape[match_ind] or match_shape[match_ind] == 1:
                result.append(in_ind)
            in_left //= match_shape[match_ind]
            match_ind += 1
        elif in_left == match_shape[match_ind]:
            in_ind += 1
            match_ind += 1
            if in_ind < len(input_shape):
                in_left = input_shape[in_ind]
        else:
            in_ind += 1
            in_left *= input_shape[in_ind]
    return result


def reverse_permute(output_shape: np.array, order: np.array):
    """
    Calculates Transpose op input shape based on output shape and permute order.
    :param output_shape: Transpose output shape
    :param order: permute order
    :return: Transpose input shape corresponding to the specified output shape
    """
    return int64_array(output_shape[PermuteAttrs.get_inverse_permutation(order)])


def set_reshape_new_output_shape(reshape_node: Node, new_output_shape: np.array):
    """
    Updates Reshape node shape to a new output shape. The function updates the second input if the node has it.
    :param reshape_node: node to update
    :param new_output_shape: new output shape
    :return: None
    """
    reshape_node.out_port(0).data.set_shape(new_output_shape)
    in_ports = [port for port in reshape_node.in_ports().values() if not port.disconnected()]
    if len(in_ports) == 2:
        reshape_node.in_port(1).data.set_value(new_output_shape)


class OptimizeTransposeReshapeSequence(BackReplacementPattern):
    """
    The transformation looks for the sequence of Reshape and Transpose operations and tries to optimize it. The procedure
    is the following:

    1. For each Reshape layer in the sequence of nodes being optimized (except leading and trailing one) make it dummy,
    i.e. not changing the input and output shape. For example, convert
    input_shape=(1,32,64,60)->Reshape->output_shape=(8,4,64,20,3) to
    input_shape=(1,8,4,64,20,3)->Reshape->output_shape=(1,8,4,64,20,3).
    2. Propagate new input/output shape forward and backward through the Transpose nodes.
    3. Remove dummy Reshapes.
    4. Fuse sequence of Transposes.
    """
    enabled = False
    run_not_recursively = True
    OPTIMIZED_NODE_FLAG = 'permute_reshape_optimized'

    def run_before(self):
        from extensions.back.ReshapeMutation import ReshapeMutation
        return [ReshapeMutation]

    def run_after(self):
        return [FuseTransposesSequence]

    def is_node_match_for_optimization(self, node: Node):
        """
        Check that the node can be added to the sequence of nodes for the Transpose-Reshape optimization
        :param node: node to check
        :return: result of the check
        """
        # TODO change to 'op' and reshape-like
        return node.has_and_set('type') and node.type in ('Transpose', 'Reshape') and \
            not node.has_and_set(self.OPTIMIZED_NODE_FLAG)

    def find_and_replace_pattern(self, graph: Graph):
        for start_node in graph.pseudo_topological_sort():
            matched_nodes = []
            if self.is_node_match_for_optimization(start_node):
                next_node = start_node
                while self.is_node_match_for_optimization(next_node):
                    matched_nodes.append(next_node)
                    next_node[self.OPTIMIZED_NODE_FLAG] = True
                    next_nodes = get_next_operation(next_node)
                    if len(next_nodes) > 1:
                        log.debug('There are two consumers of the node {}. Stop matching sequence.'.format(
                            next_node.soft_get('name')))
                        break
                    next_node = next_nodes[0]
            # optimize sequence of three or more Transpose-Reshape nodes
            if len(matched_nodes) >= 3:
                self.optimize_permute_reshape_sequence(graph, matched_nodes)

        # run the RemoveRedundantReshapes to remove dummy (NOP) reshapes. After that we can run Transposes fusing
        FuseReshapesSequence().find_and_replace_pattern(graph)
        RemoveRedundantReshapes().find_and_replace_pattern(graph)
        FuseTransposesSequence().find_and_replace_pattern(graph)

    @staticmethod
    def optimize_permute_reshape_sequence(graph: Graph, nodes: list):
        log.debug('Running permute-reshape optimization of the following nodes: {}'.format(
            [node.soft_get('name') for node in nodes]))

        # the transformation expects that the first and the last operation in the sequence is Reshape so the following
        # function adds required reshapes
        __class__.add_leading_and_trailing_reshape(graph, nodes)

        for ind in range(1, len(nodes) - 1):
            node = nodes[ind]
            input_shape = node.in_node(0).shape
            output_shape = node.out_node(0).shape
            if node.type == 'Reshape' and not np.array_equal(input_shape, output_shape):
                log.debug('The Reshape node "{}" is not NOP. Shapes: "{}" vs "{}"'.format(
                    node.soft_get('name'), input_shape, output_shape))
                __class__.make_reshape_nop(node)

    @staticmethod
    def add_leading_and_trailing_reshape(graph: Graph, nodes: list):
        """
        When the first operation in the matched list is the Transpose then add the Reshape operation which reshapes to the
        Transpose input shape. This Reshape op is needed for the optimization pass. If the optimization will not be
        applied then this dummy Reshape will be removed by the "RemoveRedundantReshapes" pass.

        :param graph: the graph with nodes
        :param nodes: the sequence of Transpose and ReshapeFF nodes
        :return: None
        """
        # add leading Reshape
        if nodes[0].type == 'Transpose':
            dummy_reshape_node = create_op_node_with_second_input(
                graph, Reshape, nodes[0].in_port(0).data.get_shape().copy(),
                {'name': nodes[0].in_port(0).get_connection().get_source().node.id + '/Reshape'})
            dummy_reshape_node[__class__.OPTIMIZED_NODE_FLAG] = True
            nodes[0].in_port(0).get_connection().insert_node(dummy_reshape_node)
            nodes.insert(0, dummy_reshape_node)
            log.debug('Added Reshape op "{}" in the beginning of the permute-reshape sequence'.format(
                dummy_reshape_node.soft_get('name')))

        # similarly add the Reshape op after the last Transpose op which reshapes to the Transpose output shape
        if nodes[-1].type == 'Transpose':
            dummy_reshape_node = create_op_node_with_second_input(
                graph, Reshape, nodes[-1].out_port(0).data.get_shape().copy(),
                {'name': nodes[0].out_port(0).get_connection().get_destination().node.id + '/Reshape'})
            dummy_reshape_node[__class__.OPTIMIZED_NODE_FLAG] = True
            nodes[-1].out_port(0).get_connection().insert_node(dummy_reshape_node)
            nodes.append(dummy_reshape_node)
            log.debug('Added Reshape op "{}" in the end of the permute-reshape sequence'.format(
                dummy_reshape_node.soft_get('name')))

    @staticmethod
    def forward_new_reshape_shape(reshape_node: Node, initial_output_shape: np.array):
        """
        Propagates the changed output shape of the Reshape node forward. The output of the Reshape node should be
        Transpose so it is necessary to update its 'order' attribute according to the updated shape and output data node.
        :param reshape_node: the Reshape node to propagate the shape
        :param initial_output_shape: old output shape of the Reshape node
        :return: None
        """
        output_shape = reshape_node.out_port(0).data.get_shape()
        if np.all(output_shape == initial_output_shape):
            log.debug('Initial output and new output shapes match for node "{}". Do nothing'.format(
                reshape_node.soft_get('name')))
            return

        dest_node = reshape_node.out_port(0).get_destination().node
        if dest_node.type == 'Transpose':
            split_dims = split_dims_indices(initial_output_shape, output_shape)
            assert dest_node.in_port(1).data.get_value() is not None, \
                'The 1st input value "order" is not set for Transpose node "{}"'.format(dest_node.soft_get('name'))
            permute_order = dest_node.in_port(1).data.get_value()
            for split_dim in split_dims:
                permute_order = split_input_permute_dimension(split_dim, permute_order)
            dest_node.in_port(1).data.set_value(permute_order)
            dest_node.infer(dest_node)
        elif dest_node.type == 'Reshape':
            log.debug('Two subsequent reshape nodes: "{}" and "{}". Nothing to optimize'.format(
                reshape_node.soft_get('name'), dest_node.soft_get('name')))
        else:
            assert False, 'Unsupported type of the node "{}" in the Transpose-Reshape optimization' \
                          ''.format(dest_node.type)

    @staticmethod
    def backward_new_reshape_shape(reshape_node: Node, initial_input_shape: np.array):
        """
        Propagates the changed input shape of the Reshape node backward.
        1. The input of the Reshape node should be Transpose so it is necessary to update its 'order' attribute according
        to the updated shape and input data node.
        2. The input of the Transpose should be a Reshape node, so it is necessary to update its 'dim' attribute.

        :param reshape_node: the Reshape node to propagate the shape
        :param initial_input_shape: old input shape of the Reshape node
        :return: None
        """
        input_shape = reshape_node.in_port(0).data.get_shape()
        if np.all(input_shape == initial_input_shape):
            log.debug('Initial input and new input shapes match for node "{}". Do nothing'.format(
                reshape_node.soft_get('name')))
            return

        src_node = reshape_node.in_port(0).get_source().node
        if src_node.type == 'Transpose':
            split_dims = split_dims_indices(initial_input_shape, input_shape)
            assert src_node.in_port(1).data.get_value() is not None, \
                'The 1st input value "order" is not set for Transpose node "{}"'.format(src_node.soft_get('name'))
            permute_order = src_node.in_port(1).data.get_value()
            for split_dim in split_dims:
                permute_order = split_output_permute_dimension(split_dim, permute_order)
            src_node.in_port(1).data.set_value(permute_order)

            # calculate a Transpose input shape based on the Transpose output shape
            new_permute_input_shape = reverse_permute(input_shape, permute_order)

            # update the Transpose input node (it should be Reshape) output shape and 'dim' attribute
            permute_source_port = src_node.in_port(0).get_source()
            permute_source_port.data.set_shape(new_permute_input_shape)
            set_reshape_new_output_shape(permute_source_port.node, new_permute_input_shape)
        elif src_node.type == 'Reshape':
            log.debug('Two subsequent reshape nodes: "{}" and "{}". Nothing to optimize'.format(
                reshape_node.soft_get('name'), src_node.soft_get('name')))
        else:
            assert False, 'Unsupported type of the node "{}" in the Transpose-Reshape optimization' \
                          ''.format(src_node.type)

    @staticmethod
    def make_reshape_nop(reshape_node: Node):
        """
        Change the node input and output shape so the Reshape node becomes dummy (NOP). Then propagate new shapes back
        and forth.

        :param reshape_node: reshape node to make it dummy
        :return: None
        """
        initial_input_shape = reshape_node.in_port(0).data.get_shape().copy()
        initial_output_shape = reshape_node.out_port(0).data.get_shape().copy()

        # calculate new shape which makes the Reshape NOP
        match_shape = match_shapes(initial_input_shape, initial_output_shape)
        if match_shape is None:  # it is not possible to optimize reshape
            return

        # update Reshape node and input/output attrs
        reshape_node.in_port(0).data.set_shape(match_shape)
        set_reshape_new_output_shape(reshape_node, match_shape)

        # propagate forward a new reshape shape by updating Transpose op consumer attributes and the following data node
        __class__.forward_new_reshape_shape(reshape_node, initial_output_shape)

        # propagate backward a new shape
        __class__.backward_new_reshape_shape(reshape_node, initial_input_shape)
