"""Test node exercising retrieval of dynamic attribute array values from various memory locations.
Add an input attribute named "array" of type "int[]", another named "simple" of type "double", and output attributes
with matching names and types to pass data through using the dynamic attribute get() and set() methods.
"""
import numpy as np
import omni.graph.core as og
from omni.graph.core._impl.dtypes import Int


class OgnTestDynamicAttributeMemory:
    @staticmethod
    def compute(db) -> bool:
        gpu_ptr_kind = og.PtrToPtrKind.NA
        on_gpu = db.inputs.onGpu
        if on_gpu:
            gpu_ptr_kind = og.PtrToPtrKind.GPU if db.inputs.gpuPtrsOnGpu else og.PtrToPtrKind.CPU

        db.set_dynamic_attribute_memory_location(on_gpu=on_gpu, gpu_ptr_kind=gpu_ptr_kind)

        try:
            array_value = db.inputs.array
            simple_value = db.inputs.simple
        except AttributeError:
            # Only verify when the input dynamic attribute was present
            db.outputs.inputMemoryVerified = False
            db.outputs.outputMemoryVerified = False
            return True

        # CPU inputs are returned as np.ndarrays, GPU inputs are og.DataWrapper
        if not on_gpu:
            db.outputs.inputMemoryVerified = isinstance(array_value, np.ndarray) and isinstance(simple_value, float)
        else:
            _, *new_shape = array_value.shape
            new_shape = tuple(new_shape)
            db.outputs.inputMemoryVerified = (
                array_value.gpu_ptr_kind == gpu_ptr_kind
                and array_value.device.cuda
                and array_value.is_array()
                and array_value.dtype == Int()
                and isinstance(simple_value, float)
            )

        try:
            db.outputs.array = array_value
            db.outputs.simple = simple_value
        except AttributeError:
            # Only verify when the output dynamic attribute was present
            db.outputs.outputMemoryVerified = False
            return True

        # CPU inputs are returned as np.ndarrays, GPU inputs are og.DataWrapper
        new_output_arr = db.outputs.array
        new_output = db.outputs.simple
        if not on_gpu:
            db.outputs.outputMemoryVerified = isinstance(new_output_arr, np.ndarray) and isinstance(new_output, float)
        else:
            _, *new_shape = new_output_arr.shape
            new_shape = tuple(new_shape)
            db.outputs.outputMemoryVerified = (
                new_output_arr.gpu_ptr_kind == gpu_ptr_kind
                and new_output_arr.device.cuda
                and new_output_arr.is_array()
                and new_output_arr.dtype == Int()
                and isinstance(new_output, float)
            )

        return True
