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

from typing import Sequence

from torch.onnx import errors
from torch.onnx._internal import registration
from torch.testing._internal import common_utils


@common_utils.instantiate_parametrized_tests
class TestGlobalHelpers(common_utils.TestCase):
    @common_utils.parametrize(
        "available_opsets, target, expected",
        [
            ((7, 8, 9, 10, 11), 16, 11),
            ((7, 8, 9, 10, 11), 11, 11),
            ((7, 8, 9, 10, 11), 10, 10),
            ((7, 8, 9, 10, 11), 9, 9),
            ((7, 8, 9, 10, 11), 8, 8),
            ((7, 8, 9, 10, 11), 7, 7),
            ((9, 10, 16), 16, 16),
            ((9, 10, 16), 15, 10),
            ((9, 10, 16), 10, 10),
            ((9, 10, 16), 9, 9),
            ((9, 10, 16), 8, 9),
            ((9, 10, 16), 7, 9),
            ((7, 9, 10, 16), 16, 16),
            ((7, 9, 10, 16), 10, 10),
            ((7, 9, 10, 16), 9, 9),
            ((7, 9, 10, 16), 8, 9),
            ((7, 9, 10, 16), 7, 7),
            ([17], 16, None),  # New op added in 17
            ([9], 9, 9),
            ([9], 8, 9),
            ([], 16, None),
            ([], 9, None),
            ([], 8, None),
            # Ops registered at opset 1 found as a fallback when target >= 9
            ([1], 16, 1),
        ],
    )
    def test_dispatch_opset_version_returns_correct_version(
        self, available_opsets: Sequence[int], target: int, expected: int
    ):
        actual = registration._dispatch_opset_version(target, available_opsets)
        self.assertEqual(actual, expected)


class TestOverrideDict(common_utils.TestCase):
    def setUp(self):
        self.override_dict: registration.OverrideDict[
            str, int
        ] = registration.OverrideDict()

    def test_get_item_returns_base_value_when_no_override(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.set_base("b", 0)

        self.assertEqual(self.override_dict["a"], 42)
        self.assertEqual(self.override_dict["b"], 0)
        self.assertEqual(len(self.override_dict), 2)

    def test_get_item_returns_overridden_value_when_override(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.set_base("b", 0)
        self.override_dict.override("a", 100)
        self.override_dict.override("c", 1)

        self.assertEqual(self.override_dict["a"], 100)
        self.assertEqual(self.override_dict["b"], 0)
        self.assertEqual(self.override_dict["c"], 1)
        self.assertEqual(len(self.override_dict), 3)

    def test_get_item_raises_key_error_when_not_found(self):
        self.override_dict.set_base("a", 42)

        with self.assertRaises(KeyError):
            self.override_dict["nonexistent_key"]

    def test_get_returns_overridden_value_when_override(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.set_base("b", 0)
        self.override_dict.override("a", 100)
        self.override_dict.override("c", 1)

        self.assertEqual(self.override_dict.get("a"), 100)
        self.assertEqual(self.override_dict.get("b"), 0)
        self.assertEqual(self.override_dict.get("c"), 1)
        self.assertEqual(len(self.override_dict), 3)

    def test_get_returns_none_when_not_found(self):
        self.override_dict.set_base("a", 42)

        self.assertEqual(self.override_dict.get("nonexistent_key"), None)

    def test_in_base_returns_true_for_base_value(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.set_base("b", 0)
        self.override_dict.override("a", 100)
        self.override_dict.override("c", 1)

        self.assertIn("a", self.override_dict)
        self.assertIn("b", self.override_dict)
        self.assertIn("c", self.override_dict)

        self.assertTrue(self.override_dict.in_base("a"))
        self.assertTrue(self.override_dict.in_base("b"))
        self.assertFalse(self.override_dict.in_base("c"))
        self.assertFalse(self.override_dict.in_base("nonexistent_key"))

    def test_overridden_returns_true_for_overridden_value(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.set_base("b", 0)
        self.override_dict.override("a", 100)
        self.override_dict.override("c", 1)

        self.assertTrue(self.override_dict.overridden("a"))
        self.assertFalse(self.override_dict.overridden("b"))
        self.assertTrue(self.override_dict.overridden("c"))
        self.assertFalse(self.override_dict.overridden("nonexistent_key"))

    def test_remove_override_removes_overridden_value(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.set_base("b", 0)
        self.override_dict.override("a", 100)
        self.override_dict.override("c", 1)

        self.assertEqual(self.override_dict["a"], 100)
        self.assertEqual(self.override_dict["c"], 1)

        self.override_dict.remove_override("a")
        self.override_dict.remove_override("c")
        self.assertEqual(self.override_dict["a"], 42)
        self.assertEqual(self.override_dict.get("c"), None)
        self.assertFalse(self.override_dict.overridden("a"))
        self.assertFalse(self.override_dict.overridden("c"))

    def test_remove_override_removes_overridden_key(self):
        self.override_dict.override("a", 100)
        self.assertEqual(self.override_dict["a"], 100)
        self.assertEqual(len(self.override_dict), 1)
        self.override_dict.remove_override("a")
        self.assertEqual(len(self.override_dict), 0)
        self.assertNotIn("a", self.override_dict)

    def test_overriden_key_precededs_base_key_regardless_of_insert_order(self):
        self.override_dict.set_base("a", 42)
        self.override_dict.override("a", 100)
        self.override_dict.set_base("a", 0)

        self.assertEqual(self.override_dict["a"], 100)
        self.assertEqual(len(self.override_dict), 1)

    def test_bool_is_true_when_not_empty(self):
        if self.override_dict:
            self.fail("OverrideDict should be false when empty")
        self.override_dict.override("a", 1)
        if not self.override_dict:
            self.fail("OverrideDict should be true when not empty")
        self.override_dict.set_base("a", 42)
        if not self.override_dict:
            self.fail("OverrideDict should be true when not empty")
        self.override_dict.remove_override("a")
        if not self.override_dict:
            self.fail("OverrideDict should be true when not empty")


class TestRegistrationDecorators(common_utils.TestCase):
    def tearDown(self) -> None:
        registration.registry._registry.pop("test::test_op", None)

    def test_onnx_symbolic_registers_function(self):
        self.assertFalse(registration.registry.is_registered_op("test::test_op", 9))

        @registration.onnx_symbolic("test::test_op", opset=9)
        def test(g, x):
            return g.op("test", x)

        self.assertTrue(registration.registry.is_registered_op("test::test_op", 9))
        function_group = registration.registry.get_function_group("test::test_op")
        assert function_group is not None
        self.assertEqual(function_group.get(9), test)

    def test_onnx_symbolic_registers_function_applied_decorator_when_provided(self):
        wrapper_called = False

        def decorator(func):
            def wrapper(*args, **kwargs):
                nonlocal wrapper_called
                wrapper_called = True
                return func(*args, **kwargs)

            return wrapper

        @registration.onnx_symbolic("test::test_op", opset=9, decorate=[decorator])
        def test():
            return

        function_group = registration.registry.get_function_group("test::test_op")
        assert function_group is not None
        registered_function = function_group[9]
        self.assertFalse(wrapper_called)
        registered_function()
        self.assertTrue(wrapper_called)

    def test_onnx_symbolic_raises_warning_when_overriding_function(self):
        self.assertFalse(registration.registry.is_registered_op("test::test_op", 9))

        @registration.onnx_symbolic("test::test_op", opset=9)
        def test1():
            return

        with self.assertWarnsRegex(
            errors.OnnxExporterWarning,
            "Symbolic function 'test::test_op' already registered",
        ):

            @registration.onnx_symbolic("test::test_op", opset=9)
            def test2():
                return

    def test_custom_onnx_symbolic_registers_custom_function(self):
        self.assertFalse(registration.registry.is_registered_op("test::test_op", 9))

        @registration.custom_onnx_symbolic("test::test_op", opset=9)
        def test(g, x):
            return g.op("test", x)

        self.assertTrue(registration.registry.is_registered_op("test::test_op", 9))
        function_group = registration.registry.get_function_group("test::test_op")
        assert function_group is not None
        self.assertEqual(function_group.get(9), test)

    def test_custom_onnx_symbolic_overrides_existing_function(self):
        self.assertFalse(registration.registry.is_registered_op("test::test_op", 9))

        @registration.onnx_symbolic("test::test_op", opset=9)
        def test_original():
            return "original"

        self.assertTrue(registration.registry.is_registered_op("test::test_op", 9))

        @registration.custom_onnx_symbolic("test::test_op", opset=9)
        def test_custom():
            return "custom"

        function_group = registration.registry.get_function_group("test::test_op")
        assert function_group is not None
        self.assertEqual(function_group.get(9), test_custom)


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