# SPDX-License-Identifier: Apache-2.0

# Copyright (c) ONNX Project Contributors

import unittest

import onnx
from onnx import checker, utils


class TestFunction(unittest.TestCase):
    def _verify_function_set(self, extracted_model, function_set, func_domain):  # type: ignore
        checker.check_model(extracted_model)
        self.assertEqual(len(extracted_model.functions), len(function_set))
        for function in function_set:
            self.assertIsNotNone(
                next(
                    (
                        f
                        for f in extracted_model.functions
                        if f.name == function and f.domain == func_domain
                    ),
                    None,
                )
            )

    def test_extract_model_with_local_function(self) -> None:
        r"""#   1. build a model with graph below. extract models with output combinations
        #   2. validate extracted models' local functions
        #
        # model graph:
        #      i0                    i1                 i2
        #      |   __________________|__________________/_________
        #      |  |                  |             |   /          |
        #      |  |                  |             |  /           |
        #   func_add        func_identity          add         identity
        #    |  ___\___________\____________________|_________    |
        #    | |    \           \                   |  _______|___|
        #    | |     \           \                  | |       |   |
        #    add     function_nested_identity_add   add     function_nested_identity_add
        #     |                 |                    |              |
        #     |                 |                    |              |
        #   o_func_add      o_all_func0           o_no_func     o_all_func1
        #
        # where function_nested_identity_add is a function that is defined with functions:
        #       a               b
        #       |               |
        #   func_identity   func_identity
        #             \       /
        #             func_add
        #                |
        #                c
        #
        """
        # function common
        func_domain = "local"
        func_opset_imports = [onnx.helper.make_opsetid("", 14)]
        func_nested_opset_imports = [
            onnx.helper.make_opsetid("", 14),
            onnx.helper.make_opsetid(func_domain, 1),
        ]

        # add function
        func_add_name = "func_add"
        func_add_inputs = ["a", "b"]
        func_add_outputs = ["c"]
        func_add_nodes = [onnx.helper.make_node("Add", ["a", "b"], ["c"])]
        func_add = onnx.helper.make_function(
            func_domain,
            func_add_name,
            func_add_inputs,
            func_add_outputs,
            func_add_nodes,
            func_opset_imports,
        )

        # identity function
        func_identity_name = "func_identity"
        func_identity_inputs = ["a"]
        func_identity_outputs = ["b"]
        func_identity_nodes = [onnx.helper.make_node("Identity", ["a"], ["b"])]
        func_identity = onnx.helper.make_function(
            func_domain,
            func_identity_name,
            func_identity_inputs,
            func_identity_outputs,
            func_identity_nodes,
            func_opset_imports,
        )

        # nested identity/add function
        func_nested_identity_add_name = "func_nested_identity_add"
        func_nested_identity_add_inputs = ["a", "b"]
        func_nested_identity_add_outputs = ["c"]
        func_nested_identity_add_nodes = [
            onnx.helper.make_node("func_identity", ["a"], ["a1"], domain=func_domain),
            onnx.helper.make_node("func_identity", ["b"], ["b1"], domain=func_domain),
            onnx.helper.make_node("func_add", ["a1", "b1"], ["c"], domain=func_domain),
        ]
        func_nested_identity_add = onnx.helper.make_function(
            func_domain,
            func_nested_identity_add_name,
            func_nested_identity_add_inputs,
            func_nested_identity_add_outputs,
            func_nested_identity_add_nodes,
            func_nested_opset_imports,
        )

        # create graph nodes
        node_func_add = onnx.helper.make_node(
            func_add_name, ["i0", "i1"], ["t0"], domain=func_domain
        )
        node_add0 = onnx.helper.make_node("Add", ["i1", "i2"], ["t2"])
        node_add1 = onnx.helper.make_node("Add", ["t0", "t2"], ["o_func_add"])
        node_func_identity = onnx.helper.make_node(
            func_identity_name, ["i1"], ["t1"], domain=func_domain
        )
        node_identity = onnx.helper.make_node("Identity", ["i1"], ["t3"])
        node_add2 = onnx.helper.make_node("Add", ["t3", "t2"], ["o_no_func"])
        node_func_nested0 = onnx.helper.make_node(
            func_nested_identity_add_name,
            ["t0", "t1"],
            ["o_all_func0"],
            domain=func_domain,
        )
        node_func_nested1 = onnx.helper.make_node(
            func_nested_identity_add_name,
            ["t3", "t2"],
            ["o_all_func1"],
            domain=func_domain,
        )

        graph_name = "graph_with_imbedded_functions"
        ir_version = 8
        opset_imports = [
            onnx.helper.make_opsetid("", 14),
            onnx.helper.make_opsetid("local", 1),
        ]
        tensor_type_proto = onnx.helper.make_tensor_type_proto(elem_type=2, shape=[5])

        graph = onnx.helper.make_graph(
            [
                node_func_add,
                node_add0,
                node_add1,
                node_func_identity,
                node_identity,
                node_func_nested0,
                node_func_nested1,
                node_add2,
            ],
            graph_name,
            [
                onnx.helper.make_value_info(name="i0", type_proto=tensor_type_proto),
                onnx.helper.make_value_info(name="i1", type_proto=tensor_type_proto),
                onnx.helper.make_value_info(name="i2", type_proto=tensor_type_proto),
            ],
            [
                onnx.helper.make_value_info(
                    name="o_no_func", type_proto=tensor_type_proto
                ),
                onnx.helper.make_value_info(
                    name="o_func_add", type_proto=tensor_type_proto
                ),
                onnx.helper.make_value_info(
                    name="o_all_func0", type_proto=tensor_type_proto
                ),
                onnx.helper.make_value_info(
                    name="o_all_func1", type_proto=tensor_type_proto
                ),
            ],
        )

        meta = {
            "ir_version": ir_version,
            "opset_imports": opset_imports,
            "producer_name": "test_extract_model_with_local_function",
            "functions": [func_identity, func_add, func_nested_identity_add],
        }
        model = onnx.helper.make_model(graph, **meta)

        checker.check_model(model)
        extracted_with_no_funcion = utils.Extractor(model).extract_model(
            ["i0", "i1", "i2"], ["o_no_func"]
        )
        self._verify_function_set(extracted_with_no_funcion, {}, func_domain)

        extracted_with_add_funcion = utils.Extractor(model).extract_model(
            ["i0", "i1", "i2"], ["o_func_add"]
        )
        self._verify_function_set(
            extracted_with_add_funcion, {func_add_name}, func_domain
        )

        extracted_with_o_all_funcion0 = utils.Extractor(model).extract_model(
            ["i0", "i1", "i2"], ["o_all_func0"]
        )
        self._verify_function_set(
            extracted_with_o_all_funcion0,
            {func_add_name, func_identity_name, func_nested_identity_add_name},
            func_domain,
        )

        extracted_with_o_all_funcion1 = utils.Extractor(model).extract_model(
            ["i0", "i1", "i2"], ["o_all_func1"]
        )
        self._verify_function_set(
            extracted_with_o_all_funcion1,
            {func_add_name, func_identity_name, func_nested_identity_add_name},
            func_domain,
        )

        extracted_with_o_all_funcion2 = utils.Extractor(model).extract_model(
            ["i0", "i1", "i2"],
            ["o_no_func", "o_func_add", "o_all_func0", "o_all_func1"],
        )
        self._verify_function_set(
            extracted_with_o_all_funcion2,
            {func_add_name, func_identity_name, func_nested_identity_add_name},
            func_domain,
        )


if __name__ == "__main__":
    unittest.main()