"""
Implementation of the Python node accessing CUDA attributes in a way that accesses the GPU arrays with a CPU pointer.

No actual computation is done here as the tutorial nodes are not set up to handle GPU computation.
"""
import ctypes

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)


def get_address(attr: og.Attribute) -> int:
    """Returns the contents of the memory the attribute points to"""
    ptr_type = ctypes.POINTER(ctypes.c_size_t)
    ptr = ctypes.cast(attr.memory, ptr_type)
    return ptr.contents.value


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

    @staticmethod
    def compute(db) -> bool:
        """Accesses the CUDA data, which for arrays exists as wrappers around CPU memory pointers to GPU pointer arrays.
        No compute is done here.
        """

        # Put accessors into local variables for convenience
        input_points = db.inputs.points
        multiplier = db.inputs.multiplier

        # Set the size to what is required for the multiplication - this can be done without accessing GPU data.
        # Notice that since this is CPU pointers to GPU data the size has to be taken from the data type description
        # rather than the usual method of taking len(input_points).
        db.outputs.points_size = input_points.dtype.size

        # After changing the size the memory isn't allocated immediately (when necessary). It is delayed until you
        # request access to it, which is what this line will do.
        output_points = db.outputs.points

        # This is a separate test to add a points attribute to the output bundle to show how when a bundle has
        # CPU pointers to the GPU data that information propagates to its children
        # Start with an empty output bundle.
        output_bundle = db.outputs.outBundle
        output_bundle.clear()
        output_bundle.add_attributes([og.Type(og.BaseDataType.FLOAT, 3, 1)], ["points"])
        bundle_attr = output_bundle.attribute_by_name("points")
        # As for the main attributes, setting the bundle member size readies the buffer of the given size on the GPU
        bundle_attr.size = input_points.dtype.size

        # The output cannot be written to here through the normal assignment mechanisms, e.g. the typical step of
        # copying input points to the output points, as the data is not accessible on the GPU through Python directly.
        # Instead you can access the GPU memory pointers through the attribute values and send it to CUDA code, either
        # generated from the Python code or accessed through something like pybind wrappers.

        print("Locations in CUDA() should be in GPU memory space")
        print(f"    CPU Location for reference = {hex(id(db))}", flush=True)
        print(f"    Input points are {input_points} at CUDA({hex(get_address(input_points))})", flush=True)
        print(f"    Multiplier is CUDA({multiplier})", flush=True)
        print(f"    Output points are {output_points} at CUDA({hex(get_address(output_points))})", flush=True)
        print(f"    Bundle {bundle_attr.gpu_value} at CUDA({hex(get_address(bundle_attr.gpu_value))})", flush=True)

        return True
