# SPDX-License-Identifier: Apache-2.0

from onnx import checker, helper, TensorProto, NodeProto, GraphProto, ValueInfoProto, ModelProto, TensorShapeProto, ONNX_ML
from onnx.helper import make_model, make_node, make_tensor, make_tensor_value_info
from typing import Sequence, Union, Tuple, List, Any, Optional
import onnx.shape_inference
import unittest


class TestSymbolicShape(unittest.TestCase):

    def _assert_valueinfo_shape(self, onnx_model: ModelProto, value_infos: List[ValueInfoProto]) -> None:
        """
            Assert onnx_model.value_info should be the same as expected value_infos
            Instead of exact symbol, use -1 to represent symbolic shape in expected value_infos
        """
        for expected_vi in value_infos:
            shape = self._get_shape_from_name(onnx_model, expected_vi.name)
            assert shape is not None, f'{onnx_model}'
            if expected_vi.type.HasField('tensor_type'):
                expected_shape = expected_vi.type.tensor_type.shape
            elif expected_vi.type.HasField('sparse_tensor_type'):
                expected_shape = expected_vi.type.sparse_tensor_type.shape
            assert len(shape.dim) == len(expected_shape.dim), f'{onnx_model}'
            for dim_i in range(len(shape.dim)):
                dim = shape.dim[dim_i]
                expected_dim = expected_shape.dim[dim_i]
                # -1 means it's a symbolic shape
                if expected_dim.dim_value == -1:
                    # symbolic dimension must exist
                    assert dim.dim_param, f'{onnx_model}'
                else:
                    assert dim.dim_value == expected_dim.dim_value, f'{onnx_model}'

    def _count_unique_dim_param_number(self, onnx_model: ModelProto) -> int:
        """
           return the total number of unique symbolic shape
        """
        symbol_shape_set = set()
        inputs = list(onnx_model.graph.input)
        outputs = list(onnx_model.graph.output)
        valueinfos = list(onnx_model.graph.value_info)
        for v in inputs + outputs + valueinfos:
            for dim in v.type.tensor_type.shape.dim:
                if dim.dim_param:
                    symbol_shape_set.add(dim.dim_param)
        return len(symbol_shape_set)

    def _get_shape_from_name(self, onnx_model: ModelProto, name: str) -> Optional[TensorShapeProto]:
        """
            Get shape from tensor_type or sparse_tensor_type according to given name
        """
        inputs = list(onnx_model.graph.input)
        outputs = list(onnx_model.graph.output)
        valueinfos = list(onnx_model.graph.value_info)
        for v in inputs + outputs + valueinfos:
            if v.name == name:
                if v.type.HasField('tensor_type'):
                    return v.type.tensor_type.shape
                elif v.type.HasField('sparse_tensor_type'):
                    return v.type.sparse_tensor_type.shape
        return None

    def test_concat_enable_symbolic(self) -> None:
        concat = helper.make_node('Concat', inputs=['A', 'B'], outputs=['C'], name='Concat', axis=1)
        cast = onnx.helper.make_node('Cast',
            inputs=['C'],
            outputs=['output'],
            to=getattr(TensorProto, 'FLOAT'))
        graph_def = helper.make_graph(name='test_graph',
            nodes=[concat, cast],
            inputs=[helper.make_tensor_value_info('A', TensorProto.FLOAT, [2, 'A']),
                helper.make_tensor_value_info('B', TensorProto.FLOAT, [2, 3])],
            outputs=[helper.make_tensor_value_info('output', TensorProto.FLOAT, [2, None])]
        )

        onnx_model = make_model(graph_def)
        inferred_model = onnx.shape_inference.infer_shapes(onnx_model, strict_mode=True)
        self._assert_valueinfo_shape(inferred_model, [make_tensor_value_info("C", TensorProto.FLOAT, (2, -1))])
        # the symbolic shape of C and output should be the same
        assert self._get_shape_from_name(inferred_model, 'C') == self._get_shape_from_name(inferred_model, 'output')

    def test_two_symbolic_concat(self) -> None:
        concat1 = helper.make_node('Concat', inputs=['A', 'B'], outputs=['C'], name='Concat', axis=1)
        concat2 = helper.make_node('Concat', inputs=['C', 'D'], outputs=['E'], name='Concat', axis=1)
        cast = onnx.helper.make_node('Cast',
            inputs=['E'],
            outputs=['output'],
            to=getattr(TensorProto, 'FLOAT'))
        graph_def = helper.make_graph(name='test_graph',
            nodes=[concat1, concat2, cast],
            inputs=[helper.make_tensor_value_info('A', TensorProto.FLOAT, [2, 'A']),
                helper.make_tensor_value_info('B', TensorProto.FLOAT, [2, 3]),
                helper.make_tensor_value_info('D', TensorProto.FLOAT, [2, 'D'])],
            outputs=[helper.make_tensor_value_info('output', TensorProto.FLOAT, [2, None])]
        )

        onnx_model = make_model(graph_def)
        inferred_model = onnx.shape_inference.infer_shapes(onnx_model, strict_mode=True)
        self._assert_valueinfo_shape(inferred_model, [
            make_tensor_value_info("C", TensorProto.FLOAT, (2, -1)),
            make_tensor_value_info("E", TensorProto.FLOAT, (2, -1))])
        # the symbolic shape of E and output should be the same
        assert self._get_shape_from_name(inferred_model, 'E') == self._get_shape_from_name(inferred_model, 'output')

    def test_duplicate_symbolic_shape(self) -> None:
        concat1 = helper.make_node('Concat', inputs=['A', 'B'], outputs=['C'], name='Concat', axis=1)
        concat2 = helper.make_node('Concat', inputs=['C', 'D'], outputs=['E'], name='Concat', axis=1)
        cast = onnx.helper.make_node('Cast',
            inputs=['E'],
            outputs=['output'],
            to=getattr(TensorProto, 'FLOAT'))
        graph_def = helper.make_graph(name='test_graph',
            nodes=[concat1, concat2, cast],
            inputs=[helper.make_tensor_value_info('A', TensorProto.FLOAT, [2, 'unk__0']),
                helper.make_tensor_value_info('B', TensorProto.FLOAT, [2, 3]),
                helper.make_tensor_value_info('D', TensorProto.FLOAT, [2, 'unk__1'])],
            outputs=[helper.make_tensor_value_info('output', TensorProto.FLOAT, [2, 'unk__0'])]
        )

        onnx_model = make_model(graph_def)
        original_count = self._count_unique_dim_param_number(onnx_model)
        inferred_model = onnx.shape_inference.infer_shapes(onnx_model, strict_mode=True)
        inferred_count = self._count_unique_dim_param_number(inferred_model)
        # to prevent duplicate so the inferred count will be count + 2
        # new symbol 'unk__2' and 'unk__3' should be generated
        # original: {'unk_0', 'unk__1'}
        # inferred: {'unk_0', 'unk__1', 'unk__2', 'unk__3'}
        assert inferred_count == original_count + 2, f'{inferred_model}{onnx_model}'

    def test_unknown_shape(self) -> None:
        concat = helper.make_node('Concat', inputs=['A', 'B'], outputs=['C'], name='Concat', axis=1)
        cast = onnx.helper.make_node('Cast',
            inputs=['C'],
            outputs=['output'],
            to=getattr(TensorProto, 'FLOAT'))
        graph_def = helper.make_graph(name='test_graph',
            nodes=[concat, cast],
            inputs=[helper.make_tensor_value_info('A', TensorProto.FLOAT, [3, None]),  # unknown shape
                helper.make_tensor_value_info('B', TensorProto.FLOAT, [3, None])],
            outputs=[helper.make_tensor_value_info('output', TensorProto.FLOAT, [3, None])]
        )

        onnx_model = make_model(graph_def)
        inferred_model = onnx.shape_inference.infer_shapes(onnx_model, strict_mode=True)
        self._assert_valueinfo_shape(inferred_model, [make_tensor_value_info("C", TensorProto.FLOAT, (3, -1))])
        # the symbolic shape of C and output should be the same
        # ('unk__0', 'unk__1')
        assert self._get_shape_from_name(inferred_model, 'C') == self._get_shape_from_name(inferred_model, 'output')


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