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

import argparse
import logging
import os
from typing import List, Union

import coloredlogs
from constants import (
    AttentionInputIDs,
    AttentionOutputIDs,
    MultiHeadAttentionInputIDs,
    MultiHeadAttentionOutputIDs,
    Operators,
)
from onnx import helper, load_model
from onnx_model import NodeProto, OnnxModel
from shape_infer_helper import SymbolicShapeInferenceHelper

logger = logging.getLogger(__name__)


class PackingAttentionBase:
    def __init__(self, model: OnnxModel, attention_op_type: str):
        self.model: OnnxModel = model
        self.nodes_to_remove: List = []
        self.nodes_to_add: List = []
        self.prune_graph: bool = False
        self.node_name_to_graph_name: dict = {}
        self.this_graph_name: str = self.model.model.graph.name
        self.attention_op_type = attention_op_type
        self.attention_nodes = self.model.get_nodes_by_op_type(attention_op_type)

    def _try_getting_attention_mask(self) -> Union[str, None]:
        mask_index = (
            AttentionInputIDs.MASK_INDEX
            if self.attention_op_type == Operators.ATTENTION
            else MultiHeadAttentionInputIDs.KEY_PADDING_MASK
        )
        first_attention_node = self._try_getting_first_attention()
        # check if attention has mask
        if not first_attention_node or len(first_attention_node.input) <= mask_index:
            return None

        attention_mask = first_attention_node.input[mask_index]

        # check if all attention nodes have same mask
        for node in self.attention_nodes:
            if len(node.input) <= mask_index or node.input[mask_index] != attention_mask:
                return None

        return attention_mask

    def _try_getting_first_attention(self) -> Union[NodeProto, None]:
        if len(self.attention_nodes) <= 0:
            return None

        return self.attention_nodes[0]

    def _try_getting_last_layernorm(self) -> Union[NodeProto, None]:
        last_layernorm_node = None
        for node in self.model.nodes():
            if node.op_type == Operators.LAYERNORM or node.op_type == Operators.SKIPLAYERNORM:
                last_layernorm_node = node
        return last_layernorm_node

    def _are_attentions_supported(self) -> bool:
        raise NotImplementedError()

    def _insert_removepadding_node(self, inputs: List[str], outputs: List[str]) -> None:
        new_node = helper.make_node(
            Operators.REMOVEPADDING,
            inputs=inputs,
            outputs=outputs,
            name=self.model.create_node_name(Operators.REMOVEPADDING),
        )

        new_node.domain = "com.microsoft"
        self.nodes_to_add.append(new_node)
        self.node_name_to_graph_name[new_node.name] = self.this_graph_name

    def _insert_restorepadding_node(self, inputs: List[str], outputs: List[str]) -> None:
        new_node = helper.make_node(
            Operators.RESTOREPADDING,
            inputs=inputs,
            outputs=outputs,
            name=self.model.create_node_name(Operators.RESTOREPADDING),
        )

        new_node.domain = "com.microsoft"
        self.nodes_to_add.append(new_node)
        self.node_name_to_graph_name[new_node.name] = self.this_graph_name

    def _replace_attention_with_packing_attention(self, token_offset: str, cumulative_sequence_length: str) -> None:
        raise NotImplementedError()

    def _get_input_to_remove_padding(self, first_attention_node) -> Union[str, None]:
        if self.attention_op_type == Operators.ATTENTION:
            return first_attention_node.input[AttentionInputIDs.INPUT]
        return None

    def convert(self, use_symbolic_shape_infer: bool = True) -> None:
        logger.debug("start converting to packing model...")

        if not self._are_attentions_supported():
            return

        attention_mask = self._try_getting_attention_mask()
        if not attention_mask:
            return

        first_attention_node = self._try_getting_first_attention()
        last_layernorm_node = self._try_getting_last_layernorm()
        if not last_layernorm_node:
            return

        # insert RemovePadding
        input_to_remove_padding = self._get_input_to_remove_padding(first_attention_node)
        if not input_to_remove_padding:
            return

        output_without_padding = input_to_remove_padding + "_no_padding"
        token_offset = input_to_remove_padding + "_token_offset"
        cumulated_seq_len = input_to_remove_padding + "_cumulated_seq_len"
        max_seq_len = input_to_remove_padding + "_max_seq_len"
        self._insert_removepadding_node(
            [input_to_remove_padding, attention_mask],
            [output_without_padding, token_offset, cumulated_seq_len, max_seq_len],
        )
        self.model.replace_input_of_all_nodes(input_to_remove_padding, output_without_padding)
        logger.debug("inserted RemovePadding before Attention")

        # insert RestorePadding
        restorepadding_input = last_layernorm_node.output[0] + "_restore_input"
        self._insert_restorepadding_node([restorepadding_input, token_offset], [last_layernorm_node.output[0]])
        self.model.replace_output_of_all_nodes(last_layernorm_node.output[0], restorepadding_input)
        logger.debug(f"inserted RestorePadding after last {last_layernorm_node.op_type} layer")

        # insert PackedAttention
        self._replace_attention_with_packing_attention(token_offset, cumulated_seq_len)
        logger.debug(f"replaced {self.attention_op_type} with Packed{self.attention_op_type}")

        self.model.remove_nodes(self.nodes_to_remove)
        self.model.add_nodes(self.nodes_to_add, self.node_name_to_graph_name)

        if self.prune_graph:
            self.model.prune_graph()
        elif self.nodes_to_remove or self.nodes_to_add:
            self.model.update_graph()
        self.model.clean_shape_infer()
        if use_symbolic_shape_infer:
            # Use symbolic shape inference since custom operators (like Gelu, SkipLayerNormalization etc)
            # are not recognized by onnx shape inference.
            shape_infer_helper = SymbolicShapeInferenceHelper(self.model.model, verbose=0)
            inferred_model = shape_infer_helper.infer_shapes(self.model.model, auto_merge=True, guess_output_rank=False)
            if inferred_model:
                self.model.model = inferred_model


class PackingAttention(PackingAttentionBase):
    def __init__(self, model: OnnxModel):
        super().__init__(model, Operators.ATTENTION)

    def _are_attentions_supported(self) -> bool:
        for node in self.attention_nodes:
            if OnnxModel.get_node_attribute(node, "past_present_share_buffer") is not None:
                return False
            if OnnxModel.get_node_attribute(node, "do_rotary") is not None:
                return False
            unidirection_attr = OnnxModel.get_node_attribute(node, "unidirectional")
            if unidirection_attr is not None and unidirection_attr != 0:
                return False
            if len(node.input) > AttentionInputIDs.PAST and not node.input[AttentionInputIDs.PAST]:
                return False
            if (
                len(node.input) > AttentionInputIDs.PAST_SEQUENCE_LENGTH
                and not node.input[AttentionInputIDs.PAST_SEQUENCE_LENGTH]
            ):
                return False
        return True

    def _replace_attention_with_packing_attention(self, token_offset: str, cumulative_sequence_length: str) -> None:
        for attention in self.attention_nodes:
            attention_bias = (
                attention.input[AttentionInputIDs.ATTENTION_BIAS]
                if len(attention.input) > AttentionInputIDs.ATTENTION_BIAS
                else ""
            )
            packed_attention = helper.make_node(
                Operators.PACKEDATTENTION,
                inputs=[
                    attention.input[AttentionInputIDs.INPUT],
                    attention.input[AttentionInputIDs.WEIGHTS],
                    attention.input[AttentionInputIDs.BIAS],
                    token_offset,
                    cumulative_sequence_length,
                    attention_bias,
                ],
                outputs=[attention.output[AttentionOutputIDs.OUTPUT]],
                name=self.model.create_node_name(Operators.PACKEDATTENTION),
            )

            attributes = []
            for attr in attention.attribute:
                if attr.name in ["num_heads", "qkv_hidden_sizes", "scale"]:
                    attributes.append(attr)

            packed_attention.attribute.extend(attributes)
            packed_attention.domain = "com.microsoft"
            self.nodes_to_add.append(packed_attention)
            self.nodes_to_remove.append(attention)
            self.node_name_to_graph_name[packed_attention.name] = self.this_graph_name

        logger.info("Converted %d Attention nodes to PackedAttention.", len(self.attention_nodes))


class PackingMultiHeadAttention(PackingAttentionBase):
    def __init__(self, model: OnnxModel):
        super().__init__(model, Operators.MULTI_HEAD_ATTENTION)

    def _check_empty_input(self, node, index: int, name: str):
        """Check a node does not have given input."""
        if len(node.input) > index:
            if len(node.input[index]) > 0:
                logger.error(f"node input {index} ({name}) is not supported in PackedMultiHeadAttention: {node}")
                return False
        return True

    def _check_empty_output(self, node, index: int, name: str):
        """Check a node does not have given input."""
        if len(node.output) > index:
            if len(node.output[index]) > 0:
                logger.error(f"node output {index} ({name}) is not supported in PackedMultiHeadAttention: {node}")
                return False
        return True

    def _are_attentions_supported(self) -> bool:
        for node in self.attention_nodes:
            for attr in node.attribute:
                if attr.name not in ["num_heads", "mask_filter_value", "scale"]:
                    logger.error(f"node attribute {attr.name} is not supported in PackedMultiHeadAttention: {node}")
                    return False

            if node.input[MultiHeadAttentionInputIDs.KEY] and not node.input[MultiHeadAttentionInputIDs.VALUE]:
                logger.error("packed kv format is not supported in PackedMultiHeadAttention")
                return False

            if not (
                self._check_empty_input(node, MultiHeadAttentionInputIDs.PAST_KEY, "past_key")
                and self._check_empty_input(node, MultiHeadAttentionInputIDs.PAST_VALUE, "past_key")
                and self._check_empty_output(node, MultiHeadAttentionOutputIDs.PRESENT_KEY, "present_key")
                and self._check_empty_output(node, MultiHeadAttentionOutputIDs.PRESENT_VALUE, "present_key")
            ):
                return False

        return True

    def _replace_attention_with_packing_attention(self, token_offset: str, cumulative_sequence_length: str) -> None:
        gated_relative_pos_bias_count = 0
        for mha in self.attention_nodes:
            attention_bias = (
                mha.input[MultiHeadAttentionInputIDs.ATTENTION_BIAS]
                if len(mha.input) > MultiHeadAttentionInputIDs.ATTENTION_BIAS
                else ""
            )
            packed_mha = helper.make_node(
                Operators.PACKED_MULTI_HEAD_ATTENTION,
                inputs=[
                    mha.input[MultiHeadAttentionInputIDs.QUERY],
                    mha.input[MultiHeadAttentionInputIDs.KEY],
                    mha.input[MultiHeadAttentionInputIDs.VALUE],
                    mha.input[MultiHeadAttentionInputIDs.BIAS],
                    token_offset,
                    cumulative_sequence_length,
                    attention_bias,
                ],
                outputs=[mha.output[MultiHeadAttentionOutputIDs.OUTPUT]],
                name=self.model.create_node_name(Operators.PACKED_MULTI_HEAD_ATTENTION),
            )

            attributes = []
            for attr in mha.attribute:
                if attr.name in ["num_heads", "mask_filter_value", "scale"]:
                    attributes.append(attr)

            packed_mha.attribute.extend(attributes)
            packed_mha.domain = "com.microsoft"
            self.nodes_to_add.append(packed_mha)
            self.nodes_to_remove.append(mha)
            self.node_name_to_graph_name[packed_mha.name] = self.this_graph_name

            # Append token_offset input to GatedRelativePositionBias
            if attention_bias:
                rel_pos_bias_node = self.model.get_parent(mha, MultiHeadAttentionInputIDs.ATTENTION_BIAS)
                if (
                    rel_pos_bias_node
                    and rel_pos_bias_node.op_type == "GatedRelativePositionBias"
                    and len(rel_pos_bias_node.input) == 6
                ):
                    rel_pos_bias_node.input.append(token_offset)
                    gated_relative_pos_bias_count += 1

        logger.info("Converted %d MultiHeadAttention nodes to PackedMultiHeadAttention.", len(self.attention_nodes))
        logger.info("Converted %d GatedRelativePositionBias nodes to packing mode.", gated_relative_pos_bias_count)

    def _get_input_to_remove_padding(self, first_attention_node) -> Union[str, None]:
        # When there are query, key and value inputs, we need to find the first input of the parent MatMul node.
        matmul = self.model.get_parent(first_attention_node, 0)
        if matmul and matmul.op_type == "MatMul":
            return matmul.input[0]
        return None


class PackingMode:
    def __init__(self, model: OnnxModel):
        self.model = model

    def convert(self, use_symbolic_shape_infer: bool = True) -> None:
        if self.model.get_nodes_by_op_type(Operators.ATTENTION):
            if self.model.get_nodes_by_op_type(Operators.MULTI_HEAD_ATTENTION):
                logger.error("Packing mode does not support both Attention and MultiHeadAttention in same graph.")
                return None
            packing = PackingAttention(self.model)
            return packing.convert(use_symbolic_shape_infer)
        elif self.model.get_nodes_by_op_type(Operators.MULTI_HEAD_ATTENTION):
            packing = PackingMultiHeadAttention(self.model)
            return packing.convert(use_symbolic_shape_infer)
        else:
            logger.error("Packing mode requires either Attention or MultiHeadAttention node in onnx graph.")
            return None


def _parse_arguments():
    parser = argparse.ArgumentParser(
        description="Convert to packing mode tool for ONNX Runtime. It converts BERT like model to use packing mode."
    )
    parser.add_argument("--input", required=True, type=str, help="input onnx model path")

    parser.add_argument("--output", required=True, type=str, help="optimized onnx model path")

    parser.add_argument("--verbose", required=False, action="store_true", help="show debug information.")
    parser.set_defaults(verbose=False)

    parser.add_argument(
        "--use_external_data_format",
        required=False,
        action="store_true",
        help="use external data format to store large model (>2GB)",
    )
    parser.set_defaults(use_external_data_format=False)

    args = parser.parse_args()

    return args


def _setup_logger(verbose):
    if verbose:
        coloredlogs.install(
            level="DEBUG",
            fmt="[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s",
        )
    else:
        coloredlogs.install(fmt="%(funcName)20s: %(message)s")


def main():
    args = _parse_arguments()

    _setup_logger(args.verbose)

    logger.debug(f"arguments:{args}")

    if os.path.realpath(args.input) == os.path.realpath(args.output):
        logger.warning("Specified the same input and output path. Note that this may overwrite the original model")

    model = load_model(args.input)
    packing_mode = PackingMode(OnnxModel(model))
    packing_mode.convert()
    packing_mode.model.save_model_to_file(args.output, use_external_data_format=args.use_external_data_format)


if __name__ == "__main__":
    main()