"""Test the math nodes"""
import math

import numpy as np
import omni.graph.core as og
import omni.graph.core.tests as ogts


# ======================================================================
class TestMathNodes(ogts.OmniGraphTestCase):
    """Run a simple unit test that exercises graph functionality"""

    TEST_GRAPH_PATH = "/World/TestGraph"

    # ----------------------------------------------------------------------
    async def test_add_node_simple(self):
        """Test the add node for a variety of simple types"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (simple_node_a, simple_node_b, simple_node_sum, add_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SimpleA", "omni.graph.tutorials.SimpleData"),
                    ("SimpleB", "omni.graph.tutorials.SimpleData"),
                    ("SimpleSum", "omni.graph.tutorials.SimpleData"),
                    ("Add", "omni.graph.nodes.Add"),
                ]
            },
        )
        input_a = add_node.get_attribute("inputs:a")
        input_b = add_node.get_attribute("inputs:b")
        output_sum = add_node.get_attribute("outputs:sum")

        # Test data is a list of lists, one per test configuration, containing (ATTRIBUTE, A, B, SUM):
        #    ATTRIBUTE: the name of the attribute to be connected from the "Simple" nodes to resolve the data types
        #    A: First value to add
        #    B: Second value to add
        #    SUM: Expected sum
        # (Note that since the simple nodes add 1 to inputs to get outputs SUM != A + B)
        test_data = [
            ["a_float", 5.25, 6.25, 13.5],
            ["a_double", 5.125, 6.125, 13.25],
            ["a_half", 5.0, 6.0, 13.0],
            ["a_int", 6, 7, 15],
            ["a_int64", 8, 9, 19],
            ["unsigned:a_uchar", 10, 11, 23],
            ["unsigned:a_uint", 12, 13, 27],
            ["unsigned:a_uint64", 14, 15, 31],
        ]

        for (attribute_name, a, b, expected) in test_data:
            connection_list = [
                ((f"outputs:{attribute_name}", simple_node_a), input_a),
                ((f"outputs:{attribute_name}", simple_node_b), input_b),
                (output_sum, (f"inputs:{attribute_name}", simple_node_sum)),
            ]
            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CONNECT: connection_list,
                    keys.SET_VALUES: [
                        ((f"inputs:{attribute_name}", simple_node_a), a),
                        ((f"inputs:{attribute_name}", simple_node_b), b),
                    ],
                },
            )
            await controller.evaluate(graph)
            actual = og.Controller.get(output_sum)
            self.assertAlmostEqual(expected, actual)

            # Disconnecting the connections we just made lets the loop work
            controller.edit(self.TEST_GRAPH_PATH, {keys.DISCONNECT: connection_list})

    # ----------------------------------------------------------------------
    async def test_add_node_mismatched(self):
        """Test the add node for simple types that have compatible but not identical types"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (_, _, _, add_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("SimpleA", "omni.graph.tutorials.SimpleData"),
                    ("SimpleB", "omni.graph.tutorials.SimpleData"),
                    ("SimpleSum", "omni.graph.tutorials.SimpleData"),
                    ("Add", "omni.graph.nodes.Add"),
                ],
                keys.CONNECT: [
                    ("SimpleA.outputs:a_double", "Add.inputs:a"),
                    ("SimpleB.outputs:a_int", "Add.inputs:b"),
                    ("Add.outputs:sum", "SimpleSum.inputs:a_float"),
                ],
                keys.SET_VALUES: [
                    ("SimpleA.inputs:a_double", 5.5),
                    ("SimpleB.inputs:a_int", 11),
                ],
            },
        )
        await controller.evaluate(graph)
        actual = og.Controller.get(controller.attribute("outputs:sum", add_node))
        self.assertAlmostEqual(18.5, actual)

    # ----------------------------------------------------------------------
    async def test_add_node_tuples(self):
        """Test the add node for a variety of tuple types"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (tuple_node_a, tuple_node_b, tuple_node_sum, add_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("TupleA", "omni.tutorials.TupleData"),
                    ("TupleB", "omni.tutorials.TupleData"),
                    ("TupleSum", "omni.tutorials.TupleData"),
                    ("Add", "omni.graph.nodes.Add"),
                ]
            },
        )
        input_a = add_node.get_attribute("inputs:a")
        input_b = add_node.get_attribute("inputs:b")
        output_sum = add_node.get_attribute("outputs:sum")

        # Test data is a list of lists, one per test configuration, containing (ATTRIBUTE, A, B, SUM):
        #    ATTRIBUTE: the name of the attribute to be connected from the "Tuple" nodes to resolve the data types
        #    A: First value to add
        #    B: Second value to add
        #    SUM: Expected sum
        # (Note that since the tuple nodes add 1 to input elements to get outputs SUM != A + B)
        test_data = [
            ["a_double2", (20.0, 30.0), (20.1, 30.1), (42.1, 62.1)],
            ["a_float2", (2.0, 3.0), (2.1, 3.1), (6.1, 8.1)],
            ["a_half2", (1.0, 2.0), (3.0, 4.0), (6.0, 8.0)],
            ["a_int2", (16, -5), (22, 5), (40, 2)],
            ["a_float3", (1.1, 2.2, 3.3), (2.2, 3.3, 4.4), (5.3, 7.5, 9.7)],
            ["a_double3", (10.1, 20.2, 30.3), (20.2, 30.3, 40.4), (32.3, 52.5, 72.7)],
        ]

        for (attribute_name, a, b, expected) in test_data:
            connection_list = [
                ((f"outputs:{attribute_name}", tuple_node_a), input_a),
                ((f"outputs:{attribute_name}", tuple_node_b), input_b),
                (output_sum, (f"inputs:{attribute_name}", tuple_node_sum)),
            ]
            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CONNECT: connection_list,
                    keys.SET_VALUES: [
                        ((f"inputs:{attribute_name}", tuple_node_a), a),
                        ((f"inputs:{attribute_name}", tuple_node_b), b),
                    ],
                },
            )
            await controller.evaluate(graph)
            actual = og.Controller.get(output_sum)
            self.assertTrue(np.allclose(expected, actual), f"{expected} != {actual}")

            # Disconnecting the connections we just made lets the loop work
            controller.edit(self.TEST_GRAPH_PATH, {keys.DISCONNECT: connection_list})

    # ----------------------------------------------------------------------
    async def test_add_node_arrays(self):
        """Test the add node for a variety of array types"""
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (array_node_a, array_node_b, array_node_sum, add_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("ArrayA", "omni.graph.tutorials.ArrayData"),
                    ("ArrayB", "omni.graph.tutorials.ArrayData"),
                    ("ArraySum", "omni.graph.tutorials.ArrayData"),
                    ("Add", "omni.graph.nodes.Add"),
                ],
                keys.SET_VALUES: [
                    # Set the array nodes up to double the values
                    ("ArrayA.inputs:multiplier", 2.0),
                    ("ArrayB.inputs:multiplier", 2.0),
                ],
            },
        )
        input_a = add_node.get_attribute("inputs:a")
        input_b = add_node.get_attribute("inputs:b")
        output_sum = add_node.get_attribute("outputs:sum")

        # Test data is a list of lists, one per test configuration, containing (A, B, SUM):
        #    A: First value to add
        #    B: Second value to add
        #    SUM: Expected sum
        # (Note that since the array nodes add 1 to input elements to get outputs SUM != A + B)
        test_data = [
            [[20.0, 30.0], [20.1, 30.1], [80.2, 120.2]],
            [[1.0, 2.0, 3.0, 4.0], [-1.0, -2.0, -3.0, -4.0], [0.0, 0.0, 0.0, 0.0]],
        ]

        for (a, b, expected) in test_data:
            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CONNECT: [
                        (("outputs:result", array_node_a), input_a),
                        (("outputs:result", array_node_b), input_b),
                        (output_sum, ("inputs:original", array_node_sum)),
                    ],
                    keys.SET_VALUES: [
                        (("inputs:gates", array_node_a), [True] * len(a)),
                        (("inputs:gates", array_node_b), [True] * len(b)),
                        (("inputs:original", array_node_a), a),
                        (("inputs:original", array_node_b), b),
                    ],
                },
            )
            await controller.evaluate(graph)
            actual = og.Controller.get(output_sum)
            self.assertTrue(np.allclose(expected, actual), f"{expected} != {actual}")

    # ----------------------------------------------------------------------
    async def test_add_node_tuple_arrays(self):
        """Test the add node for array-of-tuple types"""
        controller = og.Controller()
        keys = og.Controller.Keys

        (graph, (tuple_array_node_a, tuple_array_node_b, tuple_array_node_sum, add_node), _, _,) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("TupleArrayA", "omni.graph.test.TupleArrays"),
                    ("TupleArrayB", "omni.graph.test.TupleArrays"),
                    ("TupleArraySum", "omni.graph.test.TupleArrays"),
                    ("Add", "omni.graph.nodes.Add"),
                ],
                keys.SET_VALUES: [
                    # Set the array nodes up to double the values
                    ("TupleArrayA.inputs:multiplier", 2.0),
                    ("TupleArrayB.inputs:multiplier", 2.0),
                ],
            },
        )
        input_a = add_node.get_attribute("inputs:a")
        input_b = add_node.get_attribute("inputs:b")
        output_sum = add_node.get_attribute("outputs:sum")

        # Test data is a list of lists, one per test configuration, containing (A, B, SUM):
        #    A: First value to add
        #    B: Second value to add
        #    SUM: Expected sum
        # (Note that since the array nodes add 1 to input elements to get outputs SUM != A + B)
        test_data = [
            [
                [[1.0, 2.0, 3.0], [1.1, 2.2, 3.3]],
                [[10.0, 20.0, 30.0], [11.0, 22.0, 33.0]],
                [[22.0, 44.0, 66.0], [24.2, 48.4, 72.6]],
            ],
        ]

        for (a, b, expected) in test_data:
            controller.edit(
                self.TEST_GRAPH_PATH,
                {
                    keys.CONNECT: [
                        (("outputs:float3Array", tuple_array_node_a), input_a),
                        (("outputs:float3Array", tuple_array_node_b), input_b),
                        (output_sum, ("inputs:float3Array", tuple_array_node_sum)),
                    ],
                    keys.SET_VALUES: [
                        (("inputs:float3Array", tuple_array_node_a), a),
                        (("inputs:float3Array", tuple_array_node_b), b),
                    ],
                },
            )
            await controller.evaluate(graph)
            actual = og.Controller.get(output_sum)
            self.assertTrue(np.allclose(expected, actual), f"{expected} != {actual}")

    # ----------------------------------------------------------------------
    async def test_add_all_types(self):
        """Test the OgnAdd node with all types"""
        node_name = "omni.graph.nodes.Add"
        # Input attributes to be resolved
        input_attribute_names = ["inputs:a", "inputs:b"]
        # Output attributes to be tested
        output_attribute_names = ["outputs:sum"]
        # list of unsupported types
        unsupported_types = ["string", "token", "path", "bool", "uint64", "bundle", "target"]
        # Operation giving the expected value
        operation = np.add

        # ----------------------------------------------------------
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (test_node, data_node), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {keys.CREATE_NODES: [("TestNode", node_name), ("TestData", "omni.graph.test.TestAllDataTypes")]},
        )

        # Generate Test data consisting of the input attribute followed by the expected output values
        test_data = []
        for attribute in data_node.get_attributes():
            type_base_name = attribute.get_resolved_type().get_base_type_name()
            if (
                "output" not in attribute.get_name()
                or type_base_name in unsupported_types
                or attribute.get_type_name() in unsupported_types
                or attribute.get_resolved_type().role == og.AttributeRole.EXECUTION
            ):
                continue
            value = og.Controller.get(attribute)
            test_data.append([attribute, operation(value, value)])

        input_attributes = [test_node.get_attribute(input_attribute) for input_attribute in input_attribute_names]
        output_attributes = [test_node.get_attribute(output_attribute) for output_attribute in output_attribute_names]

        for data_attribute, expected_value in test_data:
            connections = [(data_attribute, input_attribute) for input_attribute in input_attributes]
            controller.edit(self.TEST_GRAPH_PATH, {keys.CONNECT: connections})
            await controller.evaluate(graph)
            # Check outputs
            for output_attribute in output_attributes:
                actual_value = og.Controller.get(output_attribute)
                self.assertTrue(np.allclose(expected_value, actual_value), f"{expected_value} != {actual_value}")
            controller.edit(self.TEST_GRAPH_PATH, {keys.DISCONNECT: connections})

    # ----------------------------------------------------------------------
    async def test_multiply_all_types(self):
        """Test the OgnMultiply node with all types"""
        node_name = "omni.graph.nodes.Multiply"
        # Input attributes to be resolved
        input_attribute_names = ["inputs:a", "inputs:b"]
        # Output attributes to be tested
        output_attribute_names = ["outputs:product"]
        # list of unsupported types
        unsupported_types = ["string", "token", "path", "bool", "uint64", "bundle", "target"]
        # Operation giving the expected value
        operation = np.multiply

        # ----------------------------------------------------------
        controller = og.Controller()
        keys = og.Controller.Keys
        (graph, (test_node, data_node), _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {keys.CREATE_NODES: [("TestNode", node_name), ("DataNode", "omni.graph.test.TestAllDataTypes")]},
        )

        # Generate Test data consisting of the input attribute followed by the expected output values
        test_data = []
        for attribute in data_node.get_attributes():
            type_base_name = attribute.get_resolved_type().get_base_type_name()
            if (
                "output" not in attribute.get_name()
                or type_base_name in unsupported_types
                or attribute.get_type_name() in unsupported_types
                or attribute.get_resolved_type().role == og.AttributeRole.EXECUTION
            ):
                continue
            value = og.Controller.get(attribute)
            test_data.append([attribute, operation(value, value)])

        input_attributes = [test_node.get_attribute(input_attribute) for input_attribute in input_attribute_names]
        output_attributes = [test_node.get_attribute(output_attribute) for output_attribute in output_attribute_names]

        for data_attribute, expected_value in test_data:
            connections = [(data_attribute, input_attribute) for input_attribute in input_attributes]
            controller.edit(self.TEST_GRAPH_PATH, {keys.CONNECT: connections})
            await controller.evaluate(graph)
            # Check outputs
            for output_attribute in output_attributes:
                actual_value = og.Controller.get(output_attribute)
                self.assertTrue(np.allclose(expected_value, actual_value), f"{expected_value} != {actual_value}")
            controller.edit(self.TEST_GRAPH_PATH, {keys.DISCONNECT: connections})

    # ----------------------------------------------------------------------
    async def test_constant_pi_node(self):
        """Test the constant pi node"""
        controller = og.Controller()
        keys = og.Controller.Keys()
        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH, {keys.CREATE_NODES: [("constPi", "omni.graph.nodes.ConstantPi")]}
        )

        result = nodes[0].get_attribute("outputs:value")

        setter = og.Controller(attribute=og.Controller.attribute("inputs:factor", nodes[0]))
        getter = og.Controller(attribute=result)

        # Test data is a list, one per test configuration, containing (A, SUM):
        #    A: Multiply this by Pi to get the result
        #    SUM: Expected value
        test_data = [
            [1.0, math.pi],
            [2.0, 2.0 * math.pi],
            [0.5, 0.5 * math.pi],
        ]

        for (multiplier, expected) in test_data:
            setter.set(multiplier)
            await og.Controller.evaluate(graph)
            self.assertAlmostEqual(expected, getter.get())

    # ---------------------------------------------------------------------
    async def test_nth_root_node(self):
        """Test the nthroot node"""
        controller = og.Controller()
        keys = og.Controller.Keys()
        (graph, nodes, _, _) = controller.edit(
            self.TEST_GRAPH_PATH,
            {
                keys.CREATE_NODES: [
                    ("nthRoot", "omni.graph.nodes.NthRoot"),
                    ("src", "omni.graph.nodes.ConstantDouble"),
                ],
                keys.CONNECT: [("src.inputs:value", "nthRoot.inputs:value")],
            },
        )

        result = nodes[0].get_attribute("outputs:result")
        value_setter = og.Controller(attribute=og.Controller.attribute("inputs:value", nodes[1]))
        nth_root_setter = og.Controller(attribute=og.Controller.attribute("inputs:nthRoot", nodes[0]))
        getter = og.Controller(attribute=result)
        # Test data is a list, one per test configuration, containing (nthroot, value):
        #    nthroot: Take the nthroot of the value
        #    value: The number that is to be taken roots
        test_data = [[2, -4], [4, -8.0], [5, -1]]

        for (nthroot, value) in test_data:
            nth_root_setter.set(nthroot)
            value_setter.set(value)
            await og.Controller.evaluate(graph)
            self.assertTrue(np.isnan(getter.get()))

        test_data = [[3, -8.0, -2], [3, -1, -1]]

        for (nthroot, value, expected) in test_data:
            nth_root_setter.set(nthroot)
            value_setter.set(value)
            await og.Controller.evaluate(graph)
            self.assertEqual(expected, getter.get())

    # ----------------------------------------------------------------------
    async def test_floor_node_backwards_compatibility(self):
        """Test the floor node extended outputs for backwards compatibility"""
        # Test data contains floor nodes with int output, instead of token
        (result, error) = await ogts.load_test_file("TestFloorOutput.usda", use_caller_subdirectory=True)
        self.assertTrue(result, f"{error}")
        controller = og.Controller()
        graph = controller.graph("/World/TestGraph")

        def assert_attr_is_int(node):
            attr = controller.attribute("outputs:result", node, graph)
            self.assertEqual(attr.get_resolved_type().base_type, og.BaseDataType.INT)

        assert_attr_is_int("floor_double")
        assert_attr_is_int("floor_float")
        assert_attr_is_int("floor_half")
