# flake8: noqa
# This file contains snippets of example Python code to be imported by the OGN documentation.
# It's not actual running code itself, merely sections of code that illustrate a point
# It uses the reStructuredText code import with begin/end markings to cherry-pick relevant portions.
# The sections of code being referenced are demarqued by "begin-XX" and "end-XX".

# begin-minimal
# This line isn't strictly necessary. It's only useful for more stringent type information of the compute parameter.
# Note how the extra submodule "ogn" is appended to the extension's module to find the database file.
from ogn.examples.ogn.OgnNoOpDatabase import OgnNoOpDatabase


class OgnNoOp:
    @staticmethod
    def compute(db: OgnNoOpDatabase) -> bool:
        """This comment should describe the compute algorithm.
        Running help() on the database class will provide the information from the node type description field in the
        .ogn file. The information here should be a supplement to that, consisting of implementation notes.
        """
        # This logs a warning to the console, once
        db.log_warning("This node does nothing")

        # The node is accessible through the database. Normally you don't have to check its validity as it
        # will already be checked, it's just done here to illustrate access and error logging.
        if db.node is None or not db.node.isValid():
            # This logs an error to the console and should only be used for compute failure
            db.log_error("The node being computed is not valid")
            return False

        return True


# end-minimal

# begin-node-metadata
class OgnNodeMetadata:
    @staticmethod
    def compute(db) -> bool:
        # Specifically defined metadata can be accessed by name
        print(f"The author of this node is {db.get_metadata('author')}")

        # Some metadata is automatically added; you can see it by iterating over all of the existing metadata.
        # The Python iteration interfaces with the C++ ABI to make it seem like the metadata is an iterable list.
        for metadata_name, metadata_value in db.node.get_all_metadata():
            print(f"Metadata for {metadata_name} is {metadata_value}")

        return True


# end-node-metadata

# begin-node-icon
import omni.graph.tools.ogn as ogn


class OgnNodeWithIcon:
    @staticmethod
    def compute(db) -> bool:
        # The icon path is just a special case of metadata. The hardcoded key is in the Python namespace
        path = db.get_metadata(ogn.MetadataKeys.ICON_PATH)
        color = db.get_metadata(ogn.MetadataKeys.ICON_COLOR)
        background_color = db.get_metadata(ogn.MetadataKeys.ICON_BACKGROUND_COLOR)
        border_color = db.get_metadata(ogn.MetadataKeys.ICON_BORDER_COLOR)
        if path is not None:
            print(f"Icon found at {path}")
            print(f"...color override is {color}" if color is not None else "...using default color")
            print(
                f"...backgroundColor override is {background_color}"
                if background_color is not None
                else "...using default backgroundColor"
            )
            print(
                f"...borderColor override is {border_color}"
                if border_color is not None
                else "...using default borderColor"
            )

        return True


# end-node-icon

# begin-node-scheduling
class OgnNodeSchedulingHints:
    @staticmethod
    def compute(db) -> bool:
        scheduling_hints = db.abi_node.get_node_type().get_scheduling_hints()

        # Ordinarily you would not need to access this scheduling hints as it is mainly for OmniGraph's use,
        # however it is available through the ABI so you can access it at runtime if you wish.
        print(f"Is this node threadsafe? {scheduling_hints.get_thread_safety()}")
        return True


# end-node-scheduling

# begin-node-singleton
import omni.graph.tools.ogn as ogn


class OgnNodeSingleton:
    @staticmethod
    def compute(db) -> bool:
        # The singleton value is just a special case of metadata. The hardcoded key is in the Python namespace
        singleton_value = db.get_metadata(ogn.MetadataKeys.SINGLETON)
        if singleton_value and singleton_value[0] == "1":
            print("I am a singleton")

        return True


# end-node-singleton

# begin-node-categories
import omni.graph.tools.ogn as ogn


class OgnNodeCategories:
    @staticmethod
    def compute(db) -> bool:
        # The categories value is just a special case of metadata. The hardcoded key is in the Python namespace
        categories_value = db.get_metadata(ogn.MetadataKeys.CATEGORIES)
        if categories_value:
            print(f"These are my categories {categories_value.split(',')}")

        return True


# end-node-categories

# begin-node-tags
import omni.graph.tools.ogn as ogn


class OgnNodeTags:
    @staticmethod
    def compute(db) -> bool:
        # The tags value is just a special case of metadata. The hardcoded key is in the Python namespace
        print(f"My tags are {db.get_metadata(ogn.MetadataKeys.TAGS)}")

        return True


# end-node-tags

# begin-tokens
class OgnNodeTokens:
    @staticmethod
    def compute(db) -> bool:

        print(f"The name for red is {db.tokens.red}")
        print(f"The name for green is {db.tokens.green}")
        print(f"The name for blue is {db.tokens.blue}")

        return True


# end-tokens

# begin-node-uiName
import omni.graph.core as og


class OgnNodeUiName:
    @staticmethod
    def compute(db) -> bool:
        # The uiName value is just a special case of metadata. The hardcoded key is in the Python namespace
        print("Call me ", db.get_metadata(ogn.MetadataKeys.UI_NAME))

        return True


# end-node-uiName

# begin-simple
class OgnTokenStringLength:
    @staticmethod
    def compute(db) -> bool:

        # Access pattern is "db", the database, "inputs", the attribute's access type, and "token", the name the
        # attribute was given in the .ogn file
        #
        # Local variables can be used to clarify the intent, but are not necessary. As a matter of consistency we
        # use PEP8 conventions for local variables. Attribute names may not exactly follow the naming conventions
        # since they are mutually exclusive between C++ and Python (camelCase vs. snake_case)
        token_to_measure = db.inputs.token

        # Simple assignment to the output attribute's accessor is all you need to do to set the value, as it is
        # pointing directly to that data.
        db.outputs.length = len(token_to_measure)

        return True


# end-simple

# begin-tuple
class OgnVectorMultiply:
    @staticmethod
    def compute(db) -> bool:

        # The fact that everything is in numpy makes this kind of calculation trivial
        db.outputs.product = db.inputs.vector1.reshape(4, 1) @ db.inputs.vector2.reshape(1, 4)

        # Here db.inputs.vector1.shape = db.inputs.vector1.shape = (4,), db.outputs.product.shape = (4,4)
        return True


# end-tuple

# begin-role
class OgnPointsToVector:
    @staticmethod
    def compute(db) -> bool:

        # In Python the wrapper is only a property so the role has to be extracted from a parallel structure
        if db.role.inputs.point1 != db.ROLE_POINT:
            db.log_error(f"Cannot convert role {db.role.inputs.point1} to {db.ROLE_POINT}")
            return False
        if db.role.inputs.point2 != db.ROLE_POINT:
            db.log_error(f"Cannot convert role {db.role.inputs.point2} to {db.ROLE_POINT}")
            return False
        if db.role.outputs.vector != db.ROLE_POINT:
            db.log_error(f"Cannot convert role {db.role.inputs.vector} to {db.ROLE_VECTOR}")
            return False

        # The actual calculation is a trivial numpy call
        db.outputs.vector = db.inputs.point2 - db.inputs.point1

        return True


# end-role

# begin-array
import numpy as np


class OgnPartialSums:
    @staticmethod
    def compute(db) -> bool:

        # This is a critical step, setting the size of the output array. Without this the array has no memory in
        # which to write.
        #
        # As the Python wrapper is a property, in order to get and set the size a secondary property is introduced
        # for array data types which have the same name as the attribute with "_size" appended to it. For outputs
        # this property also has a setter, which accomplishes the resizing.
        #
        db.outputs.partialSums_size = db.inputs.array_size

        # Always explicitly handle edge cases as it ensures your node doesn't disrupt evaluation
        if db.outputs.partialSums_size == 0:
            return True

        # IMPORTANT:
        #     The data value returned from accessing the property is a numpy array whose memory location was
        #     allocated by Fabric. As such you cannot use the numpy functions that resize the arrays as they will
        #     not use Fabric data.

        # However, since the array attribute data is wrapped in a numpy array you can use numpy functions that
        # modify data in place to make efficient use of memory.
        #
        db.inputs.array.cumsum(out=db.outputs.partialSums)

        # A second way you can assign array data is to collect the data externally and then do a simple list
        # assignment. This is less efficient as it does a physical copy of the entire list, though more flexible as
        # you can arbitrarily resize your data before assigning it. If you use this approach you skip the step of
        # setting the output array size as the assignment will do it for you.
        #
        #   # Using numpy
        #   output_list = np.cumsum(db.inputs.array)
        #   db.outputs.partialSums = output_list
        #
        #   # Using Python lists
        #   output_list = [value for value in db.inputs.array]
        #   for index, value in enumerate(output_list[:-1]):
        #       output_list[index + 1] = output_list[index + 1] + output_list[index]
        #   db.outputs.partialSum = output_list
        #
        #   # numpy is smart enough to do element-wise copy, but in this case you do have to preset the size
        #   output_list = np.cumsum(db.inputs.array)
        #   db.outputs.partialSums_size = db.inputs.array_size
        #   db.outputs.partialSums[:] = output_list[:]
        #

        return True


# end-array

# begin-tuple-array
class OgnCrossProducts:
    @staticmethod
    def compute(db) -> bool:

        # It usually keeps your code cleaner if you put your attribute wrappers into local variables, avoiding
        # the constant use of the "db.inputs" or "db.outputs" namespaces.
        a = db.inputs.a
        b = db.inputs.b
        crossProduct = db.outputs.crossProduct

        # This node chooses to make mismatched array lengths an error. You could also make it a warning, or just
        # simply calculate the result for the minimum number of available values.
        if db.inputs.a_size != db.inputs.b_size:
            db.log_error(f"Input array lengths do not match - '{db.inputs.a_size}' vs. '{db.inputs.b_size}'")
            return False

        # As with simple arrays, the size of the output tuple-array must be set first to allocate Fabric memory.
        db.outputs.crossProduct_size = db.inputs.a_size

        # Edge case is easily handled
        if db.inputs.a_size == 0:
            return False

        # The numpy cross product returns the result so there will be a single copy of the result onto the output.
        # numpy handles the iteration over the array so this one line does the entire calculation.
        crossProduct = np.cross(a, b)
        # This common syntax will do exactly the same thing
        #    crossProduct[:] = np.cross(a, b)[:]

        return True


# end-tuple-array

# begin-string
class OgnReverseString:
    @staticmethod
    def compute(db) -> bool:

        # In Python the wrapper to string attributes provides a standard Python string object.
        # As the wrapper is a property the assignment of a value uses the setter method to both allocate the
        # necessary space in Fabric and copy the values.

        db.outputs.result = db.inputs.original[::-1]

        return True


# end-string

import numpy as np

# begin-any
import omni.graph.core as og


class OgnAdd:
    @staticmethod
    def compute(db) -> bool:

        # The extra work in handling extended types such as "any" or "union" is in checking the resolved types to
        # first ensure they are handled and second select appropriate code paths if they differ from one type to
        # another.

        # Unlike C++ the loose typing of the Python objects mean that there is no requirement for explicit casting
        # of data into the resolve types. Instead, the much more flexible method of letting the numpy library do the
        # work for us is employed. The add operation will perform any necessary type conversion, as will the assignment
        # operator.
        try:
            db.outputs.sum.value = np.add(db.inputs.a.value, db.inputs.b.value)
        except Exception as error:
            # If numpy doesn't handle addition of the two input types, or assignment to the output type, then an
            # exception will be thrown and the node can report the error.
            db.log_error(f"Addition could not be performed: {error}")
            return False

        return True


# end-any

# begin-union
class OgnMultiplyNumbers:
    @staticmethod
    def compute(db) -> bool:

        # Full details on handling extended types can be seen in the example for the "any" type. This example
        # shows only the necessary parts to handle the two types accepted for this union type (float and double).
        # The underlying code is all the same, the main difference is in the fact that the graph only allows
        # resolving to types explicitly mentioned in the union, rather than any type at all.

        # Use the exception system to implicitly check the resolved types. Unresolved types will not have accessible
        # data and raise an exception.
        try:
            db.outputs.product = np.mult(db.inputs.a, db.inputs.b)
        except Exception as error:
            db.log_error(f"Multiplication could not be performed: {error}")
            return False

        return True


# end-union

# begin-bundle
class OgnMergeBundles:
    @staticmethod
    def compute(db) -> bool:

        bundleA = db.inputs.bundleA
        bundleB = db.inputs.bundleB
        mergedBundle = db.outputs.bundle

        # Bundle assignment means "assign all of the members of the RHS bundle to the LHS bundle". It doesn't
        # do a deep copy of the bundle members.
        mergedBundle = bundleA

        # Bundle insertion adds the contents of a bundle to an existing bundle. The bundles may not have members
        # with the same names
        mergedBundle.insert_bundle(bundleB)

        return True


# end-bundle

# begin-bundle-data
import omni.graph.core as og

FLOAT_TYPE = og.Type(og.BaseDataType.FLOAT)


class OgnCalculateBrightness:
    def brightness_from_rgb(self, r: float, g: float, b: float) -> float:
        """The actual algorithm to run using a well-defined conversion"""
        return (r * (299.0) + (g * 587.0) + (b * 114.0)) / 256.0

    @staticmethod
    def compute(db) -> bool:
        # Retrieve the bundle accessor
        color = db.inputs.color

        # Using the bundle accessor, try to retrieve the RGB color members. In this case the types have to be
        # float, though in a more general purpose node you might also allow for double, half, and int types.
        r = color.attribute_by_name(db.tokens.r)
        g = color.attribute_by_name(db.tokens.g)
        b = color.attribute_by_name(db.tokens.b)
        # Validity of a member is a boolean
        if r.type == FLOAT_TYPE and g.type == FLOAT_TYPE and b.type == FLOAT_TYPE:
            db.outputs.brightness.value = OgnCalculateBrightness.brightness_from_rgb(r.value, g.value, b.value)
            return True

        # Having failed to extract RGB members, do the same check for CMYK members
        c = color.attribute_by_name(db.tokens.c)
        m = color.attribute_by_name(db.tokens.m)
        y = color.attribute_by_name(db.tokens.y)
        k = color.attribute_by_name(db.tokens.k)
        if c.type == FLOAT_TYPE and m.type == FLOAT_TYPE and y.type == FLOAT_TYPE and k.type == FLOAT_TYPE:
            db.outputs.brightness.value = OgnCalculateBrightness.brightness_from_rgb(
                (1.0 - c / 100.0) * (1.0 - k / 100.0),
                (1.0 - m / 100.0) * (1.0 - k / 100.0),
                (1.0 - y / 100.0) * (1.0 - k / 100.0),
            )
            return True

        # You could be more verbose about the reason for the problem as there are a few different scenarios:
        #    - some but not all of r,g,b or c,m,y,k were in the bundle
        #    - none of the color components were in the bundle
        #    - some or all of the color components were found but were of the wrong data type
        db.logError("Neither the groups (r, g, b) nor (c, m, y, k) are in the color bundle. Cannot compute brightness")
        return False


# end-bundle-data

# begin-memory-type
class OgnMemoryType:
    @staticmethod
    def compute(db) -> bool:

        # The operation specifies moving the points data onto the GPU for further computation if the size of
        # the input data reaches a threshold where that will make the computation more efficient.
        # (This particular node just moves data; in practice you would perform an expensive calculation on it.)
        if db.inputs.points.size > db.inputs.sizeThreshold:
            # The gpu property forces the data onto the GPU. It may or may not perform CPU->GPU copies under the
            # covers. Fabric handles all of those details so that you don't have to.
            db.outputs.points.gpu = db.inputs.points.gpu
        else:
            # The cpu property forces the data onto the CPU. It may or may not perform GPU->CPU copies under the
            # covers. Fabric handles all of those details so that you don't have to.
            db.outputs.points.cpu = db.inputs.points.cpu
        return True


# end-memory-type

# begin-cuda-pointers
class OgnCudaPointers:
    @staticmethod
    def compute(db) -> bool:

        # When the *cudaPointers* keyword is set to *cpu* this wrapped array will contain a CPU pointer that
        # references the GPU array data. If not, it would have contained a GPU pointer that references the GPU
        # array data and not been able to be dereferenced on the CPU side.
        callCudaFunction(db.inputs.cudaPoints, db.outputs.cudaPoints)
        return True


# end-cuda-pointers

# begin-dynamic
class OgnDynamicDuo:
    @staticmethod
    def compute(db) -> bool:
        try:
            # Testing for the existence of the dynamic input boolean attribute "Robin"
            db.outputs.batman = "Duo" if db.inputs.robin else "Unknown"
        except AttributeError:
            db.outputs.batman = "Solo"
        return True


# end-dynamic

# begin-optional
import random
from contextlib import suppress


class OgnShoes:
    SHOE_TYPES = ["Runners", "Slippers", "Oxfords"]

    @staticmethod
    def compute(db) -> bool:

        shoe_index = random.randint(0, 2)
        shoe_type_name = OgnShoes.SHOE_TYPES[shoe_index]

        # If the shoe is a type that has laces then append the lace type name
        if shoe_index != 1:
            # As this is an optional value it may or may not be valid at this point.
            # This check happens automatically with required attributes. With optional ones it has to be done when used.
            if db.attributes.inputs.shoeLaceStyle.isValid():
                # The attribute may be valid but the data retrieval may still fail. In Python this is flagged in one of
                # two ways - raising an exception, or returning None. Both indicate the possibility of invalid data.
                # In this node we've chosen to silently ignore expected but invalid shoelace style values. We could
                # equally have logged an error or a warning.
                with suppress(ValueError, TypeError):
                    shoelace_style = db.inputs.shoelaceStyle
                    if shoelace_style is not None:
                        shoe_type_name += f" with {shoelace_style} laces"

        db.outputs.shoeType = shoe_type_name
        return True


# end-optional

# begin-attribute-uiName
import omni.graph.tools.ogn as ogn


class OgnAttributeUiName:
    @staticmethod
    def compute(db) -> bool:
        # The uiName value is just a special case of metadata
        print(f"Call me {db.get_metadata(ogn.MetadataKeys.UI_NAME, db.attributes.inputs.x)}")

        return True


# end-attribute-uiName

# begin-attribute-uiType
import omni.graph.tools.ogn as ogn


class OgnAttributeUiType:
    @staticmethod
    def compute(db) -> bool:
        # The uiType value is just a special case of metadata
        print(f"The property panel ui type is {db.get_metadata(ogn.MetadataKeys.UI_TYPE, '(default)')}")

        return True


# end-attribute-uiType

# begin-unvalidated
import omni.graph.core as og


class OgnABTest:
    @staticmethod
    def compute(db) -> bool:
        choice = db.outputs.choice
        out_type = choice.type
        # Check to see which input is selected and verify that its data type matches the output resolved type
        if db.inputs.selectA:
            input_a = db.inputs.a
            if not input_a.is_valid() or input_a.type != out_type:
                db.log_error(
                    f"Mismatched types at input a - '{input_a.type.get_ogn_type_name()}' versus '{out_type.get_ogn_type_name()}'"
                )
                return False
            choice.value = input_a.value
        else:
            input_b = db.inputs.b
            if not input_b.is_valid() or input_b.type != out_type:
                db.log_error(
                    f"Mismatched types at input b - '{input_b.type.get_ogn_type_name()}' versus '{out_type.get_ogn_type_name()}'"
                )
                return False
            choice.value = input_b.value
        return True


# end-unvalidated

# begin-attribute-metadata
import omni.graph.tools.ogn as ogn


class OgnStarWarsCharacters:
    @staticmethod
    def compute(db) -> bool:
        anakin_attr = db.attributes.inputs.anakin
        # Specifically defined metadata can be accessed by name
        print(f"Anakin's secret is {db.get_metadata('secret', anakin_attr)}")

        # Some metadata is automatically added; you can see it by iterating over all of the existing metadata.
        for metadata_name, metadata_value in anakin_attr.get_all_metadata():
            print(f"Metadata for {metadata_name} is {metadata_value}")

        # You can also access it directly from the database's metadata interface, either from the node type...
        print(f"Node UI Name is {db.get_metadata(ogn.MetadataKeys.UI_NAME)}")
        # ...or from a specific attribute
        print(f"Attribute UI Name is {db.get_metadata(ogn.MetadataKeys.UI_NAME, anakin_attr)}")

        return True


# end-attribute-metadata

# begin-state-node
class OgnStateNode:
    class State:
        """Container object holding the node's state information"""

        def __init__(self):
            self.counter = 0

    @staticmethod
    def internal_state():
        """Returns an object that will contain per-node state information"""
        return OgnStateNode.State()

    @staticmethod
    def compute(db) -> bool:
        print(f"This node has been evaluated {db.internal_state.counter} times")
        db.internal_state.counter += 1

        return True


# end-state-node

# begin-versioned-node
import carb
import omni.graph.core as og


class OgnMultiply:
    @staticmethod
    def compute(db) -> bool:
        db.outputs.result = db.inputs.a * db.inputs.b + db.inputs.offset

        return True

    @staticmethod
    def update_node_version(context, node, old_version, new_version):
        if old_version == 1 and new_version == 3:
            node.create_attribute("inputs:offset", og.Type(og.BaseDataType.FLOAT))
            return True

        # Always good practice to flag unknown version changes so that they are not forgotten
        carb.log_error(f"Do not know how to upgrade Multiply from version {old_version} to {new_version}")
        return False


# end-versioned-node
