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

# Types to check on bundled attributes
FLOAT_ARRAY_TYPE = og.Type(og.BaseDataType.FLOAT, array_depth=1)
FLOAT3_ARRAY_TYPE = og.Type(og.BaseDataType.FLOAT, tuple_count=3, array_depth=1, role=og.AttributeRole.POSITION)


class OgnTutorialCpuGpuBundlesPy:
    """Exercise bundle members through a Python OmniGraph node"""

    @staticmethod
    def compute(db) -> bool:
        """Implements the same algorithm as the C++ node OgnTutorialCpuGpuBundles.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.
        """
        if db.inputs.gpu:
            # Invalid data yields no compute
            if not db.inputs.gpuBundle.valid:
                return True
            db.outputs.cpuGpuBundle = db.inputs.gpuBundle
        else:
            if not db.inputs.cpuBundle.valid:
                return True
            db.outputs.cpuGpuBundle = db.inputs.cpuBundle

        # Find and verify the attributes containing the points
        cpu_points = db.inputs.cpuBundle.attribute_by_name(db.tokens.points)
        if cpu_points.type != FLOAT3_ARRAY_TYPE:
            db.log_warning(
                f"Skipping compute - No valid float[3][] attribute named '{db.tokens.points}' on the CPU bundle"
            )
            return False
        gpu_points = db.inputs.gpuBundle.attribute_by_name(db.tokens.points)
        if gpu_points.type != FLOAT3_ARRAY_TYPE:
            db.log_warning(
                f"Skipping compute - No valid float[3][] attribute named '{db.tokens.points}' on the GPU bundle"
            )
            return False

        # If the attribute is not already on the output bundle then add it
        dot_product = db.outputs.cpuGpuBundle.attribute_by_name(db.tokens.dotProducts)
        if dot_product is None:
            dot_product = db.outputs.cpuGpuBundle.insert((og.Type(og.BaseDataType.FLOAT, array_depth=1), "dotProducts"))
        elif dot_product.type != FLOAT_ARRAY_TYPE:
            # Python types do not use a cast to find out if they are the correct type so explicitly check it instead
            db.log_warning(
                f"Skipping compute - No valid float[] attribute named '{db.tokens.dotProducts}' on the output bundle"
            )
            return False

        # Set the size to what is required for the dot product calculation
        dot_product.size = cpu_points.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
            dot_product.cpu_value = np.einsum("ij,ij->i", cpu_points.value, gpu_points.cpu_value)
            # dot_product.gpu_value = np.einsum("ij,ij->i", cpu_points.gpu_value, gpu_points.value)
        else:
            dot_product.cpu_value = np.einsum("ij,ij->i", cpu_points.value, gpu_points.cpu_value)

        return True
