"""
Implementation of the Python node accessing attributes whose type is determined at runtime.
This class exercises access to the DataModel through the generated database class for all simple data types.
"""
import omni.graph.core as og

# Hardcode each of the expected types for easy comparison
FLOAT_TYPE = og.Type(og.BaseDataType.FLOAT)
TOKEN_TYPE = og.Type(og.BaseDataType.TOKEN)
BOOL_ARRAY_TYPE = og.Type(og.BaseDataType.BOOL, array_depth=1)
FLOAT_ARRAY_TYPE = og.Type(og.BaseDataType.FLOAT, array_depth=1)
FLOAT3_TYPE = og.Type(og.BaseDataType.FLOAT, tuple_count=3)
INT2_TYPE = og.Type(og.BaseDataType.INT, tuple_count=2)
FLOAT3_ARRAY_TYPE = og.Type(og.BaseDataType.FLOAT, tuple_count=3, array_depth=1)


class OgnTutorialExtendedTypesPy:
    """Exercise the runtime data types through a Python OmniGraph node"""

    @staticmethod
    def compute(db) -> bool:
        """Implements the same algorithm as the C++ node OgnTutorialExtendedTypes.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.
        """

        def __compare_resolved_types(input_attribute, output_attribute) -> og.Type:
            """Returns the resolved type if they are the same, outputs a warning and returns None otherwise"""
            resolved_input_type = input_attribute.type
            resolved_output_type = output_attribute.type
            if resolved_input_type != resolved_output_type:
                db.log_warn(f"Resolved types do not match {resolved_input_type} -> {resolved_output_type}")
                return None
            return resolved_input_type if resolved_input_type.base_type != og.BaseDataType.UNKNOWN else None

        # ---------------------------------------------------------------------------------------------------
        def _compute_simple_values():
            """Perform the first algorithm on the simple input data types"""

            # Unlike C++ code the Python types are flexible so you must check the data types to do the right thing.
            # This works out better when the operation is the same as you don't even have to check the data type. In
            # this case the "doubling" operation is slightly different for floats and tokens.
            resolved_type = __compare_resolved_types(db.inputs.floatOrToken, db.outputs.doubledResult)
            if resolved_type == FLOAT_TYPE:
                db.outputs.doubledResult.value = db.inputs.floatOrToken.value * 2.0
            elif resolved_type == TOKEN_TYPE:
                db.outputs.doubledResult.value = db.inputs.floatOrToken.value + db.inputs.floatOrToken.value

            # A Pythonic way to do the same thing by just applying an operation and checking for compatibility is:
            #    try:
            #        db.outputs.doubledResult = db.inputs.floatOrToken * 2.0
            #    except TypeError:
            #        # Gets in here for token types since multiplying string by float is not legal
            #        db.outputs.doubledResult = db.inputs.floatOrToken + db.inputs.floatOrToken

            return True

        # ---------------------------------------------------------------------------------------------------
        def _compute_array_values():
            """Perform the second algorithm on the array input data types"""

            resolved_type = __compare_resolved_types(db.inputs.toNegate, db.outputs.negatedResult)
            if resolved_type == BOOL_ARRAY_TYPE:
                db.outputs.negatedResult.value = [not value for value in db.inputs.toNegate.value]
            elif resolved_type == FLOAT_ARRAY_TYPE:
                db.outputs.negatedResult.value = [-value for value in db.inputs.toNegate.value]

            return True

        # ---------------------------------------------------------------------------------------------------
        def _compute_tuple_values():
            """Perform the third algorithm on the 'any' data types"""

            resolved_type = __compare_resolved_types(db.inputs.tuple, db.outputs.tuple)
            # Notice how, since the operation is applied the same for both recognized types, the
            # same code can handle both of them.
            if resolved_type in (FLOAT3_TYPE, INT2_TYPE):
                db.outputs.tuple.value = tuple(-x for x in db.inputs.tuple.value)
            # An unresolved type is a temporary state and okay, resolving to unsupported types means the graph is in
            # an unsupported configuration that needs to be corrected.
            elif resolved_type is not None:
                type_name = resolved_type.get_type_name()
                db.log_error(f"Only float[3] and int[2] types are supported by this node, not {type_name}")
                return False

            return True

        # ---------------------------------------------------------------------------------------------------
        def _compute_flexible_values():
            """Perform the fourth algorithm on the multi-shape data types"""

            resolved_type = __compare_resolved_types(db.inputs.flexible, db.outputs.flexible)
            if resolved_type == FLOAT3_ARRAY_TYPE:
                db.outputs.flexible.value = [(-x, -y, -z) for (x, y, z) in db.inputs.flexible.value]
            elif resolved_type == TOKEN_TYPE:
                db.outputs.flexible.value = db.inputs.flexible.value[::-1]

            return True

        # ---------------------------------------------------------------------------------------------------
        compute_success = _compute_simple_values()
        compute_success = _compute_array_values() and compute_success
        compute_success = _compute_tuple_values() and compute_success
        compute_success = _compute_flexible_values() and compute_success

        # ---------------------------------------------------------------------------------------------------
        # As Python has a much more flexible typing system it can do things in a few lines that require a lot
        # more in C++. One such example is the ability to add two arbitrary data types. Here is an example of
        # how, using "any" type inputs "a", and "b", with an "any" type output "result" you can generically
        # add two elements without explicitly checking the type, failing only when Python cannot support
        # the operation.
        #
        #    try:
        #        db.outputs.result = db.inputs.a + db.inputs.b
        #        return True
        #    except TypeError:
        #        a_type = inputs.a.type().get_type_name()
        #        b_type = inputs.b.type().get_type_name()
        #        db.log_error(f"Cannot add attributes of type {a_type} and {b_type}")
        #    return False

        return True

    @staticmethod
    def on_connection_type_resolve(node) -> None:
        # There are 4 sets of type-coupled attributes in this node, meaning that the base_type of the attributes
        # must be the same for the node to function as designed.
        # 1. floatOrToken <-> doubledResult
        # 2. toNegate <-> negatedResult
        # 3. tuple <-> tuple
        # 4. flexible <-> flexible
        #
        # The following code uses a helper function to resolve the attribute types of the coupled pairs.  Note that
        # without this logic a chain of extended-attribute connections may result in a non-functional graph, due to
        # the requirement that types be resolved before graph evaluation, and the ambiguity of the graph without knowing
        # how the types are related.
        og.resolve_fully_coupled(
            [node.get_attribute("inputs:floatOrToken"), node.get_attribute("outputs:doubledResult")]
        )
        og.resolve_fully_coupled([node.get_attribute("inputs:toNegate"), node.get_attribute("outputs:negatedResult")])
        og.resolve_fully_coupled([node.get_attribute("inputs:tuple"), node.get_attribute("outputs:tuple")])
        og.resolve_fully_coupled([node.get_attribute("inputs:flexible"), node.get_attribute("outputs:flexible")])
