import onnx

from onnx.helper import make_node, make_tensor, TensorProto

from furiosa_sdk_quantizer.interfaces.transformer import Transformer
from furiosa_sdk_quantizer.frontend.onnx.transformer import utils
from furiosa_sdk_quantizer.frontend.onnx.utils.check_model import check_model


class Convert2dSumToAdd(Transformer):
    def transform(self, model: onnx.ModelProto) -> onnx.ModelProto:
        optimized_nodes = []
        for node in model.graph.node:
            if node.op_type != "Sum":
                optimized_nodes.append(node)
                continue

            if len(node.input) != 2:
                optimized_nodes.append(node)
                continue

            new_node = make_node(
                "Add", inputs=[node.input[0], node.input[1]], outputs=[node.output[0]]
            )

            optimized_nodes.append(new_node)

        # remove duplicate node(s) in optimized nodes
        seen = []
        for op_node in optimized_nodes:
            if op_node in seen:
                continue
            seen.append(op_node)
        optimized_nodes = seen

        model = utils.rebuild_model(model, optimized_nodes)
        check_model(model)

        return model
