#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from collections import OrderedDict

import numpy as np
import onnx
import onnx.numpy_helper
import onnx.shape_inference
import pytest
from onnx_graphsurgeon.importers.onnx_importer import OnnxImporter
from onnx_graphsurgeon.ir.tensor import Constant, Tensor, Variable, SparseValues
from onnx_graphsurgeon.ir.function import Function
from onnx_graphsurgeon.ir.node import Node
from onnx_graphsurgeon.logger import G_LOGGER

from onnx_models import (
    dim_param_model,
    ext_weights,
    identity_model,
    initializer_is_output_model,
    lstm_model,
    nested_dup_names,
    scan_model,
    sparse_nnz_model,
    sparse_nnz_rank_model,
)

G_LOGGER.severity = G_LOGGER.ULTRA_VERBOSE


class TestOnnxImporter(object):
    @pytest.mark.parametrize(
        "onnx_type, expected_type",
        [
            (onnx.TensorProto.FLOAT, np.float32),
            (onnx.TensorProto.BFLOAT16, onnx.TensorProto.BFLOAT16),
            (onnx.TensorProto.FLOAT8E4M3FN, onnx.TensorProto.FLOAT8E4M3FN),
            (onnx.TensorProto.FLOAT8E4M3FNUZ, onnx.TensorProto.FLOAT8E4M3FNUZ),
            (onnx.TensorProto.FLOAT8E5M2, onnx.TensorProto.FLOAT8E5M2),
            (onnx.TensorProto.FLOAT8E5M2FNUZ, onnx.TensorProto.FLOAT8E5M2FNUZ),
            (onnx.TensorProto.FLOAT4E2M1, onnx.TensorProto.FLOAT4E2M1),
        ],
    )
    def test_import_variable_tensor(self, onnx_type, expected_type):
        name = "test0"
        shape = (1, 2, 3, 4)
        onnx_tensor = onnx.helper.make_tensor_value_info(name, onnx_type, shape)
        tensor = OnnxImporter.import_tensor(onnx_tensor)
        assert type(tensor) == Variable
        assert tensor.name == name
        assert tensor.dtype == expected_type
        assert tuple(tensor.shape) == shape

    def test_import_constant_tensor(self):
        shape = (3, 3, 3)
        dtype = np.float32
        onnx_tensor = onnx.numpy_helper.from_array(np.ones(shape=shape, dtype=dtype))
        tensor = OnnxImporter.import_tensor(onnx_tensor)
        assert type(tensor) == Constant
        assert tensor.dtype == dtype
        assert tuple(tensor.shape) == shape

    def test_import_tensor_unknown_metadata(self):
        name = "test0"
        onnx_tensor = onnx.helper.make_empty_tensor_value_info(name)
        tensor = OnnxImporter.import_tensor(onnx_tensor)
        assert type(tensor) == Variable
        assert tensor.name == name

    # An empty string in `dim_param` should be treated like a dynamic dimension
    def test_import_empty_dim_param_tensor(self):
        shape = (1, 2, "non-empty", "")
        onnx_tensor = onnx.helper.make_tensor_value_info(
            "test0", onnx.TensorProto.FLOAT, shape
        )
        tensor = OnnxImporter.import_tensor(onnx_tensor)
        assert type(tensor) == Variable
        assert tuple(tensor.shape) == shape

    # Sometimes, tensor shape is not known, in which case we shouldn't import it
    def test_import_unknown_shape_tensor(self):
        shape = None
        onnx_tensor = onnx.helper.make_tensor_value_info(
            "test0", onnx.TensorProto.FLOAT, shape
        )
        tensor = OnnxImporter.import_tensor(onnx_tensor)
        assert type(tensor) == Variable
        assert tensor.shape is None

    # Scalars can be represented in ONNX with a dim that includes neither a dim_param nor dim_value
    def test_import_empty_dim_tensor(self):
        shape = (None,)
        onnx_tensor = onnx.helper.make_tensor_value_info(
            "test0", onnx.TensorProto.FLOAT, shape
        )
        onnx_tensor.type.tensor_type.shape.dim[0].ClearField("dim_value")
        onnx_tensor.type.tensor_type.shape.dim[0].ClearField("dim_param")

        tensor = OnnxImporter.import_tensor(onnx_tensor)
        assert type(tensor) == Variable
        assert tuple(tensor.shape) == shape

    # TODO: Test all attribute types - missing graph
    def test_import_node(self):
        op = "Test"
        inputs = ["x"]
        outputs = ["y"]
        float_attr = 4.0
        int_attr = 10
        str_attr = "constant"
        tensor_vals = np.ones(shape=(1, 2, 3, 4), dtype=np.float32)
        tensor_attr = onnx.numpy_helper.from_array(tensor_vals)
        floats_attr = [1.0, 2.0, 3.0, 4.0]
        ints_attr = [4, 3, 2, 1]
        strings_attr = ["constant", "and", "variable"]

        onnx_node = onnx.helper.make_node(
            op,
            inputs,
            outputs,
            float_attr=float_attr,
            int_attr=int_attr,
            str_attr=str_attr,
            tensor_attr=tensor_attr,
            floats_attr=floats_attr,
            ints_attr=ints_attr,
            strings_attr=strings_attr,
        )
        node = OnnxImporter.import_node(
            onnx_node, OrderedDict(), OrderedDict(), opset=11, import_domains=None
        )
        assert node.op == op
        assert node.attrs["float_attr"] == float_attr
        assert node.attrs["int_attr"] == int_attr
        assert node.attrs["str_attr"] == str_attr
        # Tensor should turn into a Constant
        assert np.all(node.attrs["tensor_attr"].values == tensor_vals)
        assert node.attrs["floats_attr"] == floats_attr
        assert node.attrs["ints_attr"] == ints_attr
        assert node.attrs["strings_attr"] == strings_attr

    def test_import_node_ref_attrs(self):
        op = "Test"
        inputs = ["x"]
        outputs = ["y"]
        attrs = {"attr1": 1, "attr2": 2.0}
        referencing_attr = "attr3"
        referenced_attr = "attr4"

        onnx_node = onnx.helper.make_node(op, inputs, outputs, **attrs)
        onnx_attr_ref = onnx.helper.make_attribute_ref(
            referencing_attr, onnx.AttributeProto.FLOAT
        )
        onnx_attr_ref.ref_attr_name = referenced_attr
        onnx_node.attribute.append(onnx_attr_ref)
        node = OnnxImporter.import_node(
            onnx_node, OrderedDict(), OrderedDict(), opset=11, import_domains=None
        )
        assert node.op == op
        assert node.attrs["attr1"] == 1
        assert node.attrs["attr2"] == 2.0
        assert node.attrs["attr3"] == Node.AttributeRef(referenced_attr, float)

    def test_import_function(self):
        name = "Test"
        domain = "com.test"
        inputs = ["X", "Y"]
        outputs = ["Z"]
        nodes = [
            onnx.helper.make_node("Add", ["X", "Y"], ["V"], attr1=1),
            onnx.helper.make_node("Add", ["X", "V"], ["W"], attr2=2),
            onnx.helper.make_node("Mul", ["V", "W"], ["Z"], attr3=3),
        ]
        opset = 18
        opset_imports = [onnx.helper.make_operatorsetid("ai.onnx", opset)]
        attributes = ["attr1", "attr2"]
        attribute_protos = [onnx.helper.make_attribute("attr3", 3)]
        doc_string = "docstring"
        onnx_function = onnx.helper.make_function(
            domain,
            name,
            inputs,
            outputs,
            nodes,
            opset_imports,
            attributes=attributes,
            attribute_protos=attribute_protos,
            doc_string=doc_string,
        )
        func = OnnxImporter.import_function(onnx_function)
        assert type(func) == Function
        assert func.name == name
        assert func.domain == domain
        assert func.doc_string == doc_string
        assert list(func.import_domains) == list(opset_imports)
        assert set(func.attrs.keys()) == set(attributes) | {
            a.name for a in attribute_protos
        }
        assert func.opset == opset
        assert all([isinstance(t, Tensor) for t in func.inputs + func.outputs])
        assert sorted(inputs) == sorted([t.name for t in func.inputs])
        assert sorted(outputs) == sorted([t.name for t in func.outputs])
        assert sorted([n.op_type for n in nodes]) == sorted([n.op for n in func.nodes])
        assert attribute_protos[0].i == func.attrs[attribute_protos[0].name]

    @pytest.mark.parametrize(
        "model",
        [
            identity_model(),
            lstm_model(),
            scan_model(),
            dim_param_model(),
            initializer_is_output_model(),
            nested_dup_names(),
            ext_weights(),
        ],
        ids=lambda model: str(model),
    )
    def test_import_graph(self, model):
        graph = OnnxImporter.import_graph(model.load().graph)
        model.assert_equal(graph)

    def test_import_graph_value_info(self):
        model = onnx.shape_inference.infer_shapes(identity_model().load())
        graph = OnnxImporter.import_graph(model.graph)
        tensors = graph.tensors()
        assert all(
            [
                type(tensor) == Variable and tensor.dtype is not None and tensor.shape
                for tensor in tensors.values()
            ]
        )

    def test_import_graph_tensor_map_preserved(self):
        model = identity_model()
        tensor_map = OrderedDict()
        graph = OnnxImporter.import_graph(model.load().graph, tensor_map=tensor_map)
        assert len(tensor_map) == 0
        model.assert_equal(graph)

    def test_import_graph_with_initializer(self):
        model = lstm_model()
        graph = OnnxImporter.import_graph(model.load().graph)
        model.assert_equal(graph)

    def test_import_graph_with_dim_param(self):
        model = dim_param_model()
        graph = OnnxImporter.import_graph(model.load().graph)
        model.assert_equal(graph)

    def test_import_graph_with_sparse_nnz_rank(self):
        model = sparse_nnz_rank_model()
        graph = OnnxImporter.import_graph(model.load().graph)
        tensors = graph.tensors()

        assert "w_sparse" in tensors
        sparse_tensor = tensors["w_sparse"]

        ref_value = np.array(
            [
                1.0,
                2.0,
                3.0,
                4.0,
                5.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
            ]
        ).reshape(sparse_tensor._values.shape)
        assert (
            type(sparse_tensor) == Constant
            and type(sparse_tensor._values) == SparseValues
        )
        assert (tensors["w_sparse"]._values.load() == ref_value).all()

    def test_import_graph_with_sparse_nnz(self):
        model = sparse_nnz_model()
        graph = OnnxImporter.import_graph(model.load().graph)
        tensors = graph.tensors()

        assert "w_sparse" in tensors
        sparse_tensor = tensors["w_sparse"]

        ref_value = np.array(
            [
                0.0,
                1.0,
                0.0,
                0.0,
                2.0,
                0.0,
                0.0,
                3.0,
                0.0,
                0.0,
                0.0,
                0.0,
                4.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                5.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
            ]
        ).reshape(sparse_tensor._values.shape)
        assert (
            type(sparse_tensor) == Constant
            and type(sparse_tensor._values) == SparseValues
        )
        assert (tensors["w_sparse"]._values.load() == ref_value).all()

    def test_import_fp4_tensor(self):
        name = "test_fp4"
        shape = (2, 2)
        dtype = onnx.TensorProto.FLOAT4E2M1

        # Four 4-bit values = 16 bits = 2 bytes
        # [0x01, 0x02] should give us the values [0.5, 0, 1, 0] once unpacked
        tensor_bytes = bytes([0x01, 0x02])
        onnx_tensor = onnx.helper.make_tensor(name, dtype, shape, tensor_bytes , raw=True)
        tensor = OnnxImporter.import_tensor(onnx_tensor)
        
        assert type(tensor) == Constant
        assert tensor.name == name
        assert tuple(tensor.shape) == shape

        # Before accessing values we expect the tensor to have an ONNX type
        assert tensor.dtype == dtype

        # After loading LazyValues, numpy represents this as four separate 8-bit values
        # while retaining the shape
        assert np.array_equal(tensor.values.ravel(), [0.5, 0, 1, 0])

        # Post loading the values, we should see the ml_dtypes conversion
        import ml_dtypes
        assert tensor.dtype == ml_dtypes.float4_e2m1fn