import numpy as np
import omni.graph.core as og

# Mappings of possible numpy dtypes from the result data type and back
dtype_from_basetype = {
    og.BaseDataType.INT: np.int32,
    og.BaseDataType.INT64: np.int64,
    og.BaseDataType.HALF: np.float16,
    og.BaseDataType.FLOAT: np.float32,
    og.BaseDataType.DOUBLE: np.float64,
}

supported_basetypes = [
    og.BaseDataType.INT,
    og.BaseDataType.INT64,
    og.BaseDataType.HALF,
    og.BaseDataType.FLOAT,
    og.BaseDataType.DOUBLE,
]

basetype_resolution_table = [
    [0, 1, 3, 3, 4],  # Int
    [1, 1, 4, 4, 4],  # Int64
    [3, 4, 2, 3, 4],  # Half
    [3, 4, 3, 3, 4],  # Float
    [4, 4, 4, 4, 4],  # Double
]


class OgnTutorialGenericMathNode:
    """Node to multiple two values of any type"""

    @staticmethod
    def compute(db) -> bool:
        """Compute the product of two values, if the types are all resolved.

        When the types are not compatible for multiplication, or the result type is not compatible with the
        resolved output type, the method will log an error and fail
        """
        try:
            # To support multiplying array of vectors by array of scalars we need to broadcast the scalars to match the
            # shape of the vector array, and we will convert the result to whatever the result is resolved to
            atype = db.inputs.a.type
            btype = db.inputs.b.type
            rtype = db.outputs.product.type

            result_dtype = dtype_from_basetype.get(rtype.base_type, None)

            # Use numpy to perform the multiplication in order to automatically handle both scalar and array types
            # and automatically convert to the resolved output type
            if atype.array_depth > 0 and btype.array_depth > 0 and btype.tuple_count < atype.tuple_count:
                r = np.multiply(db.inputs.a.value, db.inputs.b.value[:, np.newaxis], dtype=result_dtype)
            else:
                r = np.multiply(db.inputs.a.value, db.inputs.b.value, dtype=result_dtype)

            db.outputs.product.value = r
        except TypeError as error:
            db.log_error(f"Multiplication could not be performed: {error}")
            return False

        return True

    @staticmethod
    def on_connection_type_resolve(node) -> None:
        # Resolves the type of the output based on the types of inputs
        atype = node.get_attribute("inputs:a").get_resolved_type()
        btype = node.get_attribute("inputs:b").get_resolved_type()
        productattr = node.get_attribute("outputs:product")
        producttype = productattr.get_resolved_type()

        # The output types can be only inferred when both inputs types are resolved.
        if (
            atype.base_type != og.BaseDataType.UNKNOWN
            and btype.base_type != og.BaseDataType.UNKNOWN
            and producttype.base_type == og.BaseDataType.UNKNOWN
        ):

            # Resolve the base type using the lookup table
            base_type = og.BaseDataType.DOUBLE

            a_index = supported_basetypes.index(atype.base_type)
            b_index = supported_basetypes.index(btype.base_type)

            if a_index >= 0 and b_index >= 0:
                base_type = supported_basetypes[basetype_resolution_table[a_index][b_index]]

            productattr.set_resolved_type(
                og.Type(base_type, max(atype.tuple_count, btype.tuple_count), max(atype.array_depth, btype.array_depth))
            )
