import unittest

from onnx_light.ext_test_case import import_or_skip

import onnx_light.onnx as onnxl
import onnx_light.onnx.shape_inference as shape_inference

# The backend test registries are only available in the full build; skip this
# module on a reduced build (ONNX_LIGHT_BUILD_KERNELS=OFF).
make_test_class = import_or_skip("onnx_light.onnx_lib.backend.test.case", "make_test_class")


def _check_match(expected, inferred) -> None:
    expected_dict = {vi.name: vi for vi in expected}
    inferred_dict = {vi.name: vi for vi in inferred}
    for name, value in expected_dict.items():
        assert name in inferred_dict, f"Missing shape for {name!r}"
        infe = inferred_dict[name]
        if not value.type.tensor_type:
            # not implemented yet
            continue
        assert (
            value.type.tensor_type and infe.type.tensor_type
        ), f"Missing tensor_type for {name!r}"
        vt = value.type.tensor_type
        it = infe.type.tensor_type
        assert vt.elem_type == vt.elem_type, f"Element type mismatch for {name!r}: {vt} != {it}"
        v_shape = tuple(d.dim_param or d.dim_value for d in vt.shape.dim)
        i_shape = tuple(d.dim_param or d.dim_value for d in it.shape.dim)
        assert v_shape == i_shape, f"Element shape mismatch for {name!r}: {vt} != {it}"


def shape_inference_check(model: onnxl.ModelProto, *inputs):
    assert isinstance(model, onnxl.ModelProto), f"Unexpected type {type(model)}"
    work = onnxl.ModelProto()
    work.CopyFrom(model)
    if work.graph.value_info:
        work.graph.value_info.clear()
    shape_inference.infer_shapes(work)
    _check_match(model.graph.value_info, work.graph.value_info)
    _check_match(model.graph.output, work.graph.output)


TestShapeInferenceBackend = make_test_class(
    shape_inference_check,
    exclude_regex=[
        "test_cc_shape_inference_add_concat_reshape.*",
        "test_cc_shape_inference_nonzero_chain_anon.*",
        "test_cc_shape_inference_nonzero_chain_named.*",
        "test_cc_shape_inference_nonzero_plus_expression.*",
        "test_cc_shape_inference_value_as_shape.*",
        "test_cc_shape_inference_concat_split.*",
        "test_cc_shape_inference_reshape_reshape.*",
        "test_cc_shape_inference_check_shape.*",
        "test_cc_shape_inference_scan_running_sum.*",
        # ONNX's built-in shape inference gives generic unk__N dims for
        # Resize output; it does not understand onnx-light's symbolic names.
        "test_cc_shape_inference_resize_tile.*",
        # ONNX's built-in shape inference emits a generic unk__N dim for the
        # Slice(axis=2, starts=0, ends=-1) symbolic-end length instead of c-1.
        "test_cc_shape_inference_slice_symbolic_end.*",
        # ONNX's built-in shape inference gives generic unk__N dims for the
        # Conv output; it does not recover the symbolic H/W spatial dims that
        # onnx-light's symbolic Pad→Conv propagation collapses back to.
        "test_cc_shape_inference_pad_canny_average.*",
        # TopK's output axis depends on the runtime input K, so ONNX's built-in
        # shape inference emits a generic unk__N dim that does not match the
        # symbolic ``k`` placeholder stored in value_info. The ``scan_``/``loop_``
        # variants additionally emit a generic ``unk__N`` dim for the Scan/Loop
        # stacked axis that does not match onnx-light's symbolic names.
        "test_cc_shape_inference_topk_pairwise_distance.*",
        "test_cc_shape_inference_scan_topk_pairwise_distance.*",
        "test_cc_shape_inference_loop_topk_pairwise_distance.*",
        # ONNX's built-in shape inference emits generic unk__N dims for chained
        # TopK nodes; it does not preserve the symbolic dim names that
        # onnx-light assigns (e.g. ``TopK_k``).
        "test_cc_shape_inference_two_topk_same_k.*",
        "test_cc_shape_inference_two_topk_different_k.*",
        # ONNX's built-in shape inference propagates "batch" through the Loop
        # body but onnx-light's inference uses auto-generated "unk__N"/"N" dims
        # for the intermediate pairwise-distance tensor.
        "test_cc_shape_inference_loop_pairwise_distance.*",
        # ONNX's built-in shape inference does not recover the symbolic B1/B2
        # dims that onnx-light propagates through the If→Abs/Neg chain.
        "test_cc_shape_inference_if_symbolic_shapes.*",
        # ONNX's built-in shape inference does not propagate value-as-shape
        # through Gather: Expand(y, Gather(Shape(x), [0])) produces unk__0
        # instead of the symbolic N from x's leading axis.
        "test_cc_shape_inference_gather_value_as_shape.*",
        # ONNX's built-in shape inference does not propagate value-as-shape
        # through the Shape→Gather→Unsqueeze→Concat chain, so Reshape produces
        # unk__0/unk__1 instead of the symbolic D1/D2 dims from the inputs.
        "test_cc_shape_inference_unsqueeze_vas_reshape.*",
    ],
)


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