# Owner(s): ["module: onnx"]
"""Unit tests for the internal registration wrapper module."""
from __future__ import annotations

import logging
import operator
from typing import TypeVar, Union

import onnxscript  # type: ignore[import]

import torch
import torch.fx
from onnxscript import BFLOAT16, DOUBLE, FLOAT, FLOAT16  # type: ignore[import]
from onnxscript.function_libs.torch_lib import ops  # type: ignore[import]
from onnxscript.onnx_opset import opset15 as op  # type: ignore[import]
from torch.onnx._internal.diagnostics import infra
from torch.onnx._internal.fx import (
    analysis,
    diagnostics,
    onnxfunction_dispatcher,
    registration,
)
from torch.testing._internal import common_utils

# TODO: this can only be global. https://github.com/microsoft/onnxscript/issues/805
TCustomFloat = TypeVar("TCustomFloat", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16])


class TestRegistration(common_utils.TestCase):
    def setUp(self) -> None:
        self.registry = torch.onnx.OnnxRegistry()
        self.custom_domain = onnxscript.values.Opset(domain="custom", version=1)

    def tearDown(self) -> None:
        internal_name_instance = registration.OpName.from_name_parts(
            namespace="test", op_name="test_op"
        )
        self.registry._registry.pop(internal_name_instance, None)

    def test_register_custom_op_registers_custom_function(self):
        self.assertFalse(self.registry.is_registered_op("test", "test_op", "default"))

        @onnxscript.script(self.custom_domain)
        def custom_add(x, y):
            return op.Add(x, y)

        self.registry.register_op(custom_add, "test", "test_op", "default")
        self.assertTrue(self.registry.is_registered_op("test", "test_op", "default"))

        # Test on get_ops
        function_group = self.registry.get_op_functions("test", "test_op", "default")
        self.assertIsNotNone(function_group)
        self.assertEqual({func.onnx_function for func in function_group}, {custom_add})  # type: ignore[arg-type]

    def test_custom_onnx_symbolic_joins_existing_function(self):
        self.assertFalse(self.registry.is_registered_op("test", "test_op"))

        @onnxscript.script(self.custom_domain)
        def test_original(x, y):
            return op.Add(x, y)

        # default has to be specified, as we are not using the registration.OpName
        internal_name_instance = registration.OpName.from_name_parts(
            namespace="test", op_name="test_op", overload="default"
        )
        symbolic_fn = registration.ONNXFunction(
            test_original, op_full_name=internal_name_instance.qualified_name()
        )
        self.registry._register(internal_name_instance, symbolic_fn)
        self.assertTrue(self.registry.is_registered_op("test", "test_op"))

        @onnxscript.script(self.custom_domain)
        def test_custom(x, y):
            return op.Add(x, y)

        self.registry.register_op(test_custom, "test", "test_op")

        function_group = self.registry.get_op_functions("test", "test_op")
        assert function_group is not None
        # The order does matter (list)
        self.assertEqual(
            [func.onnx_function for func in function_group],
            [test_original, test_custom],
        )

    def test_unsupported_nodes_analysis_with_missing_aten_op(self):
        # NOTE: simulate unsupported nodes
        aten_mul_tensor = registration.OpName.from_name_parts(
            namespace="aten", op_name="mul", overload="Tensor"
        )
        aten_mul_default = registration.OpName.from_name_parts(
            namespace="aten", op_name="mul"
        )
        aten_add_tensor = registration.OpName.from_name_parts(
            namespace="aten", op_name="add", overload="Tensor"
        )
        aten_add_default = registration.OpName.from_name_parts(
            namespace="aten", op_name="add"
        )

        self.registry._registry.pop(aten_mul_tensor)
        self.registry._registry.pop(aten_mul_default)
        self.registry._registry.pop(aten_add_tensor)
        self.registry._registry.pop(aten_add_default)

        diagnostic_context = diagnostics.DiagnosticContext(
            "torch.onnx.dynamo_export", torch.__version__
        )
        dispatcher = onnxfunction_dispatcher.OnnxFunctionDispatcher(
            self.registry, diagnostic_context
        )

        graph: torch.fx.Graph = torch.fx.Graph()
        x: torch.fx.Node = graph.create_node("placeholder", "x")
        x.meta["val"] = torch.tensor(3.0)
        b: torch.fx.Node = graph.create_node(
            "call_function", target=torch.ops.aten.mul.Tensor, args=(x, x)
        )
        c: torch.fx.Node = graph.create_node(
            "call_function", target=torch.ops.aten.add.Tensor, args=(b, b)
        )
        output: torch.fx.Node = graph.output(c)
        module = torch.fx.GraphModule(torch.nn.Module(), graph)

        with self.assertRaises(infra.RuntimeErrorWithDiagnostic):
            analysis.UnsupportedFxNodesAnalysis(
                diagnostic_context, module, dispatcher
            ).analyze(infra.levels.ERROR)

        try:
            analysis.UnsupportedFxNodesAnalysis(
                diagnostic_context, module, dispatcher
            ).analyze(infra.levels.ERROR)
        except infra.RuntimeErrorWithDiagnostic as e:
            self.assertIn(
                "Unsupported FX nodes: {'call_function': ['aten.mul.Tensor', 'aten.add.Tensor']}.",
                e.diagnostic.message,
            )


@common_utils.instantiate_parametrized_tests
class TestDispatcher(common_utils.TestCase):
    def setUp(self):
        self.registry = torch.onnx.OnnxRegistry()
        # TODO: remove this once we have a better way to do this
        logger = logging.getLogger("TestDispatcher")
        self.diagnostic_context = diagnostics.DiagnosticContext(
            "torch.onnx.dynamo_export", torch.__version__
        )
        self.dispatcher = onnxfunction_dispatcher.OnnxFunctionDispatcher(
            self.registry, self.diagnostic_context
        )

    @common_utils.parametrize(
        "node, expected_name",
        [
            common_utils.subtest(
                (
                    torch.fx.Node(
                        graph=torch.fx.Graph(),
                        name="aten::add.Tensor",
                        op="call_function",
                        target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                        args=(torch.tensor(3), torch.tensor(4)),
                        kwargs={},
                    ),
                    ("aten", "add", "Tensor"),
                ),
                name="get_Opoverload_name",
            ),
            common_utils.subtest(
                (
                    torch.fx.Node(
                        graph=torch.fx.Graph(),
                        name="aten::sym_size",
                        op="call_function",
                        target=torch.ops.aten.sym_size,
                        args=(),
                        kwargs={},
                    ),
                    ("aten", "sym_size", None),
                ),
                name="get_Opoverloadpacket_name",
            ),
            common_utils.subtest(
                (
                    torch.fx.Node(
                        graph=torch.fx.Graph(),
                        name="builtin_add",
                        op="call_function",
                        target=operator.add,
                        args=(1, 2),
                        kwargs={},
                    ),
                    ("_operator", "add", None),
                ),
                name="get_builtin_op_name",
            ),
        ],
    )
    def test_get_aten_name_on_supported_fx_node(
        self, node: torch.fx.Node, expected_name: str
    ):
        expected_name_class = registration.OpName.from_name_parts(*expected_name)
        self.assertEqual(
            self.dispatcher._get_aten_name(node, self.diagnostic_context),
            expected_name_class,
        )

    @common_utils.parametrize(
        "node",
        [
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add",
                    op="call_function",
                    target=torch.ops.aten.add,
                    args=(),
                    kwargs={},
                ),
                name="unsupported_Opoverloadpacket_name",
            ),
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="builtin_add",
                    op="call_function",
                    target=operator.add,
                    args=("A", "B"),
                    kwargs={},
                ),
                name="unsupported_input_dtypes_for_builtin_op",
            ),
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::made_up_node",
                    op="call_function",
                    target=lambda: None,
                    args=(),
                    kwargs={},
                ),
                name="unsupported_target_function",
            ),
        ],
    )
    def test_get_aten_name_on_unsupported_fx_node(self, node: torch.fx.Node):
        with self.assertRaises(RuntimeError):
            self.dispatcher._get_aten_name(node, self.diagnostic_context)

    def test_get_function_overloads_gives_overload_fall_back_default(self):
        # Test fall back to default op name
        node_overload = torch.fx.Node(
            graph=torch.fx.Graph(),
            name="aten::add.Tensor",
            op="call_function",
            target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
            args=(torch.tensor(3), torch.tensor(4)),
            kwargs={},
        )
        node_overloadpacket = torch.fx.Node(
            graph=torch.fx.Graph(),
            name="aten::add",
            op="call_function",
            target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
            args=(),
            kwargs={},
        )

        self.assertEqual(
            self.dispatcher.get_function_overloads(
                node_overload, self.diagnostic_context
            ),
            self.dispatcher.get_function_overloads(
                node_overloadpacket,
                self.diagnostic_context,
            ),
        )

        # Non-registered op
        unsupported_op_node = torch.fx.Node(
            graph=torch.fx.Graph(),
            name="aten::made_up_node",
            op="call_function",
            target=lambda: None,
            args=(),
            kwargs={},
        )
        with self.assertRaises(RuntimeError):
            self.dispatcher.get_function_overloads(
                unsupported_op_node,
                self.diagnostic_context,
            )

    @common_utils.parametrize(
        "node",
        [
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add.Tensor",
                    op="call_function",
                    target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                    args=(torch.tensor(3.0), torch.tensor(4.0)),
                    kwargs={},
                ),
                name="nearest_match",
            ),
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add.Tensor",
                    op="call_function",
                    target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                    args=(torch.tensor(3.0), torch.tensor(4.0)),
                    kwargs={"alpha": 1},
                ),
                name="perfect_match_with_kwargs",
            ),
        ],
    )
    def test_find_the_perfect_or_nearest_match_onnxfunction_gives_custom_ops_precedence(
        self, node
    ):
        custom_domain = onnxscript.values.Opset(domain="custom", version=1)

        @onnxscript.script(custom_domain)
        def test_custom_op(
            x: TCustomFloat, y: TCustomFloat, alpha: int = 1
        ) -> TCustomFloat:
            return op.Add(x, y)

        @onnxscript.script(custom_domain)
        def test_default_op(
            x: TCustomFloat, y: TCustomFloat, alpha: int = 1
        ) -> TCustomFloat:
            return op.Add(x, y)

        op_full_name = "test::test_op"

        custom_overloads = [
            registration.ONNXFunction(
                test_custom_op, op_full_name=op_full_name, is_custom=True
            )
        ]
        function_overloads = [
            registration.ONNXFunction(test_default_op, op_full_name=op_full_name)
        ] + custom_overloads

        symbolic_fn = self.dispatcher._find_the_perfect_or_nearest_match_onnxfunction(
            node,
            function_overloads,
            node.args,
            node.kwargs,
            self.diagnostic_context,
        )
        self.assertEqual(symbolic_fn, test_custom_op)

    @common_utils.parametrize(
        "node",
        [
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add.Tensor",
                    op="call_function",
                    target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                    args=(torch.tensor(3.0), torch.tensor(4.0)),
                    kwargs={"attr": None},
                ),
                name="perfect_match_with_ignoring_none_attribute",
            ),
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add.Tensor",
                    op="call_function",
                    target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                    args=(torch.tensor(3.0), torch.tensor(4.0)),
                    kwargs={"unrelated": None},
                ),
                name="perfect_match_with_ignoring_unrelated_none_attribute",
            ),
        ],
    )
    def test_find_the_perfect_or_nearest_match_onnxfunction_ignores_attribute_with_none(
        self, node
    ):
        custom_domain = onnxscript.values.Opset(domain="custom", version=1)

        @onnxscript.script(custom_domain)
        def test_op_attribute(
            x: TCustomFloat, y: TCustomFloat, attr: int
        ) -> TCustomFloat:
            return op.Add(x, y)

        @onnxscript.script(custom_domain)
        def test_op(x: TCustomFloat, y: TCustomFloat) -> TCustomFloat:
            return op.Add(x, y)

        op_full_name = "test::test_op"

        function_overloads = [
            registration.ONNXFunction(test_op_attribute, op_full_name=op_full_name),
            registration.ONNXFunction(test_op, op_full_name=op_full_name),
        ]

        symbolic_fn = self.dispatcher._find_the_perfect_or_nearest_match_onnxfunction(
            node,
            function_overloads,
            node.args,
            node.kwargs,
            self.diagnostic_context,
        )
        self.assertEqual(symbolic_fn, test_op)

    @common_utils.parametrize(
        "node",
        [
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add.Tensor",
                    op="call_function",
                    target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                    args=(torch.tensor(3.0), torch.tensor(4.0)),
                    kwargs={},
                ),
                name="nearest_match",
            ),
            common_utils.subtest(
                torch.fx.Node(
                    graph=torch.fx.Graph(),
                    name="aten::add.Tensor",
                    op="call_function",
                    target=torch.ops.aten.add.Tensor,  # type: ignore[attr-defined]
                    args=(torch.tensor(3.0), torch.tensor(4.0)),
                    kwargs={"alpha": 1},
                ),
                name="perfect_match_with_kwargs",
            ),
        ],
    )
    def test_find_the_perfect_or_nearest_match_onnxfunction_gives_tie_breaks_to_registered_order(
        self, node
    ):
        custom_domain = onnxscript.values.Opset(domain="custom", version=1)

        @onnxscript.script(custom_domain)
        def test_second_custom_op(
            x: TCustomFloat, y: TCustomFloat, alpha: int = 1
        ) -> TCustomFloat:
            return op.Add(x, y)

        @onnxscript.script(custom_domain)
        def test_third_custom_op(
            x: TCustomFloat, y: TCustomFloat, alpha: int = 1
        ) -> TCustomFloat:
            return op.Add(x, y)

        @onnxscript.script(custom_domain)
        def test_first_custom_op(
            x: TCustomFloat, y: TCustomFloat, alpha: int = 1
        ) -> TCustomFloat:
            return op.Add(x, y)

        op_full_name = "aten::add"

        function_overloads = [
            registration.ONNXFunction(
                test_first_custom_op, op_full_name=op_full_name, is_custom=True
            ),
            registration.ONNXFunction(
                test_second_custom_op, op_full_name=op_full_name, is_custom=True
            ),
            registration.ONNXFunction(
                test_third_custom_op, op_full_name=op_full_name, is_custom=True
            ),
        ]

        symbolic_fn = self.dispatcher._find_the_perfect_or_nearest_match_onnxfunction(
            node,
            function_overloads,
            node.args,
            node.kwargs,
            self.diagnostic_context,
        )
        self.assertEqual(symbolic_fn, test_third_custom_op)


@common_utils.instantiate_parametrized_tests
class TestOpSchemaWrapper(common_utils.TestCase):
    def setUp(self):
        # overload type: optional dtype
        self.onnx_function_new_full = ops.core.aten_new_full
        self.onnx_function_new_full_dtype = ops.core.aten_new_full_dtype

    @common_utils.parametrize(
        "inputs, attributes, assertion",
        [
            common_utils.subtest(
                ([torch.randn(3, 4), torch.randn(3, 4)], {"alpha": 2.0}, True),
                name="perfect_match_with_kwargs",
            ),
            common_utils.subtest(
                (["A", "B"], {}, False),
                name="non_perfect_match_due_to_non_tensor_inputs",
            ),
            common_utils.subtest(
                ([torch.randn(3, 4), torch.randn(3, 4), torch.randn(3, 4)], {}, False),
                name="non_perfect_match_due_to_too_many_inputs",
            ),
            common_utils.subtest(
                ([torch.randn(3, 4), torch.randn(3, 4)], {"wrong_kwargs": 2.0}, False),
                name="non_perfect_match_due_to_wrong_kwargs",
            ),
        ],
    )
    def test_perfect_match_inputs(self, inputs, attributes, assertion):
        # OnnxFunction with default attributes
        dummy_diagnostic = diagnostics.Diagnostic(
            rule=diagnostics.rules.find_opschema_matched_symbolic_function,
            level=diagnostics.levels.WARNING,
        )
        op_schema_wrapper_add = onnxfunction_dispatcher._OnnxSchemaChecker(
            ops.core.aten_add
        )
        self.assertEqual(
            op_schema_wrapper_add.perfect_match_inputs(
                dummy_diagnostic, inputs, attributes
            ),
            assertion,
        )

    @common_utils.parametrize(
        "inputs, kwargs, op, score",
        [
            common_utils.subtest(
                ([torch.randn(3, 4), torch.randn(3, 4)], {}, ops.core.aten_mul, 2),
                name="match_2_inputs",
            ),
            common_utils.subtest(
                (
                    [
                        torch.randint(0, 2, size=(3, 4), dtype=torch.int).bool(),
                        torch.randint(0, 2, size=(3, 4), dtype=torch.int).bool(),
                    ],
                    {},
                    ops.core.aten_mul,
                    0,
                ),
                name="match_0_inputs",
            ),
            common_utils.subtest(
                ([torch.randn(3, 4), torch.randn(3, 4)], {}, ops.core.aten_mul_bool, 0),
                name="match_0_inputs_bool",
            ),
            common_utils.subtest(
                (
                    [
                        torch.randint(0, 2, size=(3, 4), dtype=torch.int).bool(),
                        torch.randint(0, 2, size=(3, 4), dtype=torch.int).bool(),
                    ],
                    {},
                    ops.core.aten_mul_bool,
                    2,
                ),
                name="match_2_inputs_bool",
            ),
        ],
    )
    def test_matching_score_system_on_overload_dtypes(self, inputs, kwargs, op, score):
        op_schema_wrapper = onnxfunction_dispatcher._OnnxSchemaChecker(op)
        op_schema_wrapper._record_matching_score(inputs, kwargs)
        self.assertEqual(op_schema_wrapper.match_score, score)

    @common_utils.parametrize(
        "inputs, kwargs, op, score",
        [
            common_utils.subtest(
                ([torch.randn(3, 4), torch.tensor(3)], {}, ops.core.aten_new_full, 2),
                name="match_2_inputs",
            ),
            common_utils.subtest(
                (
                    [torch.randn(3, 4), torch.tensor(3)],
                    {"dtype": 2},  # at this point, dtype should be converted to int
                    ops.core.aten_new_full_dtype,
                    2,
                ),
                name="match_2_input_and_match_1_kwargs_optional",
            ),
        ],
    )
    def test_matching_score_system_on_optional_dtypes(self, inputs, kwargs, op, score):
        op_schema_wrapper = onnxfunction_dispatcher._OnnxSchemaChecker(op)
        op_schema_wrapper._record_matching_score(inputs, kwargs)
        self.assertEqual(op_schema_wrapper.match_score, score)

    @common_utils.parametrize(
        "value, expected_onnx_str_dtype",
        [
            common_utils.subtest(
                (1, {"tensor(int64)", "tensor(int16)", "tensor(int32)"}),
                name="all_ints",
            ),
            common_utils.subtest(
                (1.0, {"tensor(float)", "tensor(double)", "tensor(float16)"}),
                name="all_floats",
            ),
            common_utils.subtest(
                (torch.tensor([True]), {"tensor(bool)"}),
                name="bool",
            ),
            common_utils.subtest(
                (torch.tensor([1], dtype=torch.int64), {"tensor(int64)"}),
                name="int64",
            ),
            common_utils.subtest(
                (torch.tensor([1], dtype=torch.int32), {"tensor(int32)"}),
                name="int32",
            ),
            common_utils.subtest(
                (torch.tensor([1], dtype=torch.int16), {"tensor(int16)"}),
                name="int16",
            ),
            common_utils.subtest(
                (torch.tensor([1], dtype=torch.float), {"tensor(float)"}),
                name="float",
            ),
            common_utils.subtest(
                (torch.tensor([1], dtype=torch.float16), {"tensor(float16)"}),
                name="float16",
            ),
            common_utils.subtest(
                (torch.tensor([1], dtype=torch.double), {"tensor(double)"}),
                name="double",
            ),
            common_utils.subtest((None, set()), name="None"),  # None allows no dtype
            common_utils.subtest(
                ([], set()), name="empaty_list"
            ),  # Empty list allows no dtype
        ],
    )
    def test_find_onnx_data_type(self, value, expected_onnx_str_dtype):
        self.assertEqual(
            onnxfunction_dispatcher._find_onnx_data_type(value), expected_onnx_str_dtype
        )


if __name__ == "__main__":
    common_utils.run_tests()
