import pytest

from onnx import TensorProto
from onnx import helper as oh

import finn.core.onnx_exec as oxe
from finn.core.modelwrapper import ModelWrapper
from finn.transformation.streamline.reorder import MoveTransposePastJoinAdd
from finn.util.basic import gen_finn_dt_tensor


def create_model(perm):
    if perm == [0, 3, 1, 2]:
        in_shape = [1, 128, 1, 256]
        out_shape = [1, 256, 128, 1]
    if perm == [0, 2, 3, 1]:
        in_shape = [1, 256, 128, 1]
        out_shape = [1, 128, 1, 256]

    Transpose1_node = oh.make_node(
        "Transpose", inputs=["in_transpose1"], outputs=["out_transpose1"], perm=perm
    )

    Transpose2_node = oh.make_node(
        "Transpose", inputs=["in_transpose2"], outputs=["out_transpose2"], perm=perm
    )

    Join1_node = oh.make_node(
        "Add", inputs=["out_transpose1", "out_transpose2"], outputs=["out_join1"]
    )

    in_transpose1 = oh.make_tensor_value_info(
        "in_transpose1", TensorProto.FLOAT, in_shape
    )
    in_transpose2 = oh.make_tensor_value_info(
        "in_transpose2", TensorProto.FLOAT, in_shape
    )
    out_transpose1 = oh.make_tensor_value_info(
        "out_transpose1", TensorProto.FLOAT, out_shape
    )
    out_transpose2 = oh.make_tensor_value_info(
        "out_transpose2", TensorProto.FLOAT, out_shape
    )
    out_join1 = oh.make_tensor_value_info("out_join1", TensorProto.FLOAT, out_shape)

    graph = oh.make_graph(
        nodes=[Transpose1_node, Transpose2_node, Join1_node],
        name="test_graph",
        inputs=[in_transpose1, in_transpose2],
        outputs=[out_join1],
        value_info=[
            out_transpose1,
            out_transpose2,
        ],
    )

    onnx_model = oh.make_model(graph, producer_name="test_model")
    model = ModelWrapper(onnx_model)

    return model


# Permutation of transpose node
@pytest.mark.parametrize("perm", [[0, 3, 1, 2], [0, 2, 3, 1]])
def test_move_identical_op_past_join_op(perm):
    model = create_model(perm)

    # Create input data
    input0_tensor_name = model.graph.input[0].name
    input1_tensor_name = model.graph.input[1].name

    # Note: it is assumed that both tensors have the same shape and data type
    input_shape = model.get_tensor_shape(input0_tensor_name)
    input_dtype = model.get_tensor_datatype(input0_tensor_name)
    input_val = gen_finn_dt_tensor(input_dtype, input_shape)
    input_dict = {}
    input_dict[input0_tensor_name] = input_val
    input_dict[input1_tensor_name] = input_val

    model_transformed = model.transform(MoveTransposePastJoinAdd())

    assert oxe.compare_execution(model, model_transformed, input_dict)

    # Check if order changed
    node0_input0_model = model.find_consumers(model.graph.input[0].name)[0].op_type
    node1_input1_model = model.find_consumers(model.graph.input[1].name)[0].op_type
    node0_input0_model_transformed = model_transformed.find_consumers(
        model_transformed.graph.input[0].name
    )[0].op_type
    node1_input1_model_transformed = model_transformed.find_consumers(
        model_transformed.graph.input[1].name
    )[0].op_type
    assert node0_input0_model != node0_input0_model_transformed
    assert node1_input1_model != node1_input1_model_transformed

[((1846, 1907), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""perm"""', '[[0, 3, 1, 2], [0, 2, 3, 1]]'], {}), "('perm', [[0, 3, 1, 2], [0, 2, 3, 1]])\n", (1869, 1907), False, 'import pytest\n'), ((533, 628), 'onnx.helper.make_node', 'oh.make_node', (['"""Transpose"""'], {'inputs': "['in_transpose1']", 'outputs': "['out_transpose1']", 'perm': 'perm'}), "('Transpose', inputs=['in_transpose1'], outputs=[\n    'out_transpose1'], perm=perm)\n", (545, 628), True, 'from onnx import helper as oh\n'), ((661, 756), 'onnx.helper.make_node', 'oh.make_node', (['"""Transpose"""'], {'inputs': "['in_transpose2']", 'outputs': "['out_transpose2']", 'perm': 'perm'}), "('Transpose', inputs=['in_transpose2'], outputs=[\n    'out_transpose2'], perm=perm)\n", (673, 756), True, 'from onnx import helper as oh\n'), ((784, 876), 'onnx.helper.make_node', 'oh.make_node', (['"""Add"""'], {'inputs': "['out_transpose1', 'out_transpose2']", 'outputs': "['out_join1']"}), "('Add', inputs=['out_transpose1', 'out_transpose2'], outputs=[\n    'out_join1'])\n", (796, 876), True, 'from onnx import helper as oh\n'), ((907, 978), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', (['"""in_transpose1"""', 'TensorProto.FLOAT', 'in_shape'], {}), "('in_transpose1', TensorProto.FLOAT, in_shape)\n", (932, 978), True, 'from onnx import helper as oh\n'), ((1013, 1084), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', (['"""in_transpose2"""', 'TensorProto.FLOAT', 'in_shape'], {}), "('in_transpose2', TensorProto.FLOAT, in_shape)\n", (1038, 1084), True, 'from onnx import helper as oh\n'), ((1120, 1193), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', (['"""out_transpose1"""', 'TensorProto.FLOAT', 'out_shape'], {}), "('out_transpose1', TensorProto.FLOAT, out_shape)\n", (1145, 1193), True, 'from onnx import helper as oh\n'), ((1229, 1302), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', (['"""out_transpose2"""', 'TensorProto.FLOAT', 'out_shape'], {}), "('out_transpose2', TensorProto.FLOAT, out_shape)\n", (1254, 1302), True, 'from onnx import helper as oh\n'), ((1333, 1401), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', (['"""out_join1"""', 'TensorProto.FLOAT', 'out_shape'], {}), "('out_join1', TensorProto.FLOAT, out_shape)\n", (1358, 1401), True, 'from onnx import helper as oh\n'), ((1415, 1616), 'onnx.helper.make_graph', 'oh.make_graph', ([], {'nodes': '[Transpose1_node, Transpose2_node, Join1_node]', 'name': '"""test_graph"""', 'inputs': '[in_transpose1, in_transpose2]', 'outputs': '[out_join1]', 'value_info': '[out_transpose1, out_transpose2]'}), "(nodes=[Transpose1_node, Transpose2_node, Join1_node], name=\n    'test_graph', inputs=[in_transpose1, in_transpose2], outputs=[out_join1\n    ], value_info=[out_transpose1, out_transpose2])\n", (1428, 1616), True, 'from onnx import helper as oh\n'), ((1707, 1755), 'onnx.helper.make_model', 'oh.make_model', (['graph'], {'producer_name': '"""test_model"""'}), "(graph, producer_name='test_model')\n", (1720, 1755), True, 'from onnx import helper as oh\n'), ((1768, 1792), 'finn.core.modelwrapper.ModelWrapper', 'ModelWrapper', (['onnx_model'], {}), '(onnx_model)\n', (1780, 1792), False, 'from finn.core.modelwrapper import ModelWrapper\n'), ((2333, 2377), 'finn.util.basic.gen_finn_dt_tensor', 'gen_finn_dt_tensor', (['input_dtype', 'input_shape'], {}), '(input_dtype, input_shape)\n', (2351, 2377), False, 'from finn.util.basic import gen_finn_dt_tensor\n'), ((2573, 2632), 'finn.core.onnx_exec.compare_execution', 'oxe.compare_execution', (['model', 'model_transformed', 'input_dict'], {}), '(model, model_transformed, input_dict)\n', (2594, 2632), True, 'import finn.core.onnx_exec as oxe\n'), ((2533, 2559), 'finn.transformation.streamline.reorder.MoveTransposePastJoinAdd', 'MoveTransposePastJoinAdd', ([], {}), '()\n', (2557, 2559), False, 'from finn.transformation.streamline.reorder import MoveTransposePastJoinAdd\n')]