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

[((64, 1, 64, 62), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(64, 25, 64, 31): '"""perm"""', (64, 33, 64, 61): '[[0, 3, 1, 2], [0, 2, 3, 1]]'}, {}), "('perm', [[0, 3, 1, 2], [0, 2, 3, 1]])", False, 'import pytest\n'), ((20, 22, 22, 5), 'onnx.helper.make_node', 'oh.make_node', (), '', True, 'from onnx import helper as oh\n'), ((24, 22, 26, 5), 'onnx.helper.make_node', 'oh.make_node', (), '', True, 'from onnx import helper as oh\n'), ((28, 17, 30, 5), 'onnx.helper.make_node', 'oh.make_node', (), '', True, 'from onnx import helper as oh\n'), ((32, 20, 34, 5), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', ({(33, 8, 33, 23): '"""in_transpose1"""', (33, 25, 33, 42): 'TensorProto.FLOAT', (33, 44, 33, 52): 'in_shape'}, {}), "('in_transpose1', TensorProto.FLOAT, in_shape)", True, 'from onnx import helper as oh\n'), ((35, 20, 37, 5), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', ({(36, 8, 36, 23): '"""in_transpose2"""', (36, 25, 36, 42): 'TensorProto.FLOAT', (36, 44, 36, 52): 'in_shape'}, {}), "('in_transpose2', TensorProto.FLOAT, in_shape)", True, 'from onnx import helper as oh\n'), ((38, 21, 40, 5), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', ({(39, 8, 39, 24): '"""out_transpose1"""', (39, 26, 39, 43): 'TensorProto.FLOAT', (39, 45, 39, 54): 'out_shape'}, {}), "('out_transpose1', TensorProto.FLOAT, out_shape)", True, 'from onnx import helper as oh\n'), ((41, 21, 43, 5), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', ({(42, 8, 42, 24): '"""out_transpose2"""', (42, 26, 42, 43): 'TensorProto.FLOAT', (42, 45, 42, 54): 'out_shape'}, {}), "('out_transpose2', TensorProto.FLOAT, out_shape)", True, 'from onnx import helper as oh\n'), ((44, 16, 44, 84), 'onnx.helper.make_tensor_value_info', 'oh.make_tensor_value_info', ({(44, 42, 44, 53): '"""out_join1"""', (44, 55, 44, 72): 'TensorProto.FLOAT', (44, 74, 44, 83): 'out_shape'}, {}), "('out_join1', TensorProto.FLOAT, out_shape)", True, 'from onnx import helper as oh\n'), ((46, 12, 55, 5), 'onnx.helper.make_graph', 'oh.make_graph', (), '', True, 'from onnx import helper as oh\n'), ((57, 17, 57, 65), 'onnx.helper.make_model', 'oh.make_model', (), '', True, 'from onnx import helper as oh\n'), ((58, 12, 58, 36), 'finn.core.modelwrapper.ModelWrapper', 'ModelWrapper', ({(58, 25, 58, 35): 'onnx_model'}, {}), '(onnx_model)', False, 'from finn.core.modelwrapper import ModelWrapper\n'), ((75, 16, 75, 60), 'finn.util.basic.gen_finn_dt_tensor', 'gen_finn_dt_tensor', ({(75, 35, 75, 46): 'input_dtype', (75, 48, 75, 59): 'input_shape'}, {}), '(input_dtype, input_shape)', False, 'from finn.util.basic import gen_finn_dt_tensor\n'), ((82, 11, 82, 70), 'finn.core.onnx_exec.compare_execution', 'oxe.compare_execution', ({(82, 33, 82, 38): 'model', (82, 40, 82, 57): 'model_transformed', (82, 59, 82, 69): 'input_dict'}, {}), '(model, model_transformed, input_dict)', True, 'import finn.core.onnx_exec as oxe\n'), ((80, 40, 80, 66), 'finn.transformation.streamline.reorder.MoveTransposePastJoinAdd', 'MoveTransposePastJoinAdd', ({}, {}), '()', False, 'from finn.transformation.streamline.reorder import MoveTransposePastJoinAdd\n')]