import pytest

import numpy as np
import onnx
import onnx.helper as oh
import qonnx.core.data_layout as DataLayout
from onnx import TensorProto
from onnx.helper import make_opsetid
from qonnx.core.datatype import DataType
from qonnx.core.modelwrapper import ModelWrapper
from qonnx.transformation.infer_data_layouts import InferDataLayouts
from qonnx.transformation.infer_shapes import InferShapes
from qonnx.util.basic import gen_finn_dt_tensor, qonnx_make_model

import finn.core.onnx_exec as oxe
from finn.transformation.streamline.reorder import MakeScaleResizeNHWC


def create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt, onnx_opset):
    ofm_dim_h = ifm_dim[0] * scales[2]
    ofm_dim_w = ifm_dim[1] * scales[3]
    inp = oh.make_tensor_value_info("inp", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]])

    param = oh.make_tensor_value_info("scales", TensorProto.FLOAT, [4])

    # Not actually used, only needed for compliance with the Resize node interface
    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])

    outp_up = oh.make_tensor_value_info(
        "outp_up", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w]
    )
    outp = oh.make_tensor_value_info("outp", TensorProto.FLOAT, [1, ofm_dim_h, ofm_dim_w, ifm_ch])

    resize_node = oh.make_node(
        "Resize",
        inputs=["inp", "roi", "scales"],
        outputs=["outp_up"],
        name="Resize1",
        mode=mode,
    )

    transpose_node = onnx.helper.make_node(
        "Transpose",
        inputs=["outp_up"],
        outputs=["outp"],
        name="Transpose1",
        perm=[0, 2, 3, 1],
    )

    graph = oh.make_graph(
        nodes=[resize_node, transpose_node],
        name="resize_graph",
        inputs=[inp],
        outputs=[outp],
        value_info=[outp_up, param, roi],
    )

    model = qonnx_make_model(
        graph, producer_name="resize_model1", opset_imports=[make_opsetid("", onnx_opset)]
    )
    model = ModelWrapper(model)
    model.set_tensor_datatype("inp", idt)
    model.set_tensor_datatype("outp", idt)

    model.set_tensor_layout("inp", DataLayout.NCHW)
    model = model.transform(InferShapes())
    model = model.transform(InferDataLayouts())

    return model


def create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt, onnx_opset):
    ofm_dim_h = ifm_dim[0] * scales[2]
    ofm_dim_w = ifm_dim[1] * scales[3]
    inp = oh.make_tensor_value_info("inp", TensorProto.FLOAT, [1, ifm_dim[0], ifm_dim[1], ifm_ch])

    param = oh.make_tensor_value_info("scales", TensorProto.FLOAT, [4])

    # Not actually used, only needed for compliance with the Resize node interface
    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])

    outp = oh.make_tensor_value_info("outp", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w])
    outp_tr = oh.make_tensor_value_info(
        "outp_tr", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]]
    )

    transpose_node = onnx.helper.make_node(
        "Transpose",
        inputs=["inp"],
        outputs=["outp_tr"],
        name="Transpose1",
        perm=[0, 3, 1, 2],
    )

    resize_node = oh.make_node(
        "Resize",
        inputs=["outp_tr", "roi", "scales"],
        outputs=["outp"],
        name="Resize1",
        mode=mode,
    )

    graph = oh.make_graph(
        nodes=[transpose_node, resize_node],
        name="resize_graph",
        inputs=[inp],
        outputs=[outp],
        value_info=[outp_tr, param, roi],
    )

    model = qonnx_make_model(
        graph, producer_name="resize_model2", opset_imports=[make_opsetid("", onnx_opset)]
    )
    model = ModelWrapper(model)
    model.set_tensor_datatype("inp", idt)
    model.set_tensor_datatype("outp", idt)
    model.set_tensor_layout("inp", DataLayout.NHWC)

    model = model.transform(InferShapes())
    model = model.transform(InferDataLayouts())

    return model


def create_transpose_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt, onnx_opset):
    ofm_dim_h = ifm_dim[0] * scales[2]
    ofm_dim_w = ifm_dim[1] * scales[3]
    inp = oh.make_tensor_value_info("inp", TensorProto.FLOAT, [1, ifm_dim[0], ifm_dim[1], ifm_ch])

    param = oh.make_tensor_value_info("scales", TensorProto.FLOAT, scales)

    # Not actually used, only needed for compliance with the Resize node interface
    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])

    outp_tr = oh.make_tensor_value_info(
        "outp_tr", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]]
    )

    outp_up = oh.make_tensor_value_info(
        "outp_up", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w]
    )
    outp = oh.make_tensor_value_info("outp", TensorProto.FLOAT, [1, ofm_dim_h, ofm_dim_w, ifm_ch])

    transpose_node1 = onnx.helper.make_node(
        "Transpose",
        inputs=["inp"],
        outputs=["outp_tr"],
        name="Transpose1",
        perm=[0, 3, 1, 2],
    )

    resize_node = oh.make_node(
        "Resize",
        inputs=["outp_tr", "roi", "scales"],
        outputs=["outp_up"],
        name="Resize1",
        mode=mode,
    )

    transpose_node2 = onnx.helper.make_node(
        "Transpose",
        inputs=["outp_up"],
        outputs=["outp"],
        name="Transpose2",
        perm=[0, 2, 3, 1],
    )

    graph = oh.make_graph(
        nodes=[transpose_node1, resize_node, transpose_node2],
        name="resize_graph",
        inputs=[inp],
        outputs=[outp],
        value_info=[outp_up, outp_tr, param, roi],
    )

    model = qonnx_make_model(
        graph, producer_name="resize_model3", opset_imports=[make_opsetid("", onnx_opset)]
    )
    model = ModelWrapper(model)
    model.set_tensor_datatype("inp", idt)
    model.set_tensor_datatype("outp", idt)
    model.set_tensor_layout("inp", DataLayout.NHWC)

    model = model.transform(InferShapes())
    model = model.transform(InferDataLayouts())

    return model


def create_resize_transpose_sizes(ifm_dim, ifm_ch, sizes, mode, idt, onnx_opset):
    ofm_dim_h = sizes[2]
    ofm_dim_w = sizes[3]
    inp = oh.make_tensor_value_info("inp", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]])

    # Empty scales
    scales = oh.make_tensor_value_info("scales", TensorProto.FLOAT, [])

    # Not actually used, only needed for compliance with the Resize node interface
    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])

    param = oh.make_tensor_value_info("sizes", TensorProto.INT64, [4])

    outp_up = oh.make_tensor_value_info(
        "outp_up", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w]
    )
    outp = oh.make_tensor_value_info("outp", TensorProto.FLOAT, [1, ofm_dim_h, ofm_dim_w, ifm_ch])

    if onnx_opset == 11:
        inputs = ["inp", "roi", "scales", "sizes"]
    elif onnx_opset == 13:
        inputs = ["inp", "roi", "", "sizes"]
    else:
        raise Exception("Unsupported opset version")

    resize_node = oh.make_node(
        "Resize",
        inputs=inputs,
        outputs=["outp_up"],
        name="Resize1",
        mode=mode,
    )

    transpose_node = onnx.helper.make_node(
        "Transpose",
        inputs=["outp_up"],
        outputs=["outp"],
        name="Transpose1",
        perm=[0, 2, 3, 1],
    )

    if onnx_opset == 11:
        value_info = [outp_up, roi, scales, param]
    elif onnx_opset == 13:
        value_info = [outp_up, roi, param]
    else:
        raise Exception("Unsupported opset version")

    graph = oh.make_graph(
        nodes=[resize_node, transpose_node],
        name="resize_graph",
        inputs=[inp],
        outputs=[outp],
        value_info=value_info,
    )

    model = qonnx_make_model(
        graph, producer_name="resize_model4", opset_imports=[make_opsetid("", onnx_opset)]
    )
    model = ModelWrapper(model)
    model.set_tensor_datatype("inp", idt)
    model.set_tensor_datatype("outp", idt)

    model.set_tensor_layout("inp", DataLayout.NCHW)
    model = model.transform(InferShapes())
    model = model.transform(InferDataLayouts())

    return model


def check_transform(model):
    graph = model.graph
    node_ind = 0
    for n in graph.node:
        node_ind += 1
        if n.op_type == "Upsample" or n.op_type == "Resize":
            if model.get_tensor_layout(n.output[0]) == DataLayout.NHWC:
                return True
    return False


@pytest.mark.streamline
# input dimension
@pytest.mark.parametrize("ifm_dim", [[2**i, 2**i] for i in range(3, 6)])
# input channels
@pytest.mark.parametrize("ifm_ch", [3])
# scales
@pytest.mark.parametrize("scales", [[1, 1, i, j] for i in range(2, 5) for j in range(2, 5)])
# sizes
@pytest.mark.parametrize(
    "sizes", [[1, 3, 2**i, 2**j] for i in range(6, 7) for j in range(6, 7)]
)
# mode
@pytest.mark.parametrize("mode", ["nearest"])
# input datatype
@pytest.mark.parametrize("idt", [DataType["INT4"]])
# Onnx opset version
@pytest.mark.parametrize("onnx_opset", [11, 13])
def test_scale_resize_nhwc(ifm_dim, ifm_ch, sizes, scales, mode, idt, onnx_opset):
    # create models
    resize_model1 = create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt, onnx_opset)
    resize_model2 = create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt, onnx_opset)
    resize_model3 = create_transpose_resize_transpose(
        ifm_dim, ifm_ch, scales, mode, idt, onnx_opset
    )
    resize_model4 = create_resize_transpose_sizes(ifm_dim, ifm_ch, sizes, mode, idt, onnx_opset)

    # set initializers
    resize_model1.set_initializer("scales", np.array(scales, dtype=np.float32))
    resize_model2.set_initializer("scales", np.array(scales, dtype=np.float32))
    resize_model3.set_initializer("scales", np.array(scales, dtype=np.float32))
    resize_model4.set_initializer("sizes", np.array(sizes, dtype=np.int64))
    resize_model4.set_initializer("scales", np.array([], dtype=np.float32))  # Opset 11

    # generate input tensor for testing
    input_tensor_nchw = gen_finn_dt_tensor(idt, [1, ifm_ch, ifm_dim[0], ifm_dim[1]])
    input_tensor_nhwc = gen_finn_dt_tensor(idt, [1, ifm_dim[0], ifm_dim[1], ifm_ch])
    input_dict_nchw = {"inp": input_tensor_nchw}
    input_dict_nhwc = {"inp": input_tensor_nhwc}

    # execute first model
    output_dict1 = oxe.execute_onnx(resize_model1, input_dict_nchw)
    expected1 = output_dict1["outp"]

    # transform Resize into ResizeNHWC
    resize_model1 = resize_model1.transform(MakeScaleResizeNHWC())
    resize_model1 = resize_model1.transform(InferDataLayouts())

    # execute transformed model
    output_node_name1 = resize_model1.graph.output[0].name
    output_dict1 = oxe.execute_onnx(resize_model1, input_dict_nchw, return_full_exec_context=False)
    output1 = output_dict1[output_node_name1]

    # compare outputs
    assert (expected1 == output1).all()
    assert check_transform(resize_model1)

    # execute second model
    output_dict2 = oxe.execute_onnx(resize_model2, input_dict_nhwc)
    expected2 = output_dict2["outp"]

    # transform Resize into ResizeNHWC
    resize_model2 = resize_model2.transform(MakeScaleResizeNHWC())
    resize_model2 = resize_model2.transform(InferDataLayouts())

    # execute transformed model
    output_node_name2 = resize_model2.graph.output[0].name
    output_dict2 = oxe.execute_onnx(resize_model2, input_dict_nhwc, return_full_exec_context=False)
    output2 = output_dict2[output_node_name2]

    # compare outputs
    assert (expected2 == output2).all()
    assert check_transform(resize_model2)

    # execute third model
    output_dict3 = oxe.execute_onnx(resize_model3, input_dict_nhwc)
    expected3 = output_dict3["outp"]

    # transform Resize into ResizeNHWC
    resize_model3 = resize_model3.transform(MakeScaleResizeNHWC())
    resize_model3 = resize_model3.transform(InferDataLayouts())

    # execute transformed model
    output_node_name3 = resize_model3.graph.output[0].name
    output_dict3 = oxe.execute_onnx(resize_model3, input_dict_nhwc, return_full_exec_context=False)
    output3 = output_dict3[output_node_name3]

    # compare outputs
    assert (expected3 == output3).all()
    assert check_transform(resize_model3)

    # execute fourth model
    output_dict4 = oxe.execute_onnx(resize_model4, input_dict_nchw)
    expected4 = output_dict4["outp"]

    # transform Resize into ResizeNHWC
    resize_model4 = resize_model4.transform(MakeScaleResizeNHWC())
    resize_model4 = resize_model4.transform(InferDataLayouts())

    # execute transformed model
    output_node_name4 = resize_model4.graph.output[0].name
    output_dict4 = oxe.execute_onnx(resize_model4, input_dict_nchw, return_full_exec_context=False)
    output4 = output_dict4[output_node_name4]

    # compare outputs
    assert (expected4 == output4).all()
    assert check_transform(resize_model4)
