"""
 Copyright (c) 2018-2019 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 numpy as np

from copy import deepcopy

from extensions.middle.ConvertGroupedStridedSlice import ConvertGroupedStridedSlice
from extensions.middle.FusePermutesSequence import FusePermutesSequence
from extensions.middle.ShufflenetReshape import ReshapeSoftmaxReshape
from mo.graph.graph import Graph
from mo.middle.replacement import MiddleReplacementPattern
from mo.ops.op import Op
from mo.ops.permute import Permute


class PixelLinkReshape(MiddleReplacementPattern):
    """
      Transform adds Permutes around Reshapes that pack 4 dimensions in 2, than
      do Softmax and then unpack it back to 5 dims.
    """
    enabled = True

    def run_before(self):
        return [FusePermutesSequence, ReshapeSoftmaxReshape, ConvertGroupedStridedSlice]

    def run_after(self):
        from extensions.middle.pass_separator import MiddleStart
        return [MiddleStart]

    def pattern(self):
        return dict(nodes=[('reshape_split', dict(kind='op', type='Reshape')),
                           ('reshape_split_data', dict(kind='data')),
                           ('reshape_pack', dict(kind='op', type='Reshape')),
                           ('reshape_data', dict(kind='data')),
                           ('softmax', dict(kind='op', type='SoftMax')),
                           ('softmax_data', dict(kind='data')),
                           ('reshape_unpack', dict(kind='op', type='Reshape')),
                           ('reshape_unpack_data', dict(kind='data')),
                           ('strided_slice', dict(kind='op', op='StridedSlice')),
                           ],
                    edges=[('reshape_split', 'reshape_split_data'),
                           ('reshape_split_data', 'reshape_pack'),
                           ('reshape_pack', 'reshape_data'),
                           ('reshape_data', 'softmax'),
                           ('softmax', 'softmax_data'),
                           ('softmax_data', 'reshape_unpack'),
                           ('reshape_unpack', 'reshape_unpack_data'),
                           ('reshape_unpack_data', 'strided_slice')])

    def is_reshape_bad(self, node_pack, node_unpack, node_ss):
        shape_in = node_pack.in_node(0).shape
        shape_out = node_pack.out_node(0).shape

        if len(shape_in) == 5 and len(shape_out) == 2:
            shape_in = node_unpack.in_node(0).shape
            shape_out = node_unpack.out_node(0).shape
            if len(shape_out) == 5 and len(shape_in) == 2:
                if node_ss.slices[1].stop == shape_out[1] and node_ss.slices[1].start == 0 and node_ss.slices[
                    1].step == 1 and \
                        node_ss.slices[2].stop == shape_out[2] and node_ss.slices[2].start == 0 and node_ss.slices[
                    2].step == 1 and \
                        node_ss.slices[3].stop == shape_out[3] and node_ss.slices[3].start == 0 and node_ss.slices[
                    3].step == 1 and \
                        node_ss.shrink_axis_mask[4] and \
                        not node_ss.shrink_axis_mask[1] and not node_ss.shrink_axis_mask[2] and not \
                        node_ss.shrink_axis_mask[3]:
                    return True
                else:
                    return False
        else:
            return False

    def replace_pattern(self, graph: Graph, match: dict):
        if graph.graph['layout'] != 'NHWC':
            return

        if self.is_reshape_bad(match['reshape_pack'], match['reshape_unpack'], match['strided_slice']):
            log.info("Reshape that pack/unpack several dimensions detected {}".format(match['reshape_pack'].id))
            node_split = match['reshape_split']

            # insert Permute before reshape
            data_node = Op._create_data_node(graph, node_split.name + "/Permute_before_data")
            permute_before = Permute(graph, dict(name=node_split.name + "/Permute_before",
                                                 order=np.array([0, 2, 3, 1])))
            in_node = node_split.in_node(0)
            attrs = deepcopy(graph.get_edge_data(in_node.id, node_split.id)[0])
            graph.remove_edge(in_node.id, node_split.id)
            permute_before_node = permute_before.create_node_with_data([in_node], permute_before.attrs,
                                                                       data_nodes=[data_node])
            graph.add_edge(permute_before_node.id, node_split.id, **attrs)

            node = match['reshape_pack']
            node['nchw_layout'] = True
            new_reshape_shape = np.concatenate((np.array([node.in_node(0).shape[0]]),
                                                np.array([np.prod(node.in_node(0).shape[[1, 2, 3]])]),
                                                np.array([node.in_node(0).shape[-1]])))

            node.dim = new_reshape_shape

            # insert Permute after reshape
            data_node = Op._create_data_node(graph, node.name + "/Permute_after_data", {'shape': node.dim})
            permute_after = Permute(graph, dict(name=node.name + "/Permute_after",
                                                order=np.array([0, 2, 1])))
            out_node = node.out_node(0)
            out_node.shape = new_reshape_shape[np.array([0, 2, 1])]
            attrs = deepcopy(graph.get_edge_data(node.id, out_node.id)[0])
            graph.remove_edge(node.id, out_node.id)

            permute_after.create_node_with_data([data_node], permute_after.attrs,
                                                data_nodes=[out_node])
            graph.add_edge(node.id, data_node.id, **attrs)

            # update softmax shape
            node_softmax = match['softmax']
            node_softmax.out_node(0).shape = out_node.shape

            if ConvertGroupedStridedSlice.enabled is True:
                # revert strided slice and reshape
                node_ss = match['strided_slice']
                node_unpack = match['reshape_unpack']

                unpack_out = node_unpack.out_node(0).id
                ss_out = node_ss.out_node(0).id

                # gather edge attributes
                soft_reshape_attrs = deepcopy(graph.get_edge_data(node_softmax.out_node(0).id, node_unpack.id)[0])
                reshape_data_attrs = deepcopy(graph.get_edge_data(node_unpack.id, unpack_out)[0])
                reshape_ss_attrs = deepcopy(graph.get_edge_data(unpack_out, node_ss.id)[0])
                ss_out_attrs = deepcopy(graph.get_edge_data(node_ss.id, ss_out)[0])

                # remove all edges in Softmax->Reshape->StridedSlice chain
                graph.remove_edge(node_softmax.out_node(0).id, node_unpack.id)
                graph.remove_edge(node_unpack.id, unpack_out)
                graph.remove_edge(unpack_out, node_ss.id)
                graph.remove_edge(node_ss.id, ss_out)

                # add new edges to get chain Softmax->StridedSlice->Reshape
                graph.add_edge(node_softmax.out_node(0).id, node_ss.id, **soft_reshape_attrs)
                graph.add_edge(node_ss.id, unpack_out, **reshape_data_attrs)
                graph.add_edge(unpack_out, node_unpack.id, **reshape_ss_attrs)
                graph.add_edge(node_unpack.id, ss_out, **ss_out_attrs)

                # update output shape and parameters for StridedSlice
                node_ss.out_node(0).shape = np.zeros(3)
                node_ss.out_node(0).shape[0] = out_node.shape[0]
                node_ss.out_node(0).shape[1] = 1
                node_ss.out_node(0).shape[2] = out_node.shape[2]

                old_slices = node_ss.slices.copy()
                node_ss.slices = []
                node_ss.slices.append(old_slices[0])
                node_ss.slices.append(old_slices[-1])
                node_ss.slices.append(slice(0, out_node.shape[2], 1))
                node_ss.shrink_axis_mask = np.array([0, 0, 0], dtype=np.int64)
                node_ss.new_axis_mask = np.array([0, 0, 0], dtype=np.int64)
                node_ss.ellipsis_mask = np.array([0, 0, 0], dtype=np.int64)
                node_ss.begin_mask = np.array([0, 1, 0], dtype=np.int64)
                node_ss.end_mask = np.array([0, 1, 0], dtype=np.int64)

                # update Reshape attribute
                node_unpack.dim = np.delete(node_unpack.dim, 4)
                # prevent permute for reshape because it gives wrong result
                node_unpack['nchw_layout'] = True
                node_unpack.out_node(0)['nchw_layout'] = True
            else:
                # reshape unpack: permute correctly
                node_unpack = match['reshape_unpack']
                data_node = Op._create_data_node(graph, node.name + "/Permute_after_unpack_data", {'shape': node_unpack.out_node().shape})
                permute_after_unpack = Permute(graph, dict(name=node.name + "/Permute_after_unpack",
                                                           order=np.array([0, 3, 1, 2, 4])))
                out_node = node_unpack.out_node(0)
                out_node.shape = out_node.shape[np.array([0, 3, 1, 2, 4], dtype=np.int)]
                attrs = deepcopy(graph.get_edge_data(node_unpack.id, out_node.id)[0])
                graph.remove_edge(node_unpack.id, out_node.id)
                permute_after.create_node_with_data([data_node], permute_after_unpack.attrs,
                                                    data_nodes=[out_node])
                graph.add_edge(node_unpack.id, data_node.id, **attrs)
