"""
 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 unittest

import numpy as np

from extensions.middle.SliceConverter import ConvertSlice
from mo.front.common.partial_infer.utils import int64_array
from mo.utils.ir_engine.compare_graphs import compare_graphs
from mo.utils.unittest.graph import build_graph, regular_op_with_shaped_data, valued_const_with_data, \
    regular_op_with_empty_data, result, connect, connect_data

nodes_attributes = {
    **regular_op_with_shaped_data('input', [2, 3, 300, 300], {'type': 'Parameter', 'op': 'Parameter'}),
    **regular_op_with_empty_data('starts', {'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('ends', {'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('axes', {'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('steps', {'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('slice', {'op': 'Slice', 'type': None}),

    **regular_op_with_empty_data('ss_begin_cast', {'op': 'Cast', 'type': 'Convert', 'dst_type': np.int64}),
    **regular_op_with_empty_data('ss_begin_clamp', {'op': 'Clamp', 'type': None}),
    **regular_op_with_empty_data('ss_begin_clamp_min', {'value': np.iinfo(np.int32).min, 'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('ss_begin_clamp_max', {'value': np.iinfo(np.int32).max, 'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('ss_begin_gather_0', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_begin_gather_0_idx', int64_array([0])),
    **regular_op_with_shaped_data('ss_begin_gather_0_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_begin_gather_1', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_begin_gather_1_idx', int64_array([1])),
    **regular_op_with_shaped_data('ss_begin_gather_1_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_begin_gather_2', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_begin_gather_2_idx', int64_array([2])),
    **regular_op_with_shaped_data('ss_begin_gather_2_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_begin_gather_3', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_begin_gather_3_idx', int64_array([3])),
    **regular_op_with_shaped_data('ss_begin_gather_3_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_begin_const_0', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_begin_const_1', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_begin_const_2', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_begin_const_3', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_begin_concat', {'op': 'Concat', 'type': 'Concat'}),

    **regular_op_with_empty_data('ss_end_cast', {'op': 'Cast', 'type': 'Convert', 'dst_type': np.int64}),
    **regular_op_with_empty_data('ss_end_clamp', {'op': 'Clamp', 'type': None}),
    **regular_op_with_empty_data('ss_end_clamp_min', {'value': np.iinfo(np.int32).min, 'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('ss_end_clamp_max', {'value': np.iinfo(np.int32).max, 'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('ss_end_gather_0', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_end_gather_0_idx', int64_array([0])),
    **regular_op_with_shaped_data('ss_end_gather_0_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_end_gather_1', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_end_gather_1_idx', int64_array([1])),
    **regular_op_with_shaped_data('ss_end_gather_1_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_end_gather_2', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_end_gather_2_idx', int64_array([2])),
    **regular_op_with_shaped_data('ss_end_gather_2_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_end_gather_3', {'op': 'Gather', 'type': 'Gather'}),
    **valued_const_with_data('ss_end_gather_3_idx', int64_array([3])),
    **regular_op_with_shaped_data('ss_end_gather_3_axis', [], {'op': 'Const', 'type': 'Const', 'value': [0]}),
    **regular_op_with_empty_data('ss_end_const_0', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_end_const_1', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_end_const_2', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_end_const_3', {'op': 'Const', 'type': 'Const', 'value': int64_array([0])}),
    **regular_op_with_empty_data('ss_end_concat', {'op': 'Concat', 'type': 'Concat'}),

    **regular_op_with_empty_data('ss_strides', {'op': 'Const', 'type': 'Const'}),
    **regular_op_with_empty_data('ss', {'op': 'StridedSlice', 'type': 'StridedSlice',
                                        'new_axis_mask': np.zeros(4, dtype=np.int64),
                                        'shrink_axis_mask': np.zeros(4, dtype=np.int64),
                                        'ellipsis_mask': np.zeros(4, dtype=np.int64)}),
    **result('result')
}

pattern_graph = [
    *connect('input:0', '0:slice'),
    *connect('starts:0', '1:slice'),
    *connect('ends:0', '2:slice'),
    *connect('axes:0', '3:slice'),
    *connect('steps:0', '4:slice'),
    *connect('slice:0', '0:result')
]

pattern_ref_graph = [
    *connect('input:0', '0:ss'),
    *connect('starts:0', '0:ss_begin_clamp'),
    *connect('ss_begin_clamp:0', '0:ss_begin_cast'),
    *connect('ss_begin_clamp_min:0', '1:ss_begin_clamp'),
    *connect('ss_begin_clamp_max:0', '2:ss_begin_clamp'),
    *connect('ss_begin_concat:0', '1:ss'),
    *connect('ends:0', '0:ss_end_clamp'),
    *connect('ss_end_clamp:0', '0:ss_end_cast'),
    *connect('ss_end_clamp_min:0', '1:ss_end_clamp'),
    *connect('ss_end_clamp_max:0', '2:ss_end_clamp'),
    *connect('ss_end_concat:0', '2:ss'),
    *connect('ss_strides:0', '3:ss'),
    *connect('ss:0', '0:result'),

    *connect('ss_begin_gather_0_idx:0', '1:ss_begin_gather_0'),
    *connect('ss_begin_gather_0_axis:0', '2:ss_begin_gather_0'),
    *connect('ss_begin_gather_1_idx:0', '1:ss_begin_gather_1'),
    *connect('ss_begin_gather_1_axis:0', '2:ss_begin_gather_1'),
    *connect('ss_begin_gather_2_idx:0', '1:ss_begin_gather_2'),
    *connect('ss_begin_gather_2_axis:0', '2:ss_begin_gather_2'),
    *connect('ss_begin_gather_3_idx:0', '1:ss_begin_gather_3'),
    *connect('ss_begin_gather_3_axis:0', '2:ss_begin_gather_3'),

    *connect('ss_end_gather_0_idx:0', '1:ss_end_gather_0'),
    *connect('ss_end_gather_0_axis:0', '2:ss_end_gather_0'),
    *connect('ss_end_gather_1_idx:0', '1:ss_end_gather_1'),
    *connect('ss_end_gather_1_axis:0', '2:ss_end_gather_1'),
    *connect('ss_end_gather_2_idx:0', '1:ss_end_gather_2'),
    *connect('ss_end_gather_2_axis:0', '2:ss_end_gather_2'),
    *connect('ss_end_gather_3_idx:0', '1:ss_end_gather_3'),
    *connect('ss_end_gather_3_axis:0', '2:ss_end_gather_3'),
]


class ConvertSliceTests(unittest.TestCase):

    def test_convert_slice_to_strided_slice_one_axis(self):
        graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_graph,
            update_attributes={
                'starts': {'value': int64_array([0]), 'shape': [1]},
                'ends': {'value': int64_array([1]), 'shape': [1]},
                'axes': {'value': int64_array([0]), 'shape': [1]},
                'axes_d': {'value': int64_array([0]), 'shape': [1]},
                'steps': {'value': int64_array([1]), 'shape': [1]},
                'steps_d': {'value': int64_array([1]), 'shape': [1]}
            },
            nodes_with_edges_only=True
        )

        ref_graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_ref_graph + [
                *connect('ss_begin_cast:0', '0:ss_begin_gather_0'),
                *connect('ss_begin_gather_0:0', '0:ss_begin_concat'),
                *connect('ss_begin_const_1:0', '1:ss_begin_concat'),
                *connect('ss_begin_const_2:0', '2:ss_begin_concat'),
                *connect('ss_begin_const_3:0', '3:ss_begin_concat'),

                *connect('ss_end_cast:0', '0:ss_end_gather_0'),
                *connect('ss_end_gather_0:0', '0:ss_end_concat'),
                *connect('ss_end_const_1:0', '1:ss_end_concat'),
                *connect('ss_end_const_2:0', '2:ss_end_concat'),
                *connect('ss_end_const_3:0', '3:ss_end_concat'),
            ],
            update_attributes={
                'starts': {'value': int64_array([0]), 'shape': [1]},
                'ends': {'value': int64_array([1]), 'shape': [1]},
                'ss_strides': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]},
                'ss': {'begin_mask': int64_array([1, 0, 0, 0]), 'end_mask': int64_array([1, 0, 0, 0])}
            }
        )
        ConvertSlice().find_and_replace_pattern(graph)
        (flag, resp) = compare_graphs(graph, ref_graph, 'result', check_op_attrs=True)
        self.assertTrue(flag, resp)

    def test_convert_slice_to_strided_slice_one_axis_steps_is_2(self):
        graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_graph,
            update_attributes={
                'starts': {'value': int64_array([0]), 'shape': [1]},
                'ends': {'value': int64_array([150]), 'shape': [1]},
                'axes': {'value': int64_array([2]), 'shape': [1]},
                'axes_d': {'value': int64_array([2]), 'shape': [1]},
                'steps': {'value': int64_array([2]), 'shape': [1]},
                'steps_d': {'value': int64_array([2]), 'shape': [1]}
            },
            nodes_with_edges_only=True
        )

        ref_graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_ref_graph + [
                *connect('ss_begin_cast:0', '0:ss_begin_gather_0'),
                *connect('ss_begin_gather_0:0', '2:ss_begin_concat'),
                *connect('ss_begin_const_0:0', '0:ss_begin_concat'),
                *connect('ss_begin_const_1:0', '1:ss_begin_concat'),
                *connect('ss_begin_const_3:0', '3:ss_begin_concat'),

                *connect('ss_end_cast:0', '0:ss_end_gather_0'),
                *connect('ss_end_gather_0:0', '2:ss_end_concat'),
                *connect('ss_end_const_0:0', '0:ss_end_concat'),
                *connect('ss_end_const_1:0', '1:ss_end_concat'),
                *connect('ss_end_const_3:0', '3:ss_end_concat'),
            ],
            update_attributes={
                'starts': {'value': int64_array([0]), 'shape': [1]},
                'ends': {'value': int64_array([150]), 'shape': [1]},
                'ss_strides': {'value': int64_array([1, 1, 2, 1]), 'shape': [4]},
                'ss': {'begin_mask': int64_array([0, 0, 1, 0]), 'end_mask': int64_array([0, 0, 1, 0])}
            }
        )
        ConvertSlice().find_and_replace_pattern(graph)
        (flag, resp) = compare_graphs(graph, ref_graph, 'result', check_op_attrs=True)
        self.assertTrue(flag, resp)

    def test_convert_slice_to_strided_slice_two_axes(self):
        graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_graph,
            update_attributes={
                'starts': {'value': int64_array([0, 0]), 'shape': [2]},
                'ends': {'value': int64_array([150, 150]), 'shape': [2]},
                'axes': {'value': int64_array([2, 3]), 'shape': [2]},
                'axes_d': {'value': int64_array([2, 3]), 'shape': [2]},
                'steps': {'value': int64_array([1, 1]), 'shape': [2]},
                'steps_d': {'value': int64_array([1, 1]), 'shape': [2]}
            },
            nodes_with_edges_only=True
        )

        ref_graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_ref_graph + [
                *connect('ss_begin_cast:0', '0:ss_begin_gather_0'),
                *connect('ss_begin_gather_0:0', '2:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_1'),
                *connect('ss_begin_gather_1:0', '3:ss_begin_concat'),
                *connect('ss_begin_const_0:0', '0:ss_begin_concat'),
                *connect('ss_begin_const_1:0', '1:ss_begin_concat'),

                *connect('ss_end_cast:0', '0:ss_end_gather_0'),
                *connect('ss_end_gather_0:0', '2:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_1'),
                *connect('ss_end_gather_1:0', '3:ss_end_concat'),
                *connect('ss_end_const_0:0', '0:ss_end_concat'),
                *connect('ss_end_const_1:0', '1:ss_end_concat'),
            ],
            update_attributes={
                'starts': {'value': int64_array([0, 0]), 'shape': [2]},
                'ends': {'value': int64_array([150, 150]), 'shape': [2]},
                'ss_strides': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]},
                'ss': {'begin_mask': int64_array([0, 0, 1, 1]), 'end_mask': int64_array([0, 0, 1, 1])}
            }
        )
        ConvertSlice().find_and_replace_pattern(graph)
        (flag, resp) = compare_graphs(graph, ref_graph, 'result', check_op_attrs=True)
        self.assertTrue(flag, resp)

    def test_convert_slice_to_strided_slice_three_axes(self):
        graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_graph,
            update_attributes={
                'starts': {'value': int64_array([0, 0, 0]), 'shape': [3]},
                'ends': {'value': int64_array([2, 150, 150]), 'shape': [3]},
                'axes': {'value': int64_array([1, 2, 3]), 'shape': [3]},
                'axes_d': {'value': int64_array([1, 2, 3]), 'shape': [3]},
                'steps': {'value': int64_array([1, 1, 1]), 'shape': [3]},
                'steps_d': {'value': int64_array([1, 1, 1]), 'shape': [3]}
            },
            nodes_with_edges_only=True
        )

        ref_graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_ref_graph + [
                *connect('ss_begin_cast:0', '0:ss_begin_gather_0'),
                *connect('ss_begin_gather_0:0', '1:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_1'),
                *connect('ss_begin_gather_1:0', '2:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_2'),
                *connect('ss_begin_gather_2:0', '3:ss_begin_concat'),
                *connect('ss_begin_const_0:0', '0:ss_begin_concat'),

                *connect('ss_end_cast:0', '0:ss_end_gather_0'),
                *connect('ss_end_gather_0:0', '1:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_1'),
                *connect('ss_end_gather_1:0', '2:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_2'),
                *connect('ss_end_gather_2:0', '3:ss_end_concat'),
                *connect('ss_end_const_0:0', '0:ss_end_concat'),
            ],
            update_attributes={
                'starts': {'value': int64_array([0, 0, 0]), 'shape': [3]},
                'ends': {'value': int64_array([2, 150, 150]), 'shape': [3]},
                'ss_strides': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]},
                'ss': {'begin_mask': int64_array([0, 1, 1, 1]), 'end_mask': int64_array([0, 1, 1, 1])}
            }
        )
        ConvertSlice().find_and_replace_pattern(graph)
        (flag, resp) = compare_graphs(graph, ref_graph, 'result', check_op_attrs=True)
        self.assertTrue(flag, resp)

    def test_convert_slice_to_strided_slice_not_sorted_axes(self):
        graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_graph,
            update_attributes={
                'starts': {'value': int64_array([0, 1, 1, 0]), 'shape': [4]},
                'ends': {'value': int64_array([1, 150, 150, 2]), 'shape': [4]},
                'axes': {'value': int64_array([0, 2, 3, 1]), 'shape': [4]},
                'axes_d': {'value': int64_array([0, 2, 3, 1]), 'shape': [4]},
                'steps': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]},
                'steps_d': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]}
            },
            nodes_with_edges_only=True
        )

        ref_graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_ref_graph + [
                *connect('ss_begin_cast:0', '0:ss_begin_gather_0'),
                *connect('ss_begin_gather_0:0', '0:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_1'),
                *connect('ss_begin_gather_1:0', '2:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_2'),
                *connect('ss_begin_gather_2:0', '3:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_3'),
                *connect('ss_begin_gather_3:0', '1:ss_begin_concat'),

                *connect('ss_end_cast:0', '0:ss_end_gather_0'),
                *connect('ss_end_gather_0:0', '0:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_1'),
                *connect('ss_end_gather_1:0', '2:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_2'),
                *connect('ss_end_gather_2:0', '3:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_3'),
                *connect('ss_end_gather_3:0', '1:ss_end_concat'),
            ],
            update_attributes={
                'starts': {'value': int64_array([0, 1, 1, 0]), 'shape': [4]},
                'ends': {'value': int64_array([1, 150, 150, 2]), 'shape': [4]},
                'ss_strides': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]},
                'ss': {'begin_mask': int64_array([1, 1, 1, 1]), 'end_mask': int64_array([1, 1, 1, 1])}
            }
        )
        ConvertSlice().find_and_replace_pattern(graph)
        (flag, resp) = compare_graphs(graph, ref_graph, 'result', check_op_attrs=True)
        self.assertTrue(flag, resp)

    def test_convert_slice_to_strided_slice_without_axes_and_steps(self):
        graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=[
                *connect('input:0', '0:slice'),
                *connect('starts:0', '1:slice'),
                *connect('ends:0', '2:slice'),
                *connect('slice:0', '0:result')
            ],
            update_attributes={
                'starts': {'value': int64_array([0, 0, 0, 0]), 'shape': [4]},
                'ends': {'value': int64_array([1, 2, 150, 150]), 'shape': [4]},
            },
            nodes_with_edges_only=True
        )

        ref_graph = build_graph(
            nodes_attrs=nodes_attributes,
            edges=pattern_ref_graph + [
                *connect('ss_begin_cast:0', '0:ss_begin_gather_0'),
                *connect('ss_begin_gather_0:0', '0:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_1'),
                *connect('ss_begin_gather_1:0', '1:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_2'),
                *connect('ss_begin_gather_2:0', '2:ss_begin_concat'),
                *connect_data('ss_begin_cast:0', '0:ss_begin_gather_3'),
                *connect('ss_begin_gather_3:0', '3:ss_begin_concat'),

                *connect('ss_end_cast:0', '0:ss_end_gather_0'),
                *connect('ss_end_gather_0:0', '0:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_1'),
                *connect('ss_end_gather_1:0', '1:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_2'),
                *connect('ss_end_gather_2:0', '2:ss_end_concat'),
                *connect_data('ss_end_cast:0', '0:ss_end_gather_3'),
                *connect('ss_end_gather_3:0', '3:ss_end_concat'),
            ],
            update_attributes={
                'starts': {'value': int64_array([0, 0, 0, 0]), 'shape': [4]},
                'ends': {'value': int64_array([1, 2, 150, 150]), 'shape': [4]},
                'ss_strides': {'value': int64_array([1, 1, 1, 1]), 'shape': [4]},
                'ss': {'begin_mask': int64_array([1, 1, 1, 1]), 'end_mask': int64_array([1, 1, 1, 1])}
            }
        )
        ConvertSlice().find_and_replace_pattern(graph)
        (flag, resp) = compare_graphs(graph, ref_graph, 'result', check_op_attrs=True)
        self.assertTrue(flag, resp)