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

# Helper script that will check if the types and operators used in an ONNX model
# are supported by the pre-built ORT Mobile package.

import argparse
import logging
import pathlib
import sys

import onnx
from onnx import shape_inference

from ..onnx_model_utils import get_opsets_imported
from ..reduced_build_config_parser import parse_config

cpp_to_tensorproto_type = {
    "float": 1,
    "uint8_t": 2,
    "int8_t": 3,
    "uint16_t": 4,
    "int16_t": 5,
    "int32_t": 6,
    "int64_t": 7,
    "std::string": 8,
    "bool": 9,
    "MLFloat16": 10,
    "double": 11,
    "uint32_t": 12,
    "uint64_t": 13,
    "Complex64": 14,  # not supported by ORT
    "Complex128": 15,  # not supported by ORT
    "BFloat16": 16,
}

tensorproto_type_to_cpp = {v: k for k, v in cpp_to_tensorproto_type.items()}


def check_graph(graph, opsets, required_ops, global_types, special_types, unsupported_ops, logger):
    """
    Check the graph and any subgraphs for usage of types or operators which we know are not supported.
    :param graph: Graph to process.
    :param opsets: Map of domain to opset version that the model imports.
    :param required_ops: Operators that are included in the pre-built package.
    :param global_types: Types globally enabled in the pre-built package.
    :param special_types: Types that are always enabled for a subset of operators and are _usually_ supported but are
                          not guaranteed to be. We would need to add a lot of infrastructure to know for sure so
                          currently we treat them as supported.
    :param unsupported_ops: Set of unsupported operators that were found.
    :param logger: Logger for diagnostic output.
    :return: Returns whether the graph uses unsupported operators or types.
    """
    has_unsupported_types = False
    value_info_map = {vi.name: vi for vi in graph.value_info}

    def _is_type_supported(value_info, description):
        is_supported = True
        type_name = value_info.type.WhichOneof("value")
        if type_name == "tensor_type":
            t = value_info.type.tensor_type.elem_type
            if t not in global_types and t not in special_types:
                cpp_type = tensorproto_type_to_cpp[t]
                logger.debug(f"Element type {cpp_type} of {description} is not supported.")
                is_supported = False
        else:
            # we don't support sequences, map, sparse tensors, or optional types in the pre-built package
            logger.debug(f"Data type {type_name} of {description} is not supported.")
            is_supported = False

        return is_supported

    def _input_output_is_supported(value_info, input_output):
        return _is_type_supported(value_info, f"graph {input_output} {value_info.name}")

    # node outputs are simpler to check.
    # node inputs have a much wider mix of types, some of which come from initializers and most likely are always
    # enabled as we generally do type reduction on the user data input to the operator and not the weights/etc. which
    # come from initializers.
    def _node_output_is_supported(name):
        is_supported = True
        if name in value_info_map:
            vi = value_info_map[name]
            is_supported = _is_type_supported(vi, f"node output {name}")
        else:
            # we don't have type info so ignore
            pass

        return is_supported

    for i in graph.input:
        if not _input_output_is_supported(i, "input"):
            has_unsupported_types = True

    for o in graph.output:
        if not _input_output_is_supported(o, "output"):
            has_unsupported_types = True

    for node in graph.node:
        # required_ops are map of [domain][opset] to set of op_type names. '' == ai.onnx
        domain = node.domain or "ai.onnx"

        # special case Constant as we will convert to an initializer during model load
        if domain == "ai.onnx" and node.op_type == "Constant":
            continue

        # some models don't have complete imports. use 1 as a default as that's valid for custom domains and should
        # result in an error for any others. not sure why ONNX or ORT validation allows this though.
        opset = opsets[domain] if domain in opsets else 1
        if (
            domain not in required_ops
            or opset not in required_ops[domain]
            or node.op_type not in required_ops[domain][opset]
        ):
            unsupported_ops.add(f"{domain}:{opset}:{node.op_type}")

        for output_name in node.output:
            if not _node_output_is_supported(output_name):
                has_unsupported_types = True

        # recurse into subgraph for control flow nodes (Scan/Loop/If)
        for attr in node.attribute:
            if attr.HasField("g"):
                check_graph(attr.g, opsets, required_ops, global_types, special_types, unsupported_ops, logger)

    return has_unsupported_types or unsupported_ops


def _get_global_tensorproto_types(op_type_impl_filter, logger: logging.Logger):
    """
    Map the globally supported types (C++) to onnx.TensorProto.DataType values used in the model
    See https://github.com/onnx/onnx/blob/1faae95520649c93ae8d0b403816938a190f4fa7/onnx/onnx.proto#L485

    Additionally return a set of types we special case as being able to generally be considered as supported.
    :param op_type_impl_filter: type filter from reduced build configuration parser
    :param logger: Logger
    :return: tuple of globally enabled types and special cased types
    """
    global_cpp_types = op_type_impl_filter.global_type_list()
    global_onnx_tensorproto_types = set()

    for t in global_cpp_types:
        if t in cpp_to_tensorproto_type:
            global_onnx_tensorproto_types.add(cpp_to_tensorproto_type[t])
        else:
            logger.error(f"Error: Unexpected data type of {t} in package build config's globally enabled types.")
            sys.exit(-1)

    # a subset of operators require int32 and int64 to always be enabled, as those types are used for dimensions in
    # shapes and indices.
    # additionally we have a number of operators (e.g. Not, Where) that always require the use of bool.
    # this _may_ mean values involving these types can be processed, but without adding a lot more code we don't know
    # for sure.
    special_types = [
        cpp_to_tensorproto_type["int32_t"],
        cpp_to_tensorproto_type["int64_t"],
        cpp_to_tensorproto_type["bool"],
    ]

    return global_onnx_tensorproto_types, special_types


def get_default_config_path():
    # get default path to config that was used to create the pre-built package.
    script_dir = pathlib.Path(__file__).parent
    local_config = script_dir / "mobile_package.required_operators.config"

    # if we're running in the ORT python package the file should be local. otherwise assume we're running from the
    # ORT repo
    if local_config.exists():
        default_config_path = local_config
    else:
        ort_root = script_dir.parents[3]
        default_config_path = (
            ort_root / "tools" / "ci_build" / "github" / "android" / "mobile_package.required_operators.config"
        )

    return default_config_path


def run_check_with_model(
    model_with_type_info: onnx.ModelProto, mobile_pkg_build_config: pathlib.Path, logger: logging.Logger
):
    """
    Check if an ONNX model can be used with the ORT Mobile pre-built package.
    :param model_with_type_info: ONNX model that has had ONNX shape inferencing run on to add type/shape information.
    :param mobile_pkg_build_config: Configuration file used to build the ORT Mobile package.
    :param logger: Logger for output
    :return: True if supported
    """
    if not mobile_pkg_build_config:
        mobile_pkg_build_config = get_default_config_path()

    enable_type_reduction = True
    config_path = str(mobile_pkg_build_config.resolve(strict=True))
    required_ops, op_type_impl_filter = parse_config(config_path, enable_type_reduction)
    global_onnx_tensorproto_types, special_types = _get_global_tensorproto_types(op_type_impl_filter, logger)

    # get the opset imports
    opsets = get_opsets_imported(model_with_type_info)

    # If the ONNX opset of the model is not supported we can recommend using our tools to update that first.
    supported_onnx_opsets = set(required_ops["ai.onnx"].keys())
    # we have a contrib op that is erroneously in the ai.onnx domain with opset 1. manually remove that incorrect value
    supported_onnx_opsets.remove(1)
    onnx_opset_model_uses = opsets["ai.onnx"]
    if onnx_opset_model_uses not in supported_onnx_opsets:
        logger.info(f"Model uses ONNX opset {onnx_opset_model_uses}.")
        logger.info(f"The pre-built package only supports ONNX opsets {sorted(supported_onnx_opsets)}.")
        logger.info(
            "Please try updating the ONNX model opset to a supported version using "
            "python -m onnxruntime.tools.onnx_model_utils.update_onnx_opset ..."
        )

        return False

    unsupported_ops = set()
    logger.debug(
        "Checking if the data types and operators used in the model are supported " "in the pre-built ORT package..."
    )
    unsupported = check_graph(
        model_with_type_info.graph,
        opsets,
        required_ops,
        global_onnx_tensorproto_types,
        special_types,
        unsupported_ops,
        logger,
    )

    if unsupported_ops:
        logger.info("Unsupported operators:")
        for entry in sorted(unsupported_ops):
            logger.info("  " + entry)

    if unsupported:
        logger.info("\nModel is not supported by the pre-built package due to unsupported types and/or operators.")
        logger.info(
            "Please see https://onnxruntime.ai/docs/reference/mobile/prebuilt-package/ for information "
            "on what is supported in the pre-built package."
        )
        logger.info(
            "A custom build of ONNX Runtime will be required to run the model. Please see "
            "https://onnxruntime.ai/docs/build/custom.html for details on performing that."
        )
    else:
        logger.info("Model should work with the pre-built package.")

    logger.info("---------------\n")

    return not unsupported


def run_check(model_path: pathlib.Path, mobile_pkg_build_config: pathlib.Path, logger: logging.Logger):
    """
    Check if an ONNX model will be able to be used with the ORT Mobile pre-built package.
    :param model_path: Path to ONNX model.
    :param mobile_pkg_build_config: Configuration file used to build the ORT Mobile package.
    :param logger: Logger for output
    :return: True if supported
    """
    logger.info(
        f"Checking if pre-built ORT Mobile package can be used with {model_path} once model is "
        "converted from ONNX to ORT format using onnxruntime.tools.convert_onnx_models_to_ort..."
    )

    model_file = model_path.resolve(strict=True)
    model = onnx.load(str(model_file))

    # we need to run shape inferencing to populate that type info for node outputs.
    # we will get warnings if the model uses ORT contrib ops (ONNX does not have shape inferencing for those),
    # and shape inferencing will be lost downstream of those.
    # TODO: add support for checking ORT format model as it will have full type/shape info for all nodes
    model_with_type_info = shape_inference.infer_shapes(model)

    return run_check_with_model(model_with_type_info, mobile_pkg_build_config, logger)


def main():
    parser = argparse.ArgumentParser(
        description="Check if model can be run using the ONNX Runtime Mobile Pre-Built Package",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--config_path",
        help="Path to required operators and types configuration used to build " "the pre-built ORT mobile package.",
        required=False,
        type=pathlib.Path,
        default=get_default_config_path(),
    )

    parser.add_argument("model_path", help="Path to ONNX model to check", type=pathlib.Path)

    args = parser.parse_args()

    logger = logging.getLogger("default")
    logger.setLevel(logging.INFO)
    run_check(args.model_path, args.config_path, logger)


if __name__ == "__main__":
    main()
