# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from __future__ import annotations

import logging
from dataclasses import dataclass

import onnx

from ...quant_utils import QuantType
from ...tensor_quant_overrides import QuantTypeInfo, TensorQuantOverridesHelper


@dataclass
class TensorTypeRequest:
    """
    Bundles desired quantization type requests for a tensor. A distinction is made between the
    produced type and the consumed type.
    """

    # The tensor's quant type at the producer end. If None, assumed to be the default activation quant type.
    producer: QuantTypeInfo | None

    # The tensor's quant type received by a set of consumer nodes.
    # If None, assumed to be the default activation quant type for all consumers.
    # consumers[1] is a set of consumer node names.
    consumers: tuple[QuantTypeInfo, set[str]] | None


class MixedPrecisionTensorQuantOverridesFixer:
    """
    Helper that generates tensor quantization overrides for mixed-precision QDQ models.

    Specifically, this helper fixes an initial set of quantization overrides that assign a non-default
    activation quantization type to one or more tensors by doing the following:
     - Inferring which other tensors need to be overridden to the non-default activation quantization type.
     - Inserting quantization data type conversions.

    Example:
    --------

    Float model:

    input_0 --> Op1 --> Op3 --> Op5 --> Op6 --> output_0
                                 ^
                                 |
    input_1 --> Op2 -+-> Op4 ----+
                     |
                     +-> Op7 --> output_1
                     |
                     +-> Op8 --> output_2

    If we'd like to quantize this model to uint8 precision, but would like to make sure tensor "Op4_out"
    is quantized to 16-bit, then we would specify the following initial tensor quantization overrides:

    ```
    init_overrides = {"Op4_out": [{"quant_type": QuantType.QUInt16}]}
    ```

    These initial overrides may not create a valid model because Op4 and Op5 may require both the input and output
    to be the same type (e.g., uint16). This helper fixes the overrides so that input/output data types
    are valid:

    ```
    overrides = TensorQuantOverridesHelper(init_overrides)

    fixer = MixedPrecisionTensorQuantOverridesFixer.create_from_model(overrides, model, QuantType.QUInt8)
    fixer.apply(
        default_activation_qtype=QuantType.QUInt8,
        default_activation_symmetric=False,
    )
    ```

    The above snippet generates the following "fixed" overrides (get via overrides.get_dict()):

    {
      "Op2_out": [{"quant_type": QUInt8, "convert": {"quant_type": QUInt16, "recv_nodes": {"Op4"}}}],
      "Op3_out": [{"quant_type": QUInt8, "convert": {"quant_type": QUInt16, "recv_nodes": {"Op5"}}}],
      "Op4_out": [{"quant_type": QUInt16}],
      "Op5_out": [{"quant_type": QUInt16, "convert": {"quant_type": QUInt8, "recv_nodes": {"Op6"}}}]
    }

    How to interpret the fixed overrides:
    - Op2's output is consumed by Op4, Op7, and Op8. Op4 consumes the converted u16 type,
      but Op7 and Op8 consume the original u8 type.
    - Op3's output is converted from u8 to u16. Op5 consumes the converted u16 type.
    - Op4's output is just u16 (not converted). All consumers of Op4_out get the u16 type.
    - Op5's output is converted from u16 to u8. Op6 consumes the u8 type.
    """

    def __init__(
        self,
        overrides: TensorQuantOverridesHelper,
        producers: dict[str, onnx.NodeProto],
        consumers: dict[str, list[onnx.NodeProto]],
        value_infos: dict[str, onnx.ValueInfoProto],
        initializers: dict[str, onnx.TensorProto],
    ):
        """
        Params:
            overrides: The initial tensor quantization overrides to fix.
            producers: Dictionary that maps a tensor name to the producer node that generates the tensor.
            consumers: Dictionary that maps a tensor name to the consumer nodes that take the tensor as input.
            value_infos: Dictionary that maps a tensor name to its onnx.ValueInfoProto.
            initializers: Dictionary that maps an initializer name to its onnx.TensorProto.
        """
        self.overrides = overrides
        self.consumers = consumers
        self.producers = producers
        self.value_infos = value_infos
        self.initializers = initializers

    @staticmethod
    def create_from_model(
        overrides: TensorQuantOverridesHelper, model: onnx.ModelProto, default_activation_qtype: QuantType
    ) -> MixedPrecisionTensorQuantOverridesFixer:
        """
        Helper function that creates an instance of this class from a loaded ONNX model.

        Params:
            overrides: The initial tensor quantization overrides to fix.
            model: Loaded ONNX model
            default_activation_qtype: The intended default activation quantization type.
                                      Used to validate the initial overrides.

        Returns:
            Initialized MixedPrecisionTensorQuantOverridesFixer object
        """
        model = onnx.shape_inference.infer_shapes(model)  # Need to infer shapes to get value_infos

        # Build dictionaries that enable convenient lookups of initializers and value_infos by name.
        initializers = {initializer.name: initializer for initializer in model.graph.initializer}
        value_infos = {vi.name: vi for vi in model.graph.value_info}
        value_infos.update({ot.name: ot for ot in model.graph.output})
        value_infos.update({it.name: it for it in model.graph.input})

        # Ensure that the user-provided initial overrides are actually valid.
        valid, err = overrides.is_valid(initializers, set(value_infos), default_activation_qtype)
        if not valid:
            pprint_overrides = overrides.pprint_str(indent=4)
            logging.error(f"Provided invalid tensor quantization overrides:\n{pprint_overrides}")
            raise ValueError(err)

        consumers = {}
        producers = {}

        # Build dictionaries that map a tensor name to the consumer or producer nodes.
        for node in model.graph.node:
            for input_name in node.input:
                if input_name:
                    if input_name not in consumers:
                        consumers[input_name] = []

                    consumers[input_name].append(node)

            for output_name in node.output:
                producers[output_name] = node

        return MixedPrecisionTensorQuantOverridesFixer(overrides, producers, consumers, value_infos, initializers)

    def apply(
        self,
        default_activation_qtype: QuantType,
        default_activation_symmetric: bool,
    ):
        """
        Fixes the initial tensor quantization overrides (in-place) for use in mixed-precision QDQ models.

        Params:
            default_activation_qtype: The intended default activation quantization type.
            default_activation_symmetric: The intended default symmetry used to quantize activations.
        """
        type_requests = self.get_desired_tensor_types(default_activation_qtype, default_activation_symmetric)

        # Use type requests to "fix" tensor quantization overrides by adding
        # quantization type conversions where necessary.
        for tensor_name, type_req in type_requests.items():
            all_consumers = set([node.name for node in self.consumers.get(tensor_name, [])])
            has_producer_req = type_req.producer is not None
            has_consumer_req = bool(type_req.consumers)

            # Only producer type: Add conversion back to default activation type
            if has_producer_req and not has_consumer_req:
                self._update_converted_tensor(
                    tensor_name, type_req.producer, QuantTypeInfo(default_activation_qtype), all_consumers
                )
            # Only consumers
            elif not has_producer_req and has_consumer_req:
                prod_type_info = self.overrides.get_node_output_qtype_info(tensor_name, default_activation_qtype)
                consumer_type_info = type_req.consumers[0]

                if prod_type_info != consumer_type_info:
                    self._update_converted_tensor(
                        tensor_name, prod_type_info, consumer_type_info, type_req.consumers[1]
                    )
                else:
                    if not self._check_nodes_are_not_convert_consumers(tensor_name, type_req.consumers[1]):
                        raise ValueError(
                            f"Tensor override for '{tensor_name}' converts the type for consumers that need the original type."
                        )
            # Both producer and consumers
            elif has_producer_req and has_consumer_req:
                prod_type_info = type_req.producer
                consumer_type_info = type_req.consumers[0]

                if prod_type_info != consumer_type_info:
                    self._update_converted_tensor(
                        tensor_name, prod_type_info, consumer_type_info, type_req.consumers[1]
                    )
                else:
                    consumers_for_original_type = all_consumers.difference(type_req.consumers[1])

                    if len(consumers_for_original_type) == 0:
                        # All consumers want the overridden type, so no need for convert nodes!
                        # Just add the override to the new new if not already present.
                        if tensor_name not in self.overrides:
                            self.overrides[tensor_name] = [{}]
                            prod_type_info.save_to_dict(self.overrides[tensor_name][0])

                        assert "convert" not in self.overrides[tensor_name][0]
                    else:
                        # Some consumers don't want the overridden type.
                        self._update_converted_tensor(
                            tensor_name,
                            prod_type_info,
                            QuantTypeInfo(default_activation_qtype),
                            consumers_for_original_type,
                        )
            else:
                raise ValueError(f"TypeRequest for tensor {tensor_name} has no producer or consumers.")

        # Done. Check if the overrides are valid.
        valid, err = self.overrides.is_valid(self.initializers, set(self.value_infos), default_activation_qtype)
        if not valid:
            pprint_overrides = self.overrides.pprint_str(indent=4)
            logging.error(
                f"Generated invalid tensor quantization overrides for mixed-precision QDQ model:\n{pprint_overrides}"
            )
            raise ValueError(err)

    def get_desired_tensor_types(
        self,
        default_activation_qtype: QuantType,
        default_activation_symmetric: bool,
    ) -> dict[str, TensorTypeRequest]:
        """
        Iterates through the initial tensor quantization overrides and builds a set of TensorTypeRequests objects
        that describe the quantization types required at each tensor. These TensorTypeRequests objects are ultimately
        used to generated the "fixed" overrides.

        Params:
            default_activation_qtype: The intended default activation quantization type.
            default_activation_symmetric: The intended default symmetry used to quantize activations.

        Returns:
            TensorTypeRequest objects as a dict that maps a tensor name to its requested types.
        """
        type_requests = {}
        default_activation_type_info = QuantTypeInfo(default_activation_qtype, default_activation_symmetric)

        # Scan tensor overrides for type conversion requests.
        for tensor_name, override_list in self.overrides.items():
            if not self.__is_tensor_quantizable(tensor_name):
                continue  # Skip non-quantizable tensors (e.g., not a float)

            if tensor_name in self.initializers:
                continue  # Skip initializers

            if not override_list or len(override_list) > 1:
                continue  # Skip per-channel stuff

            override_dict = override_list[0]
            quant_type_info = QuantTypeInfo.load_from_dict(override_dict, default_activation_type_info.quant_type)
            producer_node = self.producers.get(tensor_name)  # None if this is a model input

            if quant_type_info != default_activation_type_info and "convert" not in override_dict:
                if producer_node is not None:
                    self._add_type_requests_for_node(type_requests, quant_type_info, producer_node)

                # Find all consumer nodes of `tensor_name` and update their inputs/outputs to the new type.
                for consumer_node in self.consumers.get(tensor_name, []):
                    self._add_type_requests_for_node(type_requests, quant_type_info, consumer_node)

        return type_requests

    def _add_type_requests_for_node(
        self,
        type_requests: dict[str, TensorTypeRequest],
        quant_type_info: QuantTypeInfo,
        node: onnx.NodeProto,
    ):
        """
        Adds TensorTypeRequest objects for a given node, assuming that we want all its inputs and outputs
        to have the same quantization type (as specified by the `quant_type_info` parameter).

        Params:
            type_requests: Dictionary of type requests to append to for this node.
            quant_type_info: The quantization type to use for inputs and outputs.
            node: The node for which the TensorTypeRequest objects are created and added to type_requests.
        """
        # Add output side
        for output_name in node.output:
            if not self.__is_tensor_quantizable(output_name):
                continue

            if output_name not in type_requests:
                type_requests[output_name] = TensorTypeRequest(quant_type_info, None)
            else:
                if (
                    type_requests[output_name].producer is not None
                    and type_requests[output_name].producer != quant_type_info
                ):
                    raise ValueError(f"Tensor {output_name} has multiple types.")

                type_requests[output_name].producer = quant_type_info

        # Add the consumer side
        for input_name in node.input:
            if input_name and input_name not in self.initializers and self.__is_tensor_quantizable(input_name):
                if input_name not in type_requests:
                    type_requests[input_name] = TensorTypeRequest(None, None)

                if type_requests[input_name].consumers is None:
                    type_requests[input_name].consumers = (quant_type_info, set())

                if type_requests[input_name].consumers[0] != quant_type_info:
                    raise ValueError(f"Tensor {input_name} has consumers requesting different types.")

                if not node.name:
                    raise ValueError(
                        f"Node of type {node.op_type} with output 0 {node.output[0]} does not have a name!"
                    )

                type_requests[input_name].consumers[1].add(node.name)

    def _update_converted_tensor(
        self,
        tensor_name: str,
        producer_type_info: QuantTypeInfo,
        consumer_type_info: QuantTypeInfo,
        consumer_names: set[str],
    ):
        """
        Updates the tensor quantization overrides for a tensor that is converted from one type to another.

        Params:
            tensor_name: The name of the tensor for which to update overrides.
            producer_type_info: Info for the tensor's produced type.
            consumer_type_info: Info for the tensor's consumed (i.e., converted) type.
            consumer_names: Nodes names of consumers that consume the converted type.
        """
        if tensor_name not in self.overrides or not self.overrides[tensor_name]:
            self.overrides[tensor_name] = [{}]
            producer_type_info.save_to_dict(self.overrides[tensor_name][0])

        overrides = self.overrides[tensor_name][0]
        if producer_type_info != QuantTypeInfo.load_from_dict(overrides):
            raise ValueError(f"Desired producer quant_type for {tensor_name} doesn't match existing type.")

        if consumer_names:
            if "convert" not in overrides:
                overrides["convert"] = {}
                consumer_type_info.save_to_dict(overrides["convert"])

            convert_dict = overrides["convert"]
            if consumer_type_info != QuantTypeInfo.load_from_dict(convert_dict):
                raise ValueError(f"Desired consumer quant_type for {tensor_name} doesn't match existing type.")

            if "recv_nodes" not in convert_dict:
                convert_dict["recv_nodes"] = set()

            convert_dict["recv_nodes"].update(consumer_names)

    def _check_nodes_are_not_convert_consumers(self, tensor_name: str, node_names: set[str]):
        """
        Returns true if the given nodes do not consume/receive a converted quantization type.

        Params:
            tensor_name: The name of the tensor to check.
            node_names: Set of node names that should not be consumers of the converted type.
        """
        if tensor_name not in self.overrides or not self.overrides[tensor_name]:
            return True

        overrides = self.overrides[tensor_name][0]

        if "convert" not in overrides:
            return True

        convert_dict = overrides["convert"]

        if "recv_nodes" not in convert_dict:
            return False

        return not convert_dict["recv_nodes"].intersection(node_names)

    def __is_tensor_quantizable(self, tensor_name):
        weight = self.initializers.get(tensor_name)
        if weight is not None:
            if weight.data_type in (onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16):
                return True
        elif tensor_name in self.value_infos:
            vi = self.value_infos[tensor_name]
            if vi.type.HasField("tensor_type") and vi.type.tensor_type.elem_type in (
                onnx.TensorProto.FLOAT,
                onnx.TensorProto.FLOAT16,
            ):
                return True

        return False