#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import copy
from collections import OrderedDict

from polygraphy import mod, util
from polygraphy.common import TensorMetadata
from polygraphy.datatype import DataType
from polygraphy.logger import G_LOGGER, LogMode

gs = mod.lazy_import("onnx_graphsurgeon")
onnx = mod.lazy_import("onnx")
onnx_numpy_helper = mod.lazy_import("onnx.numpy_helper")


def get_num_nodes(model):
    def _get_num_graph_nodes(graph):
        num_nodes = len(graph.node)
        for node in graph.node:
            for attr in node.attribute:
                if attr.type == onnx.AttributeProto.GRAPH:
                    num_nodes += _get_num_graph_nodes(attr.g)
                elif attr.type == onnx.AttributeProto.GRAPHS:
                    for subgraph in attr.graphs:
                        num_nodes += _get_num_graph_nodes(subgraph)
        return num_nodes

    return _get_num_graph_nodes(model.graph)


def all_tensor_names(model, include_inputs=None):
    include_inputs = util.default(include_inputs, False)

    all_outputs = [
        output
        for node in model.graph.node
        if node.op_type != "Constant"
        for output in node.output
    ]
    if include_inputs:
        all_outputs += [inp.name for inp in model.graph.input]
    all_outputs = util.unique_list(all_outputs)
    return all_outputs


def _check_has_tensors(model, outputs):
    all_outputs = all_tensor_names(model, include_inputs=True)
    util.check_sequence_contains(
        all_outputs, outputs, name="the model", items_name="outputs", check_extra=False
    )


def mark_outputs(model, outputs):
    # Clear the old outputs
    while model.graph.output:
        model.graph.output.pop()

    outputs = util.unique_list(outputs)
    _check_has_tensors(model, outputs)

    value_info_map = {t.name: t for t in model.graph.value_info}
    out_tensors = []
    for output in outputs:
        value_info = value_info_map.get(
            output, onnx.helper.make_empty_tensor_value_info(output)
        )
        out_tensors.append(value_info)

    G_LOGGER.ultra_verbose(f"Marked output tensors in ONNX model: {out_tensors}")
    model.graph.output.extend(out_tensors)
    return model


def mark_layerwise(model):
    # Add all non-constant node outputs as graph outputs
    model = mark_outputs(model, all_tensor_names(model))
    return model


def unmark_outputs(model, outputs):
    outputs = util.unique_list(outputs)
    _check_has_tensors(model, outputs)

    cur_outputs = []
    while model.graph.output:
        cur_outputs.append(model.graph.output.pop())
    cur_outputs = list(reversed(cur_outputs))  # Preserve ordering

    for out in cur_outputs:
        if out.name not in outputs:
            model.graph.output.extend([out])

    return model


def get_shape(tensor):
    shape = []
    if isinstance(tensor, onnx.TensorProto):
        shape = tensor.dims
    else:
        for dim in tensor.type.tensor_type.shape.dim:
            if dim.HasField("dim_param"):
                shape.append(dim.dim_param)
            elif dim.HasField("dim_value"):
                shape.append(dim.dim_value)
            else:
                shape.append(-1)
    return shape


def get_dtype(tensor):
    if isinstance(tensor, onnx.TensorProto):
        onnx_type = tensor.data_type
    else:
        onnx_type = tensor.type.tensor_type.elem_type
    return DataType.from_dtype(onnx_type, source_module="onnx")


def get_values(tensor):
    try:
        return onnx_numpy_helper.to_array(tensor)
    except Exception as err:
        G_LOGGER.error(
            f"Failed to load weights.\nNote: Error was: {err}", mode=LogMode.ONCE
        )
    return "<error: failed to load weights>"


def get_tensor_metadata(tensors):
    metadata = TensorMetadata()
    for tensor in tensors:
        metadata.add(name=tensor.name, dtype=get_dtype(tensor), shape=get_shape(tensor))
    return metadata


def get_input_metadata(graph):
    # Some "inputs" are actually weights with initalizers, so we need to eliminate those.
    initializer_names = {tensor.name for tensor in graph.initializer}
    input_tensors = [
        tensor for tensor in graph.input if tensor.name not in initializer_names
    ]
    return get_tensor_metadata(input_tensors)


def get_output_metadata(graph):
    return get_tensor_metadata(graph.output)


def str_from_onnx(model, show_layers=None, show_attrs=None, show_weights=None):
    """
    Converts an ONNX Graph to a human-readable representation

    Args:
        graph (onnx.GraphProto): The onnx graph.
        show_layers (bool): Whether to display per-layer information.
        show_attrs (bool): Whether to display per-layer attributes.
        show_weights (bool): Whether to display the value of weights.

    Returns:
        str
    """
    show_layers = util.default(show_layers, False)
    show_attrs = util.default(show_attrs, False)
    show_weights = util.default(show_weights, False)

    def get_opset():
        default_opset = "Unknown"
        other_opsets = {}
        for info in model.opset_import:
            if not info.domain:
                default_opset = info.version
            else:
                other_opsets[info.domain] = info.version
        return default_opset, other_opsets

    default_opset, other_opsets = get_opset()
    onnx_str = ""
    onnx_str += f"Name: {model.graph.name} | ONNX Opset: {default_opset}"
    if other_opsets:
        onnx_str += f" | Other Opsets: {other_opsets}"
    onnx_str += "\n\n"

    onnx_str += str_from_onnx_graph(
        model.graph,
        tensors={},
        show_layers=show_layers,
        show_attrs=show_attrs,
        show_weights=show_weights,
    )
    return onnx_str


def str_from_onnx_graph(
    graph, tensors, show_layers, show_attrs, show_weights, indent_level=0
):
    input_metadata = get_input_metadata(graph)
    output_metadata = get_output_metadata(graph)
    initializer_metadata = get_tensor_metadata(graph.initializer)

    # Subgraph inputs should remain separate from each other, hence copy the tensors map
    tensors = copy.copy(tensors)
    tensors.update(get_tensor_metadata(graph.value_info))
    tensors.update(initializer_metadata)
    tensors.update(input_metadata)
    tensors.update(output_metadata)

    graph_type = "Graph" if indent_level == 0 else "Subgraph"

    onnx_str = ""
    if show_attrs and graph.doc_string:
        onnx_str += f"---- Docstring ----\n{graph.doc_string}\n\n"

    onnx_str += (
        f"---- {len(input_metadata)} {graph_type} Input(s) ----\n{input_metadata}\n\n"
    )
    onnx_str += f"---- {len(output_metadata)} {graph_type} Output(s) ----\n{output_metadata}\n\n"

    onnx_str += f"---- {len(initializer_metadata)} Initializer(s) ----\n"
    if show_weights:
        for init in graph.initializer:
            onnx_str += f"Initializer | {init.name} [dtype={get_dtype(init)}, shape={get_shape(init)}] | Values:\n{util.indent_block(str(get_values(init)))}\n\n"
        if not graph.initializer:
            onnx_str += "{}\n\n"
    elif show_layers:
        onnx_str += str(initializer_metadata)
        onnx_str += "\n\n"
    else:
        onnx_str += "\n"

    def get_names_and_meta(names):
        names_lst = []
        metadata = TensorMetadata()
        for name in names:
            dtype, shape = tensors.get(name, (None, None))
            if name in initializer_metadata:
                name = f"Initializer | {name}"
            names_lst.append(name)
            metadata.add(name=name, dtype=dtype, shape=shape)
        return names_lst, metadata

    # Maps values from the AttributeType enum to their string representations, e.g., {1: "FLOAT"}
    ATTR_TYPE_MAPPING = dict(
        zip(
            onnx.AttributeProto.AttributeType.values(),
            onnx.AttributeProto.AttributeType.keys(),
        )
    )

    # Maps an ONNX attribute to the corresponding Python property
    ONNX_PYTHON_ATTR_MAPPING = {
        "FLOAT": "f",
        "INT": "i",
        "STRING": "s",
        "TENSOR": "t",
        "GRAPH": "g",
        "FLOATS": "floats",
        "INTS": "ints",
        "STRINGS": "strings",
    }

    def attrs_to_dict(attrs):
        attr_dict = OrderedDict()
        for attr in attrs:

            def process_attr(attr_str: str):
                processed = getattr(attr, ONNX_PYTHON_ATTR_MAPPING[attr_str])
                if attr_str == "STRING":
                    processed = processed.decode()
                elif attr_str == "TENSOR":
                    tensor_str = f"Tensor: [dtype={get_dtype(processed)}, shape={get_shape(processed)}]"
                    if show_weights:
                        tensor_str += " | Values:\n" + util.indent_block(
                            str(get_values(processed))
                        )
                    processed = tensor_str
                elif attr_str == "GRAPH":
                    processed = "\n" + str_from_onnx_graph(
                        processed,
                        tensors,
                        indent_level=indent_level + 2,
                        show_layers=show_layers,
                        show_attrs=show_attrs,
                        show_weights=show_weights,
                    )
                elif attr_str == "FLOATS" or attr_str == "INTS":
                    # Proto hacky list to normal Python list
                    processed = [p for p in processed]
                elif attr_str == "STRINGS":
                    processed = [p.decode() for p in processed]
                return processed

            if attr.type in ATTR_TYPE_MAPPING:
                attr_str = ATTR_TYPE_MAPPING[attr.type]
                if attr_str in ONNX_PYTHON_ATTR_MAPPING:
                    attr_dict[attr.name] = process_attr(attr_str)
                else:
                    G_LOGGER.warning(
                        f"Attribute of type {attr_str} is currently unsupported. Skipping attribute."
                    )
            else:
                G_LOGGER.warning(
                    f"Attribute type: {attr.type} was not recognized. Was the graph generated with a newer IR version than the installed `onnx` package? Skipping attribute."
                )
        return attr_dict

    onnx_str += f"---- {len(graph.node)} Node(s) ----\n"
    if show_layers:
        for index, node in enumerate(graph.node):
            input_names, input_meta = get_names_and_meta(node.input)
            output_names, output_meta = get_names_and_meta(node.output)

            onnx_str += util.str_from_layer(
                "Node",
                index,
                node.name,
                node.op_type,
                input_names,
                input_meta,
                output_names,
                output_meta,
            )

            if show_attrs:
                attrs = attrs_to_dict(node.attribute)
                if attrs:
                    onnx_str += util.indent_block("---- Attributes ----") + "\n"
                for key, val in attrs.items():
                    attr_str = ""
                    if node.name:
                        attr_str += f"{node.name}."
                    onnx_str += util.indent_block(f"{attr_str}{key} = {val}") + "\n"
            onnx_str += "\n"

    return util.indent_block(onnx_str, indent_level)


##
## ONNX-GraphSurgeon utilities
##


def meta_from_gs_tensors(tensors):
    """Get TensorMetadata from a list of ONNX-GraphSurgeon tensors"""
    meta = TensorMetadata()
    for tensor in tensors:
        meta.add(tensor.name, tensor.dtype, tensor.shape)
    return meta


def set_shapes_from_layerwise_meta(graph, layerwise_meta):
    """
    Args:
        graph (gs.Graph): An ONNX graphsurgeon graph.
        layerwise_meta (TensorMetadata): Metadata for tensors in the graph.
    """
    for tensor in graph.tensors().values():
        if isinstance(tensor, gs.Variable) and tensor.name in layerwise_meta:
            tensor.shape = layerwise_meta[tensor.name].shape
            tensor.dtype = DataType.to_dtype(
                DataType.from_dtype(layerwise_meta[tensor.name].dtype), "onnx"
            )


def lower_constant_nodes(graph):
    """Converts the outputs of Constant nodes into constant tensors, removing the nodes"""
    remove_nodes = set()
    with graph.node_ids():
        for node in graph.nodes:
            if node.op == "Constant" and "value" in node.attrs:
                node.outputs[0].to_constant(node.attrs["value"].values)
                remove_nodes.add(node.id)
        # Iterate from the end so we don't shift the list under us.
        for node_id in sorted(remove_nodes, reverse=True):
            del graph.nodes[node_id]
    return graph


def get_unbounded_dds_tensors(graph):
    graph.toposort()
    # A dict of operators that might produce a output tensor with unbounded DDS, when the value of the input tensor
    # at the corresponding index is a runtime value. For example, "Range" => "1" means that if the input 1 of the Range
    # operator is a runtime value, e.g. not a const tensor or an initializer, then the Range output tensor size is unbounded.
    dispatcher_dict = {
        "Range": [1],  # the limit input of the Range operator
        "Pad": [1],  # the pads input of the Pad operator
        "Resize": [3],  # the sizes input of the Resize operator
        "Tile": [1],  # the repeats input of the Tile operator
        "Expand": [1],  # the shape input of the Expand operator
    }

    # Check if the given operator produces a output tensor with unbounded DDS.
    def check_op(node, const_tensor_set):
        # Check if the operator is inside the dispatcher dict.
        if node.op in dispatcher_dict:
            input_idx_list = dispatcher_dict[node.op]
            for input_idx in input_idx_list:
                if input_idx < len(node.inputs):
                    input_tensor = node.inputs[input_idx]
                    # Check if the corresponding input tensor is a runtime value and its producer is not Min operator.
                    # If a tensor is produced by a Min operator, its upper bound has already been set.
                    if (
                        input_tensor.name not in const_tensor_set
                        and len(input_tensor.inputs) >= 1
                        and input_tensor.inputs[0].op != "Min"
                    ):
                        return input_tensor
        return None

    # Find all constant tensors.
    def get_const_tensors(graph):
        return {
            tensor.name
            for tensor in graph.tensors().values()
            if isinstance(tensor, gs.Constant)
        }

    # Find all dynamic shape symbols, customers will set upper bounds for these symbols when building the model in TensorRT.
    def get_dynamic_shapes(graph):
        dynamic_shape_set = set()
        for tensor in graph.inputs:
            for shape in tensor.shape:
                if isinstance(shape, str):
                    dynamic_shape_set.add(shape)
        return dynamic_shape_set

    # Find all tensors with unbounded DDS.
    def get_target_tensors(graph):
        # Find dynamic shapes, these shapes should have upper bounds in TensorRT.
        dynamic_shape_set = get_dynamic_shapes(graph)

        # Find const tensors. For those operators in the dispatch dict, constant inputs will not introduce outputs with unbounded DDS.
        const_tensor_set = get_const_tensors(graph)

        # Our target is to find those input tensors that cause its consumer nodes generated unbounded outputs.
        # If a tensor has named dimensions that appeared before in its symbolic shape, it means that the shape is *not* data dependent,
        # and so will have an upper bound.
        target_tensor_names = set()
        target_tensor_list = []
        for node in graph.nodes:
            check_node = False
            # Check if the node's output contains a new introduced dynamic shape.
            for tensor in node.outputs:
                # Always check nodes if tensor.shape is None.
                # This happens when the symbolic inference does not work correctly due to some restrictions.
                if tensor.shape is None:
                    check_node = True
                else:
                    for shape in tensor.shape:
                        # If a shape is a dynamic shape, then it is a str.
                        # Only check the node that first introduced the dynamic shape.
                        if isinstance(shape, str) and shape not in dynamic_shape_set:
                            dynamic_shape_set.add(shape)
                            check_node = True
            # Check if the node will generate an unbounded output size.
            if check_node:
                target_tensor = check_op(node, const_tensor_set)
                # Avoid duplication.
                if (
                    target_tensor is not None
                    and target_tensor.name not in target_tensor_names
                ):
                    target_tensor_names.add(target_tensor.name)
                    target_tensor_list.append(target_tensor)
        return target_tensor_list

    return get_target_tensors(graph)