"""
Implementation of a node handling complex attribute data
"""
# This class exercises access to the DataModel through the generated database class for a representative set of
# complex data types, including tuples, arrays, arrays of tuples, and role-based attributes. More details on
# individual type definitions can be found in the earlier C++ tutorial nodes where each of those types are
# explored in detail.

# Any Python node with array attributes will receive its data wrapped in a numpy array for efficiency.
# Unlike C++ includes, a Python import is not transitive so this has to be explicitly imported here.
import numpy
import omni.graph.core as og


class OgnTutorialComplexDataPy:
    """Exercise a sample of complex data types through a Python OmniGraph node"""

    @staticmethod
    def compute(db) -> bool:
        """
        Multiply a float array by a float[3] to yield a float[3] array, using the point3f role.
        Practically speaking the data in the role-based attributes is no different than the underlying raw data
        types. The role only helps you understand what the intention behind the data is, e.g. to differentiate
        surface normals and colours, both of which might have float[3] types.
        """

        # Verify that the output array was correctly set up to have a "point" role
        assert db.role.outputs.a_productArray == og.AttributeRole.POSITION

        multiplier = db.inputs.a_vectorMultiplier
        input_array = db.inputs.a_inputArray
        input_array_size = len(db.inputs.a_inputArray)

        # The output array should have the same number of elements as the input array.
        # Setting the size informs fabric that when it retrieves the data it should allocate this much space.
        db.outputs.a_productArray_size = input_array_size

        # The assertions illustrate the type of data that should have been received for inputs and set for outputs
        assert isinstance(multiplier, numpy.ndarray)  # numpy.ndarray is the underlying type of tuples
        assert multiplier.shape == (3,)
        assert isinstance(input_array, numpy.ndarray)  # numpy.ndarray is the underlying type of simple arrays
        assert input_array.shape == (input_array_size,)

        # If the input array is empty then the output is empty and does not need any computing
        if input_array.shape[0] == 0:
            db.outputs.a_productArray = []
            assert db.outputs.a_productArray.shape == (0, 3)
            return True

        # numpy has a nice little method for replicating the multiplier vector the number of times required
        # by the size of the input array.
        #   e.g. numpy.tile( [1, 2], (3, 1) ) yields [[1, 2], [1, 2], [1, 2]]
        product = numpy.tile(multiplier, (input_array_size, 1))

        # Multiply each of the tiled vectors by the corresponding constant in the input array
        for i in range(0, product.shape[0]):
            product[i] = product[i] * input_array[i]
        db.outputs.a_productArray = product

        # Make sure the correct type of array was produced
        assert db.outputs.a_productArray.shape == (input_array_size, 3)

        # Create the output token array by copying the input array with the elements changed to strings
        db.outputs.a_tokenArray = numpy.array([str(x) for x in input_array])

        # Note that at the time of this writing you are not able to assign token arrays element-by-element as you
        # might do for other types of arrays. So this coding method, usable for other types of arrays such as float[],
        # would issue a warning and then fail to set the values on the output:
        #    db.outputs.a_tokenArray_size = input_array_size
        #    for index, element in enumerate(input_array):
        #        db.outputs.a_tokenArray[index] = str(element)

        return True
