import unittest
from onnx_light.ext_test_case import ExtTestCase
import onnx
import onnx.defs as onnx_defs
import onnx_light.onnx.defs


class TestSchemaSyncWithOnnxDefs(ExtTestCase):
    # SwiGLU (onnx#8202) may not yet be present in the installed onnx; the empty
    # domain string is the default ai.onnx domain used by the registered schemas.
    SWIGLU_KEY = ("", "SwiGLU", 28)

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        onnx_light.onnx.defs.register_onnx_operator_set_schema()

    def test_onnx_light_ir_and_opset_versions_match_onnx(self):
        self.assertEqual(onnx_light.onnx.defs.onnx_ir_version(), onnx.IR_VERSION)
        self.assertGreaterEqual(
            onnx_light.onnx.defs.onnx_opset_version(), onnx_defs.onnx_opset_version()
        )

    def test_registered_onnx_ops_match_onnx_registered(self):
        flex_attention_key = ("ai.onnx.preview", "FlexAttention", 1)
        onnx_light_schema_keys = {
            (schema.domain, schema.name, schema.since_version)
            for schema in onnx_light.onnx.defs.get_all_schemas_with_history()
        }
        onnx_schema_keys = {
            (schema.domain, schema.name, schema.since_version)
            for schema in onnx_defs.get_all_schemas_with_history()
        }

        if flex_attention_key not in onnx_schema_keys:
            onnx_light_schema_keys.discard(flex_attention_key)
        if self.SWIGLU_KEY not in onnx_schema_keys:
            onnx_light_schema_keys.discard(self.SWIGLU_KEY)

        self.assertEqual(onnx_light_schema_keys, onnx_schema_keys)

    def test_registered_onnx_ops_match_onnx_match_input_output(self):
        light_hist = onnx_light.onnx.defs.get_all_schemas_with_history()
        hist = onnx.defs.get_all_schemas_with_history()
        light_dict = {(s.domain, s.name, s.since_version): s for s in light_hist}
        onnx_dict = {(s.domain, s.name, s.since_version): s for s in hist}
        if self.SWIGLU_KEY not in onnx_dict:
            light_dict.pop(self.SWIGLU_KEY, None)
        self.assertEqual(set(light_dict), set(onnx_dict))
        for key, schema in light_dict.items():
            with self.subTest(key=key):
                lights = light_dict[key]
                self.assertGreater(len(schema.outputs), 0)
                self.assertGreater(len(lights.outputs), 0)
                self.assertGreater(schema.line, 0)
                self.assertGreater(lights.line, 0)
                self.assertEqual(len(schema.inputs), len(lights.inputs))
                self.assertEqual(len(schema.outputs), len(lights.outputs))
                self.assertEqual(len(schema.type_constraints), len(lights.type_constraints))
                self.assertEqual(schema.file, lights.file)
                self.assertEqual(schema.has_function, lights.has_function)
                self.assertEqual(schema.deprecated, lights.deprecated)
                self.assertEqual(schema.is_infinite(0), lights.is_infinite(0))
                self.assertEqual(schema.max_input, lights.max_input)
                self.assertEqual(schema.max_output, lights.max_output)
                # file and line are build-environment-specific and intentionally not compared
                self.assertEqual(schema.non_deterministic, lights.non_deterministic)
                self.assertEqual(schema.support_level, lights.support_level)
                self.assertEqual(schema.doc, lights.doc)

    def test_registered_onnx_ops_match_onnx_match_input_output_doc(self):
        light_hist = onnx_light.onnx.defs.get_all_schemas_with_history()
        hist = onnx.defs.get_all_schemas_with_history()
        light_dict = {(s.domain, s.name, s.since_version): s for s in light_hist}
        onnx_dict = {(s.domain, s.name, s.since_version): s for s in hist}
        if self.SWIGLU_KEY not in onnx_dict:
            light_dict.pop(self.SWIGLU_KEY, None)
        self.assertEqual(set(light_dict), set(onnx_dict))

        def _normalize(text):
            # Whitespace differences (spaces, tabs, newlines) are tolerated.
            return " ".join((text or "").split())

        for key, lights in light_dict.items():
            schema = onnx_dict[key]
            with self.subTest(key=key):
                for index, (onnx_in, light_in) in enumerate(zip(schema.inputs, lights.inputs)):
                    with self.subTest(input=index, name=onnx_in.name):
                        self.assertEqual(onnx_in.name, light_in.name)
                        self.assertEqual(
                            _normalize(onnx_in.description), _normalize(light_in.description)
                        )
                for index, (onnx_out, light_out) in enumerate(
                    zip(schema.outputs, lights.outputs)
                ):
                    with self.subTest(output=index, name=onnx_out.name):
                        self.assertEqual(onnx_out.name, light_out.name)
                        self.assertEqual(
                            _normalize(onnx_out.description), _normalize(light_out.description)
                        )

    def test_registered_onnx_ops_match_onnx_match_attributes(self):
        light_hist = onnx_light.onnx.defs.get_all_schemas_with_history()
        hist = onnx.defs.get_all_schemas_with_history()
        light_dict = {(s.domain, s.name, s.since_version): s for s in light_hist}
        onnx_dict = {(s.domain, s.name, s.since_version): s for s in hist}
        if self.SWIGLU_KEY not in onnx_dict:
            light_dict.pop(self.SWIGLU_KEY, None)
        self.assertEqual(set(light_dict), set(onnx_dict))
        for key, lights in light_dict.items():
            schema = onnx_dict[key]
            with self.subTest(key=key):
                self.assertEqual(set(schema.attributes), set(lights.attributes))
                for name in schema.attributes:
                    onnx_attr = schema.attributes[name]
                    light_attr = lights.attributes[name]
                    with self.subTest(attribute=name):
                        self.assertEqual(onnx_attr.name, light_attr.name)
                        self.assertEqual(onnx_attr.required, light_attr.required)
                        # AttrType in onnx and AttributeType in onnx_light share enum
                        # names ("INT", "FLOAT", ...).
                        self.assertEqual(onnx_attr.type.name, light_attr.type.name)
                        self.assertEqual(onnx_attr.description, light_attr.description)


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