# This file contains the set of example scripts for the script node.
# User can click on the Code Snippets button in the UI to display these scripts.
# To add new example scripts to the script node,
# simply add the delimiter to the bottom of this file, followed by the new script.

# Declare og to suppress linter warnings about undefined variables
og = None


# # # DELIMITER # # #
Title = "Default Script"
# This script is executed the first time the script node computes, or the next time it computes after this script
# is modified or the 'Reset' button is pressed.

# The following callback functions may be defined in this script:
#     setup(db): Called immediately after this script is executed
#     compute(db): Called every time the node computes (should always be defined)
#     cleanup(db): Called when the node is deleted or the reset button is pressed (if setup(db) was called before)

# Defining setup(db) and cleanup(db) is optional, but if compute(db) is not defined then this script node will run
# in legacy mode, where the entire script is executed on every compute and the callback functions above are ignored.

# Available variables:
#    db: og.Database The node interface, attributes are db.inputs.data, db.outputs.data.
#                    Use db.log_error, db.log_warning to report problems.
#                    Note that this is available outside of the callbacks only to support legacy mode.
#    og: The OmniGraph module

# Import statements, function/class definitions, and global variables may be placed outside of the callbacks.
# Variables may also be added to the db.internal_state state object.

# Example code snippet:

import math

UNITS = "cm"


def calculate_circumfrence(radius):
    return 2 * math.pi * radius


def setup(db):
    state = db.internal_state
    state.radius = 1


def compute(db):
    state = db.internal_state
    circumfrence = calculate_circumfrence(state.radius)
    print(f"{circumfrence} {UNITS}")
    state.radius += 1


# To see more examples, click on the Code Snippets button below.


# # # DELIMITER # # #
Title = "Compute Count"
# In this example, we retrieve the number of times this script node has been computed
# and assign it to Output Data so that downstream nodes can use this information.
def compute(db):
    compute_count = db.node.get_compute_count()
    db.outputs.my_output_attribute = compute_count


# # # DELIMITER # # #
Title = "Fibonacci"
# In this example, we produce the Fibonacci sequence 0, 1, 1, 2, 3, 5, 8, 13, 21...
# Each time this node is evaluated, the next Fibonacci number will be set as the output value.
# This illustrates how variables declared in the setup script can be used to keep persistent information.
# Remember to add an output attribute of type 'int' named my_output_attribute first.


def setup(db):
    state = db.internal_state
    state.num1 = 0
    state.num2 = 1


def compute(db):
    state = db.internal_state
    total = state.num1 + state.num2
    state.num1 = state.num2
    state.num2 = total

    db.outputs.my_output_attribute = state.num1


# # # DELIMITER # # #
Title = "Controller"
# In this example, we use omni.graph.core.Controller to create cube prims.
# Each time this node is evaluated, it will create a new cube prim on the scene.
# When the 'Reset' button is pressed or the node is deleted, the created cube prims will be deleted.

import omni.kit.commands


def setup(db):
    state = db.internal_state
    state.cube_count = 0


def compute(db):
    state = db.internal_state
    state.cube_count += 1
    og.Controller.edit(
        db.node.get_graph(), {og.Controller.Keys.CREATE_PRIMS: [(f"/World/Cube{state.cube_count}", "Cube")]}
    )


def cleanup(db):
    state = db.internal_state
    omni.kit.commands.execute("DeletePrims", paths=[f"/World/Cube{i}" for i in range(1, state.cube_count + 1)])


# # # DELIMITER # # #
Title = "Random Vectors with Warp"
# In this example, we compute the lengths of random 3D vectors using Warp

import numpy as np
import warp as wp

wp.init()

NUM_POINTS = 1024
DEVICE = "cuda"


@wp.kernel
def length(points: wp.array(dtype=wp.vec3), lengths: wp.array(dtype=float)):
    # thread index
    tid = wp.tid()

    # compute distance of each point from origin
    lengths[tid] = wp.length(points[tid])


def compute(db):
    # allocate an array of 3d points
    points = wp.array(np.random.rand(NUM_POINTS, 3), dtype=wp.vec3, device=DEVICE)
    lengths = wp.zeros(NUM_POINTS, dtype=float, device=DEVICE)

    # launch kernel
    wp.launch(kernel=length, dim=len(points), inputs=[points, lengths], device=DEVICE)

    print(lengths)


# # # DELIMITER # # #
Title = "Value Changed Callbacks"
# In this example, we register a value changed callback function for inputs:my_input_attribute.
# The callback is called when the value of inputs:my_input_attribute is changed from the property panel.
# Remember to add an input attribute named my_input_attribute first.


def on_my_input_attribute_changed(attr):
    print(f"inputs:my_input_attribute = {attr.get_attribute_data().get()}")


def setup(db):
    attr = db.node.get_attribute("inputs:my_input_attribute")
    attr.register_value_changed_callback(on_my_input_attribute_changed)


def compute(db):
    pass
