"""Basic tests of the compute graph"""

from unittest import skip

import numpy
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.graph.tools as ogt
import omni.kit.stage_templates
import omni.kit.test
import omni.usd
from omni.graph.core.tests.omnigraph_test_utils import create_scope_node
from pxr import Sdf

# ======================================================================
# Set of attribute types that can have default values set
ATTR_INT = Sdf.ValueTypeNames.Int
ATTR_BOOL = Sdf.ValueTypeNames.Bool
ATTR_FLOAT = Sdf.ValueTypeNames.Float
ATTR_DOUBLE = Sdf.ValueTypeNames.Double
ATTR_INT_ARRAY = Sdf.ValueTypeNames.IntArray
ATTR_FLOAT_ARRAY = Sdf.ValueTypeNames.FloatArray
ATTR_DOUBLE_ARRAY = Sdf.ValueTypeNames.DoubleArray
ATTR_FLOAT2_ARRAY = Sdf.ValueTypeNames.Float2Array
ATTR_FLOAT3_ARRAY = Sdf.ValueTypeNames.Float3Array
ATTR_STRING = Sdf.ValueTypeNames.String
ATTR_TOKEN = Sdf.ValueTypeNames.Token
ATTR_DEFAULTS = {
    ATTR_INT: 0,
    ATTR_BOOL: False,
    ATTR_FLOAT: 0.0,
    ATTR_DOUBLE: 0.0,
    ATTR_INT_ARRAY: [],
    ATTR_FLOAT_ARRAY: [],
    ATTR_DOUBLE_ARRAY: [],
    ATTR_FLOAT2_ARRAY: [],
    ATTR_FLOAT3_ARRAY: [],
    ATTR_STRING: "",
    ATTR_TOKEN: "",
}


# ======================================================================
def initialize_attribute(prim, attribute_name, attribute_type):
    """Set up an attribute named "attribute_name" on the given prim with a default value of the specified type"""
    ogt.dbg(f"    Initialize {attribute_name} with value {ATTR_DEFAULTS[attribute_type]}")
    prim.CreateAttribute(attribute_name, attribute_type).Set(ATTR_DEFAULTS[attribute_type])


# ======================================================================
def float3_array_of(array_order: int):
    """
    Returns an array of float3 the size of the array_order, with values equal to it
        e.g. array_order(1) = [[1.0, 1.0, 1.0]]
        e.g. array_order(2) = [[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]
    """
    return [[float(array_order)] * 3] * array_order


# ======================================================================
def generate_results(array_order1: int, array_order2: int):
    """
    Returns expected test data for the concatenate node consisting of:
        - an array of float3, in a sequence of "array_order1" elements with values
            (array_order1, array_order1, array_order1) followed by a similar sequence for "array_order2"
            e.g. values 1 and 2 would yield:
            [[1.0,1.0,1.0],[2.0,2.0,2.0],[2.0,2.0,2.0]]
    """
    return float3_array_of(array_order1) + float3_array_of(array_order2)


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

    # ----------------------------------------------------------------------
    def compare_values(self, expected_value, actual_value, error_message):
        """Generic assert comparison which uses introspection to choose the correct method to compare the data values"""
        ogt.dbg(f"----Comparing {actual_value} with expected {expected_value} as '{error_message}'")
        if isinstance(expected_value, (list, tuple)):
            # If the numpy module is not available array values cannot be tested so silently succeed
            if numpy is None:
                return
            # Values are returned as numpy arrays so convert the expected value and use numpy to do the comparison
            numpy.testing.assert_almost_equal(
                actual_value, numpy.array(expected_value), decimal=3, err_msg=error_message
            )
        elif isinstance(expected_value, float):
            self.assertAlmostEqual(actual_value, expected_value, error_message)
        else:
            self.assertEqual(actual_value, expected_value, error_message)

    # ----------------------------------------------------------------------
    @skip("TODO: Enable when gather nodes are fully supported")
    async def test_concat_node(self):
        """
        Run evaluation tests on the builtin math concat node. This node has its own test because it cannot be
        tested in isolation, it requires a gather node to provide a meaningful input.
        """
        concat_node = "ConcatFloat3Arrays"
        gather_node = "GatherFloat3Arrays"
        scope_node1 = "ScopeWithArrays1"
        scope_node2 = "ScopeWithArrays2"
        gather_output = f"/World/{gather_node}.createdOutput"
        cached_prims = {}
        scene_data = {
            # The concatFloat3Arrays node flattens and array of arrays of float3
            "omni.graph.nodes.ConcatenateFloat3Arrays": {
                "name": [concat_node],
                "inputs": [["inputs:inputArrays", Sdf.ValueTypeNames.Float3Array]],
                "outputs": [
                    ["outputs:outputArray", Sdf.ValueTypeNames.Float3Array, [[-1.0, -1.0, -1.0]]],
                    ["outputs:arraySizes", Sdf.ValueTypeNames.IntArray, [-1]],
                ],
            },
            # The gather node concatenates lists of inputs together so that a node can process them on mass
            "Gather": {
                "name": [gather_node],
                "inputs": [],
                "outputs": [["createdOutput", Sdf.ValueTypeNames.Float3Array, [[-1.0, -1.0, -1.0]]]],
            },
            # Generic scope nodes to supply the float3 arrays to the gather node
            "Scope": {
                "name": [scope_node1, scope_node2],
                "inputs": [],
                "outputs": [["createdOutput", Sdf.ValueTypeNames.Float3Array, [[-1.0, -1.0, -1.0]]]],
            },
        }
        # Create a stage consisting of one of each type of node to be tested with their attributes
        for node_type, node_configuration in scene_data.items():
            for node_name in node_configuration["name"]:
                if node_type == "Scope":
                    prim = create_scope_node(node_name)
                else:
                    stage = omni.usd.get_context().get_stage()
                    path = omni.usd.get_stage_next_free_path(stage, "/" + node_name, True)
                    _ = og.GraphController.create_node(path, node_type)
                    prim = stage.GetPrimAtPath(path)
                cached_prims[node_name] = prim
                for attribute_name, attribute_type in node_configuration["inputs"]:
                    prim.CreateAttribute(attribute_name, attribute_type)
                for attribute_name, attribute_type, _ in node_configuration["outputs"]:
                    prim.CreateAttribute(attribute_name, attribute_type)

        # Create the connections that get gathered arrays to the concat node
        cached_prims[concat_node].GetAttribute("inputs:inputArrays").AddConnection(gather_output)
        cached_prims[scope_node1].GetAttribute("createdOutput").Set([[1.0, 1.0, 1.0]])
        cached_prims[scope_node2].GetAttribute("createdOutput").Set([[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]])

        # Pairs of inputs for the create nodes to generate results (extracted from generated_results())
        test_data = [[1, 1], [1, 2], [1, 3], [3, 1], [3, 3]]

        await omni.kit.app.get_app().next_update_async()

        graph = omni.graph.core.get_current_graph()

        # Get the attributes that will be used for the computations
        concat_compute = graph.get_node("/World/" + concat_node)
        output_attr_array = concat_compute.get_attribute("outputs:outputArray")
        output_attr_counts = concat_compute.get_attribute("outputs:arraySizes")

        for test_values in test_data:

            # Set the test inputs
            cached_prims[scope_node1].GetAttribute("createdOutput").Set(float3_array_of(test_values[0]))
            cached_prims[scope_node2].GetAttribute("createdOutput").Set(float3_array_of(test_values[1]))

            # Run the compute
            await omni.kit.app.get_app().next_update_async()

            # Check the contents of the flattened array
            expected_output = generate_results(test_values[0], test_values[1])
            output_array_value = og.Controller.get(output_attr_array)
            self.compare_values(
                expected_output, output_array_value, "ConcatNode test - outputArray attribute value error"
            )

            # Check the contents of the array size counts
            output_counts_value = og.Controller.get(output_attr_counts)
            self.compare_values(test_values, output_counts_value, "ConcatNode test - arraySizes attribute value error")
