"""
Implementation of the Python node accessing extended attributes whose memory location is determined at runtime.
"""
import omni.graph.core as og

# Only one type of data is handled by the compute - pointf[3][]
POINT_ARRAY_TYPE = og.Type(og.BaseDataType.FLOAT, tuple_count=3, array_depth=1, role=og.AttributeRole.POSITION)


class OgnTutorialCpuGpuExtendedPy:
    """Exercise GPU access for extended attributes through a Python OmniGraph node"""

    @staticmethod
    def compute(db) -> bool:
        """Implements the same algorithm as the C++ node OgnTutorialCpuGpuExtended.cpp.

        It follows the same code pattern for easier comparison, though in practice you would probably code Python
        nodes differently from C++ nodes to take advantage of the strengths of each language.
        """
        # Find and verify the attributes containing the points
        if db.attributes.inputs.cpuData.get_resolved_type() != POINT_ARRAY_TYPE:
            db.log_warning("Skipping compute - CPU attribute type did not resolve to pointf[3][]")
            return False
        if db.attributes.inputs.gpuData.get_resolved_type() != POINT_ARRAY_TYPE:
            db.log_warning("Skipping compute - GPU attribute type did not resolve to pointf[3][]")
            return False
        if db.attributes.outputs.cpuGpuSum.get_resolved_type() != POINT_ARRAY_TYPE:
            db.log_warning("Skipping compute - Sum attribute type did not resolve to pointf[3][]")
            return False

        # Put accessors into local variables for convenience
        gpu_data = db.inputs.gpuData
        cpu_data = db.inputs.cpuData
        sums = db.outputs.cpuGpuSum

        # Mismatched sizes cannot be computed
        if gpu_data.size != cpu_data.size:
            db.log_warning(f"Skipping compute - Point arrays are different sizes ({gpu_data.size} and {cpu_data.size})")

        # Set the size to what is required for the dot product calculation
        sums.size = cpu_data.size

        # Use the correct data access based on whether the output is supposed to be on the GPU or not
        if db.inputs.gpu:
            # The second line is how the values would be extracted if Python supported GPU data extraction.
            # When it does this tutorial will be updated
            sums.cpu_value = cpu_data.value + gpu_data.cpu_value
            # sums.gpu_value = cpu_data.gpu_value + gpu_data.value
        else:
            sums.cpu_value = cpu_data.value + gpu_data.cpu_value
        return True

    @staticmethod
    def on_connection_type_resolve(node: og.Node) -> None:
        """Whenever any of the inputs or the output get a resolved type the others should get the same resolution"""
        attribs = [
            node.get_attribute("inputs:cpuData"),
            node.get_attribute("inputs:gpuData"),
            node.get_attribute("outputs:cpuGpuSum"),
        ]
        og.resolve_fully_coupled(attribs)
