"""Support for simplified access to data on nodes of type omni.graph.tutorials.CudaData

This is a tutorial node. It performs different functions on the GPU to illustrate different types of data access. The first
adds inputs 'a' and 'b' to yield output 'sum', all of which are on the GPU. The second is a sample expansion deformation
that multiplies every point on a set of input points, stored on the GPU, by a constant value, stored on the CPU, to yield
a set of output points, also on the GPU. The third is an assortment of different data types illustrating how different data
is passed to the GPU. This particular node uses CUDA for its GPU computations, as indicated in the memory type value. Normal
use case for GPU compute is large amounts of data. For testing purposes this node only handles a very small amount but the
principle is the same.
"""

import carb
import numpy

import omni.graph.core as og
import omni.graph.core._omni_graph_core as _og
import omni.graph.tools.ogn as ogn



class OgnTutorialCudaDataDatabase(og.Database):
    """Helper class providing simplified access to data on nodes of type omni.graph.tutorials.CudaData

    Class Members:
        node: Node being evaluated

    Attribute Value Properties:
        Inputs:
            inputs.a
            inputs.b
            inputs.color
            inputs.half
            inputs.matrix
            inputs.multiplier
            inputs.points
        Outputs:
            outputs.color
            outputs.half
            outputs.matrix
            outputs.points
            outputs.sum
    """

    # Imprint the generator and target ABI versions in the file for JIT generation
    GENERATOR_VERSION = (1, 41, 3)
    TARGET_VERSION = (2, 139, 12)

    # This is an internal object that provides per-class storage of a per-node data dictionary
    PER_NODE_DATA = {}

    # This is an internal object that describes unchanging attributes in a generic way
    # The values in this list are in no particular order, as a per-attribute tuple
    #     Name, Type, ExtendedTypeIndex, UiName, Description, Metadata,
    #     Is_Required, DefaultValue, Is_Deprecated, DeprecationMsg
    # You should not need to access any of this data directly, use the defined database interfaces
    INTERFACE = og.Database._get_interface([
        ('inputs:a', 'float', 0, None, 'First value to be added in algorithm 1', {ogn.MetadataKeys.DEFAULT: '0.0'}, True, 0.0, False, ''),
        ('inputs:b', 'float', 0, None, 'Second value to be added in algorithm 1', {ogn.MetadataKeys.DEFAULT: '0.0'}, True, 0.0, False, ''),
        ('inputs:color', 'color3d', 0, None, 'Input with three doubles as a color for algorithm 3', {ogn.MetadataKeys.DEFAULT: '[1.0, 0.5, 1.0]'}, True, [1.0, 0.5, 1.0], False, ''),
        ('inputs:half', 'half', 0, None, 'Input of type half for algorithm 3', {ogn.MetadataKeys.DEFAULT: '1.0'}, True, 1.0, False, ''),
        ('inputs:matrix', 'matrix4d', 0, None, 'Input with 16 doubles interpreted as a double-precision 4d matrix', {ogn.MetadataKeys.DEFAULT: '[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]'}, True, [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]], False, ''),
        ('inputs:multiplier', 'float3', 0, None, 'Amplitude of the expansion for the input points in algorithm 2', {ogn.MetadataKeys.MEMORY_TYPE: 'cpu', ogn.MetadataKeys.DEFAULT: '[1.0, 1.0, 1.0]'}, True, [1.0, 1.0, 1.0], False, ''),
        ('inputs:points', 'float3[]', 0, None, 'Points to be moved by algorithm 2', {ogn.MetadataKeys.DEFAULT: '[]'}, True, [], False, ''),
        ('outputs:color', 'color3d', 0, None, 'Output with three doubles as a color for algorithm 3', {}, True, None, False, ''),
        ('outputs:half', 'half', 0, None, 'Output of type half for algorithm 3', {}, True, None, False, ''),
        ('outputs:matrix', 'matrix4d', 0, None, 'Output with 16 doubles interpreted as a double-precision 4d matrix', {}, True, None, False, ''),
        ('outputs:points', 'float3[]', 0, None, 'Final positions of points from algorithm 2', {}, True, None, False, ''),
        ('outputs:sum', 'float', 0, None, 'Sum of the two inputs from algorithm 1', {}, True, None, False, ''),
    ])

    @classmethod
    def _populate_role_data(cls):
        """Populate a role structure with the non-default roles on this node type"""
        role_data = super()._populate_role_data()
        role_data.inputs.color = og.AttributeRole.COLOR
        role_data.inputs.matrix = og.AttributeRole.MATRIX
        role_data.outputs.color = og.AttributeRole.COLOR
        role_data.outputs.matrix = og.AttributeRole.MATRIX
        return role_data

    class ValuesForInputs(og.DynamicAttributeAccess):
        LOCAL_PROPERTY_NAMES = { }
        """Helper class that creates natural hierarchical access to input attributes"""
        def __init__(self, node: og.Node, attributes, dynamic_attributes: og.DynamicAttributeInterface):
            """Initialize simplified access for the attribute data"""
            context = node.get_graph().get_default_graph_context()
            super().__init__(context, node, attributes, dynamic_attributes)
            self._batchedReadAttributes = []
            self._batchedReadValues = []

        @property
        def a(self):
            data_view = og.AttributeValueHelper(self._attributes.a)
            return data_view.get(on_gpu=True)

        @a.setter
        def a(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.a)
            data_view = og.AttributeValueHelper(self._attributes.a)
            data_view.set(value, on_gpu=True)

        @property
        def b(self):
            data_view = og.AttributeValueHelper(self._attributes.b)
            return data_view.get(on_gpu=True)

        @b.setter
        def b(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.b)
            data_view = og.AttributeValueHelper(self._attributes.b)
            data_view.set(value, on_gpu=True)

        @property
        def color(self):
            data_view = og.AttributeValueHelper(self._attributes.color)
            return data_view.get(on_gpu=True)

        @color.setter
        def color(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.color)
            data_view = og.AttributeValueHelper(self._attributes.color)
            data_view.set(value, on_gpu=True)

        @property
        def half(self):
            data_view = og.AttributeValueHelper(self._attributes.half)
            return data_view.get(on_gpu=True)

        @half.setter
        def half(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.half)
            data_view = og.AttributeValueHelper(self._attributes.half)
            data_view.set(value, on_gpu=True)

        @property
        def matrix(self):
            data_view = og.AttributeValueHelper(self._attributes.matrix)
            return data_view.get(on_gpu=True)

        @matrix.setter
        def matrix(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.matrix)
            data_view = og.AttributeValueHelper(self._attributes.matrix)
            data_view.set(value, on_gpu=True)

        @property
        def multiplier(self):
            data_view = og.AttributeValueHelper(self._attributes.multiplier)
            return data_view.get()

        @multiplier.setter
        def multiplier(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.multiplier)
            data_view = og.AttributeValueHelper(self._attributes.multiplier)
            data_view.set(value)

        @property
        def points(self):
            data_view = og.AttributeValueHelper(self._attributes.points)
            return data_view.get(on_gpu=True)

        @points.setter
        def points(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.points)
            data_view = og.AttributeValueHelper(self._attributes.points)
            data_view.set(value, on_gpu=True)
            self.points_size = data_view.get_array_size()

        def _prefetch(self):
            readAttributes = self._batchedReadAttributes
            newValues = _og._prefetch_input_attributes_data(readAttributes)
            if len(readAttributes) == len(newValues):
                self._batchedReadValues = newValues

    class ValuesForOutputs(og.DynamicAttributeAccess):
        LOCAL_PROPERTY_NAMES = { }
        """Helper class that creates natural hierarchical access to output attributes"""
        def __init__(self, node: og.Node, attributes, dynamic_attributes: og.DynamicAttributeInterface):
            """Initialize simplified access for the attribute data"""
            context = node.get_graph().get_default_graph_context()
            super().__init__(context, node, attributes, dynamic_attributes)
            self.points_size = None
            self._batchedWriteValues = { }

        @property
        def color(self):
            data_view = og.AttributeValueHelper(self._attributes.color)
            return data_view.get(on_gpu=True)

        @color.setter
        def color(self, value):
            data_view = og.AttributeValueHelper(self._attributes.color)
            data_view.set(value, on_gpu=True)

        @property
        def half(self):
            data_view = og.AttributeValueHelper(self._attributes.half)
            return data_view.get(on_gpu=True)

        @half.setter
        def half(self, value):
            data_view = og.AttributeValueHelper(self._attributes.half)
            data_view.set(value, on_gpu=True)

        @property
        def matrix(self):
            data_view = og.AttributeValueHelper(self._attributes.matrix)
            return data_view.get(on_gpu=True)

        @matrix.setter
        def matrix(self, value):
            data_view = og.AttributeValueHelper(self._attributes.matrix)
            data_view.set(value, on_gpu=True)

        @property
        def points(self):
            data_view = og.AttributeValueHelper(self._attributes.points)
            return data_view.get(reserved_element_count=self.points_size, on_gpu=True)

        @points.setter
        def points(self, value):
            data_view = og.AttributeValueHelper(self._attributes.points)
            data_view.set(value, on_gpu=True)
            self.points_size = data_view.get_array_size()

        @property
        def sum(self):
            data_view = og.AttributeValueHelper(self._attributes.sum)
            return data_view.get(on_gpu=True)

        @sum.setter
        def sum(self, value):
            data_view = og.AttributeValueHelper(self._attributes.sum)
            data_view.set(value, on_gpu=True)

        def _commit(self):
            _og._commit_output_attributes_data(self._batchedWriteValues)
            self._batchedWriteValues = { }

    class ValuesForState(og.DynamicAttributeAccess):
        """Helper class that creates natural hierarchical access to state attributes"""
        def __init__(self, node: og.Node, attributes, dynamic_attributes: og.DynamicAttributeInterface):
            """Initialize simplified access for the attribute data"""
            context = node.get_graph().get_default_graph_context()
            super().__init__(context, node, attributes, dynamic_attributes)

    def __init__(self, node):
        super().__init__(node)
        dynamic_attributes = self.dynamic_attribute_data(node, og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT)
        self.inputs = OgnTutorialCudaDataDatabase.ValuesForInputs(node, self.attributes.inputs, dynamic_attributes)
        dynamic_attributes = self.dynamic_attribute_data(node, og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT)
        self.outputs = OgnTutorialCudaDataDatabase.ValuesForOutputs(node, self.attributes.outputs, dynamic_attributes)
        dynamic_attributes = self.dynamic_attribute_data(node, og.AttributePortType.ATTRIBUTE_PORT_TYPE_STATE)
        self.state = OgnTutorialCudaDataDatabase.ValuesForState(node, self.attributes.state, dynamic_attributes)
