#
# 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.
#


import copy

import ml_dtypes
import numpy as np
import onnx
import pytest

from onnx_graphsurgeon.ir.graph import Graph
from onnx_graphsurgeon.ir.node import Node
from onnx_graphsurgeon.ir.tensor import Constant, LazyValues, Variable
from onnx_graphsurgeon.logger import G_LOGGER
from onnx_graphsurgeon.util.misc import SynchronizedList

G_LOGGER.severity = G_LOGGER.ULTRA_VERBOSE


class TensorBaseTests(object):
    def test_can_convert_in_place_to_constant(self):
        tensor = self.tensor.to_constant(
            values=np.ones((1, 3, 5, 5), dtype=ml_dtypes.bfloat16)
        )
        assert tensor is self.tensor
        assert isinstance(tensor, Constant)
        assert isinstance(self.input_node.outputs[0], Constant)
        assert isinstance(self.output_node.inputs[0], Constant)
        assert tensor.shape == (1, 3, 5, 5)
        assert tensor.dtype == ml_dtypes.bfloat16
        assert np.all(self.input_node.outputs[0].values == tensor.values)
        assert np.all(self.output_node.inputs[0].values == tensor.values)

    def test_can_convert_in_place_to_variable(self):
        tensor = self.tensor.to_variable(dtype=np.float32, shape=(1, 3, 224, 224))
        assert tensor is self.tensor
        assert isinstance(tensor, Variable)
        assert isinstance(self.input_node.outputs[0], Variable)
        assert tensor.dtype == np.float32
        assert tensor.shape == (1, 3, 224, 224)
        assert self.input_node.outputs[0].dtype == tensor.dtype
        assert self.input_node.outputs[0].shape == tensor.shape

    def test_equals(self):
        assert self.tensor == self.tensor

    def test_set_inputs_updates_old_inputs(self):
        dummy = Node(op="dummy")
        self.tensor.inputs = [dummy]
        assert len(self.input_node.outputs) == 0
        assert dummy.outputs[0] == self.tensor

    def test_set_outputs_updates_old_outputs(self):
        dummy = Node(op="dummy")
        self.tensor.outputs = [dummy]
        assert len(self.output_node.inputs) == 0
        assert dummy.inputs[0] == self.tensor

    def test_can_copy_inputs_from_other_node(self):
        tensor = Variable(name="other_test_tensor")
        tensor.inputs = self.tensor.inputs
        assert tensor.inputs == self.tensor.inputs
        # Contents should be the same, but it should not just be a reference to the existing SynchronizedList
        assert tensor.inputs is not self.tensor.inputs

    def test_can_copy_outputs_from_other_node(self):
        tensor = Variable(name="other_test_tensor")
        tensor.outputs = self.tensor.outputs
        assert tensor.outputs == self.tensor.outputs
        assert tensor.outputs is not self.tensor.outputs

    # copy.copy/deepcopy should yield a regular list instead of a synchronized list
    @pytest.mark.parametrize("copy_func", [copy.copy, copy.deepcopy])
    def test_copy_makes_normal_list(self, copy_func):
        assert isinstance(self.tensor.inputs, SynchronizedList)

        inputs = copy_func(self.tensor.inputs)
        assert not isinstance(inputs, SynchronizedList)
        assert isinstance(inputs, list)

    def test_i(self):
        x = Variable(name="x")
        y = Variable(name="y")
        node = Node(op="Add", name="Input", inputs=[x], outputs=[y])
        assert y.i() == x

    def test_i_multiple_inputs(self):
        x = Variable(name="x")
        x2 = Variable(name="x2")
        y = Variable(name="y")
        node = Node(op="Add", name="Input", inputs=[x, x2], outputs=[y])
        assert y.i() == x
        assert y.i(1) == x2

    def test_o(self):
        x = Variable(name="x")
        y = Variable(name="y")
        node = Node(op="Add", name="Input", inputs=[x], outputs=[y])
        assert x.o() == y

    def test_o_multiple_outputs(self):
        x = Variable(name="x")
        y = Variable(name="y")
        y2 = Variable(name="y2")
        node = Node(op="Add", name="Input", inputs=[x], outputs=[y])
        node2 = Node(op="Add", name="Input", inputs=[x], outputs=[y2])
        assert x.o() == y
        assert x.o(1) == y2


class TestVariable(TensorBaseTests):
    def setup_method(self):
        self.tensor = Variable(
            name="test_tensor", dtype=np.float32, shape=(1, 3, 224, 224)
        )
        self.input_node = Node(op="Add", outputs=[self.tensor])
        self.output_node = Node(op="Add", inputs=[self.tensor])

    def test_equals_name_mismatch(self):
        tensor = Variable(name="test_tensor0", dtype=np.float32, shape=(1, 3, 224, 224))
        assert not self.tensor == tensor


class TestConstant(TensorBaseTests):
    def setup_method(self):
        self.tensor = Constant(
            name="test_tensor",
            values=np.ones((1, 3, 5, 5), dtype=np.float64),
        )
        self.input_node = Node(
            op="Add", outputs=[self.tensor]
        )  # Doesn't make sense for Constants, but needed to make base tests happy.
        self.output_node = Node(op="Add", inputs=[self.tensor])

    def test_can_get_shape(self):
        assert self.tensor.shape == (1, 3, 5, 5)

    def test_can_get_dtype(self):
        assert self.tensor.dtype == np.float64


@pytest.fixture
def node_with_nested_subgraphs():
    inner_subgraph = Graph(name="inner_subgraph")
    outer_subgraph_1 = Graph(name="subgraph1")
    outer_subgraph_2 = Graph(
        name="subgraph2", nodes=[Node("Add", attrs={"x": inner_subgraph})]
    )
    node = Node(op="Add", attrs={"x": outer_subgraph_1, "y": outer_subgraph_2, "z": 5})
    return node


class TestNode(object):
    def setup_method(self):
        self.input_tensor = Variable(name="x")
        self.output_tensor = Variable(name="y")
        self.node = Node(
            op="Add",
            name="Test",
            inputs=[self.input_tensor],
            outputs=[self.output_tensor],
        )

    def test_equals(self):
        assert self.node == self.node

    def test_equals_name_mismatch(self):
        node = Node(op="Add", name="OtherTest")
        assert not self.node == node

    def test_equals_op_mismatch(self):
        node = Node(op="Subtract", name="Test")
        assert not self.node == node

    def test_equals_num_inputs_mismatch(self):
        node = Node(op="Subtract", name="Test")
        assert not self.node == node

    def test_equals(self):
        assert self.node == self.node

    def test_equals_inputs_mismatch(self):
        tensor = Variable(name="other_tensor")
        assert not self.input_tensor == tensor

        node = Node(op="Add", name="Test", inputs=[tensor])
        assert not self.node == node

    def test_set_inputs_updates_old_inputs(self):
        dummy = Variable(name="dummy")
        self.node.inputs = [dummy]
        assert len(self.input_tensor.outputs) == 0
        assert dummy.outputs[0] == self.node

    def test_set_outputs_updates_old_outputs(self):
        dummy = Variable(name="dummy")
        self.node.outputs = [dummy]
        assert len(self.output_tensor.inputs) == 0
        assert dummy.inputs[0] == self.node

    def test_can_copy_inputs_from_other_node(self):
        node = Node(op="Subtract")
        node.inputs = self.node.inputs
        assert node.inputs == self.node.inputs
        # Contents should be the same, but it should not just be a reference to the existing SynchronizedList
        assert node.inputs is not self.node.inputs

    def test_can_copy_outputs_from_other_node(self):
        node = Node(op="Subtract")
        node.outputs = self.node.outputs
        assert node.outputs == self.node.outputs
        assert node.outputs is not self.node.outputs

    def test_i(self):
        intermediate_tensor = Variable(name="intermediate")
        input_node = Node(
            op="Add",
            name="Input",
            inputs=[self.input_tensor],
            outputs=[intermediate_tensor],
        )
        output_node = Node(
            op="Add",
            name="Out",
            inputs=[intermediate_tensor],
            outputs=[self.output_tensor],
        )
        assert output_node.i() == input_node

    def test_i_multiple_inputs(self):
        intermediate_tensor = Variable(name="intermediate")
        intermediate_tensor2 = Variable(name="intermediate2")
        input_node = Node(
            op="Add",
            name="Input",
            inputs=[self.input_tensor],
            outputs=[intermediate_tensor],
        )
        input_node2 = Node(
            op="Add",
            name="Input2",
            inputs=[self.input_tensor],
            outputs=[intermediate_tensor2],
        )
        output_node = Node(
            op="Add",
            name="Out",
            inputs=[intermediate_tensor, intermediate_tensor2],
            outputs=[self.output_tensor],
        )
        assert output_node.i() == input_node
        assert output_node.i(1) == input_node2

    def test_o(self):
        intermediate_tensor = Variable(name="intermediate")
        input_node = Node(
            op="Add",
            name="Input",
            inputs=[self.input_tensor],
            outputs=[intermediate_tensor],
        )
        output_node = Node(
            op="Add",
            name="Out",
            inputs=[intermediate_tensor],
            outputs=[self.output_tensor],
        )
        assert input_node.o() == output_node

    def test_o_multiple_outputs(self):
        intermediate_tensor = Variable(name="intermediate")
        intermediate_tensor2 = Variable(name="intermediate2")
        input_node = Node(
            op="Add",
            name="Input",
            inputs=[self.input_tensor],
            outputs=[intermediate_tensor],
        )
        output_node = Node(
            op="Add",
            name="Out",
            inputs=[intermediate_tensor],
            outputs=[self.output_tensor],
        )
        output_node2 = Node(
            op="Add",
            name="Input2",
            inputs=[intermediate_tensor],
            outputs=[intermediate_tensor2],
        )
        assert input_node.o() == output_node
        assert input_node.o(1) == output_node2

    def test_domain(self):
        node = Node(op="Add", domain="test")
        assert node.domain == "test"

    def test_subgraphs_not_recursive(self, node_with_nested_subgraphs):
        unrelated_graph = Graph(name="unrelated")
        subgraph_names = {
            subgraph.name for subgraph in node_with_nested_subgraphs.subgraphs()
        }
        assert subgraph_names == {"subgraph1", "subgraph2"}

    def test_subgraphs_recursive(self, node_with_nested_subgraphs):
        unrelated_graph = Graph(name="unrelated")
        subgraph_names = {
            subgraph.name
            for subgraph in node_with_nested_subgraphs.subgraphs(recursive=True)
        }
        assert subgraph_names == {"subgraph1", "subgraph2", "inner_subgraph"}


class TestNodeIO(object):
    def setup_method(self, field_names):
        self.tensors = [
            Variable(
                name="test_tensor_{:}".format(i),
                dtype=np.float32,
                shape=(1, 3, 224, 224),
            )
            for i in range(10)
        ]
        self.node = Node(op="Dummy")

    def get_lists(self, field_names):
        return getattr(self.node, field_names[0]), field_names[1]

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_append(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.append(self.tensors[0])
        assert nlist[0] == self.tensors[0]
        assert getattr(self.tensors[0], tensor_field)[0] == self.node

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_extend(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.extend(self.tensors)
        for tensor in self.tensors:
            assert tensor in nlist
            assert getattr(tensor, tensor_field)[0] == self.node

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_insert(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.append(self.tensors[1])
        nlist.insert(0, self.tensors[0])
        assert nlist[0] == self.tensors[0]
        assert getattr(self.tensors[0], tensor_field)[0] == self.node

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_remove(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.append(self.tensors[0])
        nlist.remove(self.tensors[0])
        assert len(nlist) == 0
        assert len(getattr(self.tensors[0], tensor_field)) == 0

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_pop(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.append(self.tensors[0])
        tensor = nlist.pop()
        assert len(nlist) == 0
        assert len(getattr(tensor, tensor_field)) == 0

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_pop_index(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.extend(self.tensors)
        tensor = nlist.pop(1)
        assert self.tensors[1] not in nlist
        assert len(getattr(tensor, tensor_field)) == 0

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_del_index(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.extend(self.tensors)
        tensor = nlist[1]
        del nlist[1]
        assert self.tensors[1] not in nlist
        assert len(getattr(tensor, tensor_field)) == 0

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_clear(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.extend(self.tensors)
        nlist.clear()
        assert len(nlist) == 0
        assert all([len(getattr(tensor, tensor_field)) == 0 for tensor in self.tensors])

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_add(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist = nlist + self.tensors
        for tensor in self.tensors:
            assert tensor in nlist

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_iadd(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist += self.tensors
        for tensor in self.tensors:
            assert tensor in nlist
            assert getattr(tensor, tensor_field)[0] == self.node

    @pytest.mark.parametrize(
        "field_names", [("inputs", "outputs"), ("outputs", "inputs")]
    )
    def test_setitem(self, field_names):
        nlist, tensor_field = self.get_lists(field_names)
        nlist.append(self.tensors[0])
        new_tensor = Variable("new_tensor")
        nlist[0] = new_tensor
        assert nlist[0] == new_tensor
        assert len(getattr(self.tensors[0], tensor_field)) == 0
        assert getattr(new_tensor, tensor_field)[0] == self.node

    def test_iadd_on_node_directly(self):
        t0 = Variable("t0")
        n0 = Node("", inputs=[])

        n0.inputs += [t0]
        assert len(n0.inputs) == 1
        assert n0.inputs[0] == t0


class TestTensorIO(object):
    def test_iadd_on_tensor_directly(self):
        n0 = Node("")
        t0 = Variable("t0")

        t0.inputs += [n0]
        assert len(t0.inputs) == 1
        assert t0.inputs[0] == n0


class TestLazyValues(object):
    def test_basic(self):
        shape = (1, 5, 5)
        onnx_tensor = onnx.helper.make_tensor_value_info(
            "test", onnx.TensorProto.FLOAT, shape
        )
        values = LazyValues(onnx_tensor)

        assert values.dtype == np.float32
        assert tuple(values.shape) == shape
        assert values.nbytes == 100