# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence

import numpy as np
import onnx
import onnx.helper  # noqa: TID251
from onnx.defs import OpSchema

from onnxscript import ir, tensor

if TYPE_CHECKING:
    from onnxscript import converter

# Conversions from python values to ONNX are used by both the script converter as well
# as the eager-mode runtime and both need to be consistent. The script converter converts
# python values into ONNX TensorProto, while the runtime converts python values into
# ONNXScript runtime's value-representation (based on Tensor).


# Utilities to convert a python value to TensorProto (for use by the script converter)


def pyvalue_to_onnx_tensor(tensor_name: str, pyvalue):
    return ir.serde.serialize_tensor(ir.tensor(pyvalue, name=tensor_name))


_REPEATED_ATTRIBUTE_TYPES = frozenset(
    {
        onnx.AttributeProto.FLOATS,
        onnx.AttributeProto.INTS,
        onnx.AttributeProto.STRINGS,
        onnx.AttributeProto.TENSORS,
        onnx.AttributeProto.GRAPHS,
        onnx.AttributeProto.SPARSE_TENSORS,
        onnx.AttributeProto.TYPE_PROTOS,
    }
)


def pyvalue_to_onnx_attribute(
    key: str,
    value: Any,
    name_generator: Callable[[], str],
    attr_type: onnx.AttributeProto.AttributeType | None = None,
) -> onnx.AttributeProto:
    """Helper function to create an ONNX AttributeProto.

    This is a refinement of onnx.helper.make_attribute that works with ONNX Script
    conventions for allowed types for attribute-values. In particular, it allows
    * Empty lists as attribute values, provided the attribute type is specified
    and is a list type.
    * Scalar-values like 1.0 as well as lists like [1, -1] to be specified
    when the attribute type is TensorProto by automatically converting the value
    into a 0-D or 1-D tensor respectively.
    """
    if isinstance(value, list) and not value:
        # Empty list value:
        if attr_type is None:
            raise ValueError("Attribute type must be specified for empty list value.")
        if attr_type not in _REPEATED_ATTRIBUTE_TYPES:
            raise ValueError("Empty list value is only allowed for repeated attribute types.")
        return onnx.AttributeProto(name=key, type=attr_type)
    elif attr_type == onnx.AttributeProto.TENSOR and not isinstance(value, onnx.TensorProto):
        return onnx.AttributeProto(
            name=key, type=attr_type, t=pyvalue_to_onnx_tensor(name_generator(), value)
        )
    else:
        # When the value is a subgraph, ONNX IR will complain that some values are
        # not found from the scope.
        return onnx.helper.make_attribute(key, value)  # noqa: TID251


# Utilities to convert python values into onnxscript tensors.


def _promotable(x) -> bool:
    """Checks if a runtime parameter value needs to be promoted into an onnxscript value.
    This is the runtime-equivalent of the promotion of literal constants into ONNX values
    in the static converter.
    """
    if isinstance(x, (bool, int, float)):
        return True
    if isinstance(x, list) and x:
        # Note: This is meant to handle valid scenarios correctly. No attempt is
        # made yet to capture all invalid usages in runtime mode.
        return _promotable(x[0])
    return False


def _get_dtype(pyvalue):
    """Return np.dtype to use when converting a python value to an onnxscript tensor.
    Note that int constants are treated as int64, as that is the common type in ONNX
    for shape/index values.
    """
    if isinstance(pyvalue, bool):
        return np.bool_
    elif isinstance(pyvalue, int):
        return np.int64
    elif isinstance(pyvalue, float):
        return np.float32
    elif isinstance(pyvalue, list):
        if pyvalue:
            # TODO: What to do about lists with mixed value types, like [1, 2.0]?
            # Should at least produce an error/warning message.
            return _get_dtype(pyvalue[0])
        raise ValueError("Cannot determine target type for empty list")
    raise TypeError(f"Value of unexpected type {type(pyvalue)}")


def cast_pyvalue_to_os_tensor(pyvalue, dtype=None):
    """Promotes python values into onnxscript tensors.
    The optional argument dtype specifies the desired np.dtype of the tensor,
    used only when a non-standard onnxscript-value is promoted into one.
    """
    if _promotable(pyvalue):
        if dtype is None:
            dtype = _get_dtype(pyvalue)
        return tensor.Tensor(np.array(pyvalue, dtype=dtype))
    return pyvalue


def cast_inputs(
    get_type_info: Callable[[Any], Any],
    cast: Callable[[Any, Any], Any],
    op_schema: OpSchema | None,
    args,
) -> tuple[Any, ...]:
    """Uses schema specification to support a limited form of auto-casting.

    * Scalars are promoted to tensors.
    * Further. they are cast to the required type when used in ops with other
    tensor inputs that are required to be of same type.
    Thus, in "A+1" or "Add(A, 1)", the value 1 will be converted to the same
    type as A.

    This is used by the converter in a static-mode, as well as by the eager-mode
    execution in a dynamic-mode.
    """
    if op_schema is None:
        # Either an error or a custom op.
        # No checks/casts in this case.
        return tuple(cast(x, None) for x in args)

    expected_inputs = op_schema.inputs
    # We make two passes. In the first pass, we identify known type-bindings for
    # type-variables: eg., {'T1' : np.float32, 'T2' : np.int32}.
    # In the second pass, we use these bindings to cast scalar-values to
    # tensors of appropriate types. The two passes are needed to handle cases
    # like "Add(1, X)" where 1 must be cast to the same type as X.
    type_bindings: dict[Optional[str], np.dtype] = {}
    args_typevars: list[tuple[str, Optional[str]]] = []
    for i, x in enumerate(args):
        if i < len(expected_inputs):
            expected = expected_inputs[i]
        elif expected_inputs[-1].option == OpSchema.FormalParameterOption.Variadic:
            expected = expected_inputs[-1]
            if not expected.is_homogeneous:
                args_typevars.append((x, None))
                continue
        else:
            raise ValueError(
                f"Number of actual parameters {len(args)} "
                f"exceeds number of formal parameters {len(expected_inputs)}."
            )
        typevar = expected.type_str
        if "(" not in typevar:
            # typevar is an identifier, like "T"
            typeinfo = get_type_info(x)
            if typeinfo is not None:
                type_bindings[typevar] = typeinfo
        args_typevars.append((x, typevar))
    cast_args = [cast(x, type_bindings.get(typevar)) for x, typevar in args_typevars]
    return tuple(cast_args)


def dynamic_cast_inputs(op_schema: OpSchema, args):
    """Used for autocast during eager-mode execution."""

    def get_type_info(x):
        return x.dtype if isinstance(x, tensor.Tensor) else None

    return cast_inputs(get_type_info, cast_pyvalue_to_os_tensor, op_schema, args)


def static_cast_inputs(
    converter_: converter.Converter,
    op_schema: Optional[OpSchema],
    args: Sequence[Optional[converter.Variable]],
) -> tuple[str, ...]:
    """Used for autocast during script-translation.
    This is meant to transform expressions like "Add(X, 1)" to "Add(X, CastLike(1, X))"
    Polymorphic constants (like 0 and 1) are cast to the type of other operands as needed.
    """

    def get_type_info(x: Optional[converter.Variable]) -> Optional[converter.Variable]:
        """Returns x back if x can serve as the target-type for a cast (as the second
        argument of CastLike) and None otherwise. In the expression "Add(X, 1), 1 is
        castable, while X can serve as the target-type.
        """
        return None if x is None or x.is_castable else x

    def cast_like(
        x: Optional[converter.Variable], y: Optional[converter.Variable]
    ) -> Optional[str]:
        if x is None:
            return None
        if x.is_castable and y is not None:
            # Polymorphic constant x is cast to the type of y:
            x_cast = converter_.generate_unique_name(f"{x.name}_cast")
            converter_.emit([x_cast], "CastLike", [x.name, y.name])
            return x_cast
        return x.name

    return cast_inputs(get_type_info, cast_like, op_schema, args)
