"""
Implementation of a Python node that handles a simple version upgrade where the different versions contain
different attributes:
    Version 0: multiplier
    Version 1: multiplier, wavelength
    Version 2: wavelength
Depending on the version that was received the upgrade will remove the obsolete "multiplier" attribute and/or
insert the new "wavelength" attribute.
"""
import numpy as np
import omni.graph.core as og


class OgnVersionedDeformerPy:
    @staticmethod
    def compute(db):
        input_points = db.inputs.points
        if len(input_points) <= 0:
            return False

        wavelength = db.inputs.wavelength
        if wavelength <= 0.0:
            wavelength = 310.0

        db.outputs.points_size = input_points.shape[0]
        # Nothing to evaluate if there are no input points
        if db.outputs.points_size == 0:
            return True
        output_points = db.outputs.points

        pt = input_points.copy()  # we still need to do a copy here because we do not want to modify points
        tx = (pt[:, 0] - wavelength) / wavelength
        ty = (pt[:, 1] - wavelength) / wavelength

        disp = 10.0 * (np.sin(tx * 10.0) + np.cos(ty * 15.0))
        pt[:, 2] += disp * 5

        output_points[:] = pt[:]  # we modify output_points in memory so we do not need to set the value again
        return True

    @staticmethod
    def update_node_version(context, node, old_version, new_version):
        if old_version < new_version:
            if old_version < 1:
                node.create_attribute(
                    "inputs:wavelength",
                    og.Type(og.BaseDataType.FLOAT),
                    og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
                    50.0,
                )
            if old_version < 2:
                node.remove_attribute("inputs:multiplier")
            return True
        return False

    @staticmethod
    def initialize_type(node_type):
        node_type.add_input("test_input", "int", True)
        node_type.add_output("test_output", "float", True)
        node_type.add_extended_input(
            "test_union_input", "float,double", True, og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION
        )
        return True
