import enum
import json
from collections import OrderedDict
from typing import Callable, Optional
import inspect
import numbers

import torch
import torch.nn as nn
from torch import _VF
import torch.nn.functional as F
from torch.quantization.qconfig import QConfig
from intel_extension_for_pytorch.nn.functional import interaction

from ._quantization_state_utils import QTensorInfo


add_and_mul_ops = set([
    torch.add,
    torch.Tensor.add,
    # TODO, mul op?
])

quantized_modules_has_weights = set([
    torch.nn.Conv2d,
    torch.nn.Conv3d,
    torch.nn.Linear,
    torch.nn.EmbeddingBag,
    torch.nn.ConvTranspose2d,
    torch.nn.ConvTranspose3d,
    torch.nn.LSTM,
    ])

# those ops only support int8->int8, not int8->fp32/bf16
int8_int8_ops = set([
    str(F.adaptive_avg_pool2d),
    str(F.adaptive_avg_pool3d),
    str(F.avg_pool2d),
    str(F.avg_pool3d),
    str(F.max_pool2d),
    str(F.max_pool3d),
    str(nn.MaxPool2d),
    str(nn.MaxPool3d),
    str(nn.AvgPool2d),
    str(nn.AvgPool3d),
    str(nn.AdaptiveAvgPool2d),
    str(nn.AdaptiveAvgPool3d),
    str(torch.Tensor.relu),
    str(torch.relu),
    str(F.relu),
    str(nn.ReLU),
    str(torch.flatten),
    str(torch.Tensor.flatten),
    str(torch.nn.Flatten),
    # the following op will be supported at next step.
    #str(torch.Tensor.sigmoid),
    #str(torch.sigmoid),
    #str(F.sigmoid),
    #str(nn.Sigmoid),
    #str(F.gelu),
    #str(nn.GELU),
    # ipex customer op
    str(interaction),
    str(torch.ops.torch_ipex.interaction_forward),
    str(torch.embedding_bag),
    str(F.embedding_bag),
    str(torch.nn.EmbeddingBag),
    str(torch.nn.LSTM),
])

class OpQuantizeabilityType(enum.Enum):
    QUANTIZEABLE = 0
    NOT_QUANTIZEABLE = 1

class FuncOutputObsType(enum.Enum):
    NONE = 0
    NEW_OBS = 1
    REUSES_FIRST_INPUT_OBS = 2

# quantizeable modules
def is_leaf(
    m: torch.nn.Module,
) -> bool:

    # TODO(future PR): extend to the rest of the container classes
    container_classes = (
        torch.nn.Sequential,
        torch.nn.ModuleList,
    )
    return (
        # allowlist everything in torch.nn except containers
        (m.__module__.startswith('torch.nn') and (
            not isinstance(m, container_classes)
        ))
    )

def get_fqn_valid_for_module_dict_key(fqn: str) -> str:
    """
    Modifies `fqn`(fully qualified name) to make it a valid key to a ModuleDict.
    """
    if fqn == '':
        fqn = ' '
    return fqn.replace('.', ':')

class HookType(enum.Enum):
    """
    Describes the various types of function and module hooks that are used
    to implement quantization syntax transforms.
    """
    # Hooks which are run before, during and after a quantizeable op.
    # Usually used for op input and output observation, subsituating
    # quantized kernels, and dynamically looking up arguments to quantized
    # kernels.
    OP_HOOKS = 0
    # Hooks which are run before or after a `torch.nn.Module` which
    # is a non-leaf. Usually used for dtype transforms if the user requests
    # that the inputs or outputs of a certain module are of some dtype.
    MODULE_IO_HOOKS = 1
    # Hooks which are run before a non-quantizeable op which requires
    # `torch.float` inputs. Any inputs which are not floats are converted
    # back to floats.
    ARG_DEQUANTS = 2
    # Everything else
    NONE = 2

def get_torch_function_hook_type(
    parent_module: Optional[torch.nn.Module],
    func: Callable,
) -> HookType:
    # the direct __dict__ accesses are for performance, because
    # the default `torch.nn.Module.__getattr__` has overhead.
    parent_module_has_qstate = parent_module is not None and \
        '_auto_quant_state' in parent_module.__dict__
    needs_op_hooks = parent_module_has_qstate and \
        parent_module.__dict__['_auto_quant_state'].cur_op_needs_hooks(func)  # type: ignore[union-attr, operator]

    if needs_op_hooks:
        return HookType.OP_HOOKS
    elif (
        parent_module_has_qstate
    ):
        return HookType.ARG_DEQUANTS
    else:
        return HookType.NONE

def get_module_hook_type(
    parent_module: Optional[torch.nn.Module],
    cur_module: torch.nn.Module,
) -> HookType:
    cached_hook_type = getattr(cur_module, '_auto_quant_module_hook_type', None)
    if cached_hook_type is not None:
        return cached_hook_type
    parent_module_has_qstate = parent_module is not None and \
        '_auto_quant_state' in parent_module.__dict__
    needs_op_hooks = parent_module_has_qstate and \
        parent_module.__dict__['_auto_quant_state'].cur_op_needs_hooks(cur_module)  # type: ignore[union-attr, operator]
    # We need IO hooks if
    # * we are calling forward on a module (always True here)
    # * that module has quant state
    # * that module does not need op hooks for the parent
    needs_io_hooks = (
        '_auto_quant_state' in cur_module.__dict__ and
        (not needs_op_hooks)
    )
    needs_arg_dequants = parent_module_has_qstate and not needs_op_hooks

    if needs_op_hooks:
        result = HookType.OP_HOOKS
    elif needs_io_hooks:
        result = HookType.MODULE_IO_HOOKS
    elif needs_arg_dequants:
        result = HookType.ARG_DEQUANTS
    else:
        result = HookType.NONE
    cur_module._auto_quant_module_hook_type = result  # type: ignore[assignment]
    return result


def attach_scale_zp_values_to_model(
    module: torch.nn.Module,
) -> None:
    """
    Calculates the scale and zero_point from each observer and attaches
    these values to the parent module. This is done to avoid recalculating
    these values at inference.
    """
    if hasattr(module, '_auto_quant_state'):
        qstate: AutoQuantizationState = module._auto_quant_state  # type: ignore[assignment]
        quantized_dtype =  [torch.quint8, torch.qint8]
        for tensor_id, observer in qstate.tensor_id_to_observer.items():
            if observer.dtype in quantized_dtype:
                scale, zp = observer.calculate_qparams()
                qstate.tensor_id_to_scale_zp[int(tensor_id)] = (scale, zp)
            else:
                assert False, "The observer's dtype only can be torch.quint8 or torch.qint8"
        for tensor_id, observer in qstate.weight_tensor_id_to_observer.items():
            if observer.dtype in quantized_dtype:
                scale, zp = observer.calculate_qparams()
                qstate.weight_tensor_id_to_scale_zp[tensor_id] = (scale, zp)
            else:
                assert False, "The observer's dtype only can be torch.quint8 or torch.qint8"
        qstate.tensor_id_to_observer.clear()
        qstate.weight_tensor_id_to_observer.clear()

    for _, child in module.named_children():
        attach_scale_zp_values_to_model(child)


def _check_observer_has_run(observer):
    if observer.min_val.numel() == 0 or observer.max_val.numel() == 0:
        return False
    if (observer.min_val.dim() == 0 or observer.max_val.dim() == 0) and \
            observer.min_val == float("inf") and observer.max_val == float("-inf"):
        return False
    return True


def check_model_obsever_has_run(
    module: torch.nn.Module,
) -> None:
    """
    This function is about check whether the module's observer has been run by checking the
    observer's min_value and max_max_value is the init value or not.
    """
    if hasattr(module, '_auto_quant_state'):
        qstate: AutoQuantizationState = module._auto_quant_state  # type: ignore[assignment]
        quantized_dtype =  [torch.quint8, torch.qint8]
        for tensor_id, observer in qstate.tensor_id_to_observer.items():
            if observer.dtype in quantized_dtype:
                return _check_observer_has_run(observer)
            else:
                assert False, "The observer's dtype only can be torch.quint8 or torch.qint8"
        for tensor_id, observer in qstate.weight_tensor_id_to_observer.items():
            if observer.dtype in quantized_dtype:
                return _check_observer_has_run(observer)
            else:
                assert False, "The observer's dtype only can be torch.quint8 or torch.qint8"

    for _, child in module.named_children():
        check_model_obsever_has_run(child)

    return True


def attach_op_convert_info_to_model(
    module: torch.nn.Module,
) -> None:
    """
    Calculates the info needed to convert each op and attaches
    it to the parent module. This is done to avoid recalculating these values
    at inference.
    """
    if hasattr(module, '_auto_quant_state'):
        qstate: AutoQuantizationState = module._auto_quant_state  # type: ignore[assignment]
        for _, seen_q_op_info in qstate.idx_to_seen_q_op_infos.items():
            qstate.idx_to_op_convert_info[seen_q_op_info.idx] = \
                qstate.calculate_op_convert_info(seen_q_op_info)
            qstate.idx_to_op_weight_convert_info[seen_q_op_info.idx] = \
                qstate.calculate_op_weight_convert_info(seen_q_op_info)

    for _, child in module.named_children():
        attach_op_convert_info_to_model(child)

class Node:
    def __init__(self, op_infos, input_scale_zero=None, weight_scale_zero=None, output_scale_zero=None, qconfig=None):
        self.idx = op_infos.idx if hasattr(op_infos, 'idx') else None
        self.type = op_infos.type
        self.fqn = op_infos.fqn
        self.input_tensor_infos = op_infos.input_tensor_infos
        self.input_tensor_force_inf_dtype = [] if qconfig is None else op_infos.input_tensor_force_inf_dtype
        self.output_tensor_infos = op_infos.output_tensor_infos
        self.insert_fake_quant_after_outputs = [] if qconfig is None else op_infos.insert_fake_quant_after_outputs
        self.weight_tensor_infos = [] if qconfig is None else op_infos.weight_tensor_infos
        self.qconfig = qconfig
        self.input_scale_zero = input_scale_zero
        self.weight_scale_zero = weight_scale_zero
        self.output_scale_zero = output_scale_zero
        self.pre_nodes = []
        self.post_nodes = []

class ParentNode:
    def __init__(self, output_info):
        self.output_info = output_info

def convert_quant_state_map_to_nodes(quant_state_map):
    nodes = []
    # step1: create nodes
    for _, v in quant_state_map.items():
        new_parent_node = ParentNode(v.output_qtensor_infos)
        nodes.append(new_parent_node)
        for idx, q_op_info in v.idx_to_seen_q_op_infos.items():
            input_scale_zero = {}
            weight_scale_zero = {}
            output_scale_zero = {}
            for tensor_info in q_op_info.input_tensor_infos:
                if tensor_info is not None and tensor_info.id in v.tensor_id_to_scale_zp:
                    input_scale_zero[tensor_info.id] = v.tensor_id_to_scale_zp[tensor_info.id]
            for tensor_info in q_op_info.weight_tensor_infos:
                if tensor_info is not None and (str(idx) + "_" + str(tensor_info.id) in v.weight_tensor_id_to_scale_zp):
                    weight_id = str(idx) + "_" + str(tensor_info.id)
                    weight_scale_zero[weight_id] = v.weight_tensor_id_to_scale_zp[weight_id]
            for tensor_info in q_op_info.output_tensor_infos:
                if tensor_info is not None and tensor_info.id in v.tensor_id_to_scale_zp:
                    output_scale_zero[tensor_info.id] = v.tensor_id_to_scale_zp[tensor_info.id]
            new_node = Node(q_op_info, input_scale_zero, weight_scale_zero, output_scale_zero, qconfig=q_op_info.qconfig)
            nodes.append(new_node)
        for nonq_op_infos in v.seen_nonq_op_infos:
            new_node = Node(nonq_op_infos)
            nodes.append(new_node)
    # create connection between nodess
    for cur in nodes:
        if isinstance(cur, ParentNode):
            continue
        for n in nodes:
            if isinstance(n, ParentNode):
                continue
            # find pre_node:
            for input_info in cur.input_tensor_infos:
                if input_info in n.output_tensor_infos:
                    cur.pre_nodes.append(n)
            # find nex_node:
            for output_info in cur.output_tensor_infos:
                if output_info in n.input_tensor_infos:
                    cur.post_nodes.append(n)
    return nodes

def sync_pool_and_lstm_input_output_scale_zp(quant_state_map, nodes):
    pool_ops = [str(F.adaptive_avg_pool2d), str(F.adaptive_avg_pool3d), str(F.avg_pool2d), str(F.avg_pool3d), \
        str(F.max_pool2d), str(F.max_pool3d), str(nn.MaxPool2d), str(nn.MaxPool3d), str(nn.AvgPool2d), str(nn.AvgPool3d), \
        str(nn.AdaptiveAvgPool2d), str(nn.AdaptiveAvgPool3d)]
    shape_ops = [str(torch.flatten), str(torch.nn.Flatten), str(torch.Tensor.flatten)]
    rnn_ops = [str(torch.nn.LSTM)]
    def _sync_scale_zp_given_id(quant_state_map, id, scale_zp):
        for _, v in quant_state_map.items():
            if id in v.tensor_id_to_scale_zp:
                v.tensor_id_to_scale_zp[id] = scale_zp

    def _find_sync_op_from_given_node(cur_node, ids):
        for next in cur_node.post_nodes:
            if next.type in (pool_ops + shape_ops + rnn_ops):
                ids.append(next.output_tensor_infos[0].id)
                _find_sync_op_from_given_node(next, ids)

    for node in nodes:
        if isinstance(node, ParentNode):
            continue
        if node.qconfig is not None and node.type in (pool_ops + rnn_ops):
            if node.input_scale_zero == node.output_scale_zero:
                continue
            sync_node_begin = node
            # fistly, find the fist sync op before the cur pooling(lstm) op.
            # like: pooling->pool->shape->cur_node,
            while len(sync_node_begin.pre_nodes) == 1 and \
                sync_node_begin.pre_nodes[0].type in (pool_ops + shape_ops + rnn_ops):
                sync_node_begin = sync_node_begin.pre_nodes[0]
            tensor_ids = [sync_node_begin.output_tensor_infos[0].id]
            scale_zp = sync_node_begin.input_scale_zero[sync_node_begin.input_tensor_infos[0].id]
            # find all need sync op from sync_node_begin.
            _find_sync_op_from_given_node(sync_node_begin, tensor_ids)
            for id in tensor_ids:
                _sync_scale_zp_given_id(quant_state_map, id, scale_zp)


def _check_after_nodes_all_quantized_give_node(node):
    r"""
    This function is about check whether given node's post nodes are all quantized,
    """
    if len(node.post_nodes) == 0:
        return False
    else:
        # make sure all post nodes are quantizabled. 
        for next in node.post_nodes:
            if next.qconfig is None:
                if next.type == str(nn.Identity):
                    return _check_after_nodes_all_quantized_give_node(next)
                else:
                    return False
            else:
                int8_int8_symmetric_ops = [str(interaction), str(torch.ops.torch_ipex.interaction_forward), str(torch.embedding_bag), \
                    str(F.embedding_bag), str(torch.nn.EmbeddingBag)]
                if next.type in int8_int8_symmetric_ops:
                    if next.type in [str(interaction), str(torch.ops.torch_ipex.interaction_forward)]:
                        # node.input_tensor_infos may be set, we can use force_inf_dtype to check whether this op is quantizabled.
                        for force_inf_dtype in next.input_tensor_force_inf_dtype:
                            if force_inf_dtype != torch.qint8:
                                return False
                    else:
                        # embeddingBag
                        if next.weight_tensor_infos[0].inf_dtype != torch.qint8:
                            return False
                else:
                    for force_inf_dtype in next.input_tensor_force_inf_dtype:
                        if force_inf_dtype not in [torch.qint8, torch.quint8]:
                            return False
        # all post nodes are quantizabled.
        return True

def set_node_output_quantized(nodes):
    r"""
    # For interation EmbeddingBag, pooling and elt_wise, we only support INT8->INT*, if those ops have quantized their inputs,
    # we need make sure their output also have falk quant to make them call in INT8 kernel.
    # this function will check whether the output inf dtype is int8 dtype if its' input is set to quantized, if the
    # output's infe dtype is not int8, set it and also set insert_fake_quant_after_output to True.
    """
    def _reset_post_node_input_infos(node):
        # make sure the post node will insert fake quant if we add fake quant by cur node' output
        if len(node.post_nodes) > 0:
            for post_node in node.post_nodes:
                if post_node.qconfig is not None:
                    for idx, tensor_info in enumerate(post_node.input_tensor_infos):
                        if tensor_info in node.output_tensor_infos:
                            post_node.input_tensor_force_inf_dtype[idx] = tensor_info.orig_dtype
                elif post_node.type == str(nn.Identity):
                    _reset_post_node_input_infos(post_node)

    for node in nodes:
        if isinstance(node, ParentNode):
            continue
        if node.qconfig is not None and node.type in int8_int8_ops:
            post_node_are_quantized = _check_after_nodes_all_quantized_give_node(node)
            if node.type in str(torch.nn.EmbeddingBag):
                if node.weight_tensor_infos[0].inf_dtype == torch.qint8 and not post_node_are_quantized:
                    node.output_tensor_infos[0].inf_dtype = torch.qint8
                    node.insert_fake_quant_after_outputs[0] = True
                    _reset_post_node_input_infos(node)
            elif node.type  == str(F.embedding_bag):
                if node.input_tensor_force_inf_dtype[1] == torch.qint8 and not post_node_are_quantized:
                    node.output_tensor_infos[0].inf_dtype = torch.qint8
                    node.insert_fake_quant_after_outputs[0] = True
                    _reset_post_node_input_infos(node)
            elif node.type in [str(interaction), str(torch.ops.torch_ipex.interaction_forward)]:
                if node.input_tensor_force_inf_dtype[0] == torch.qint8 and not post_node_are_quantized:
                    node.output_tensor_infos[0].inf_dtype = torch.qint8
                    node.insert_fake_quant_after_outputs[0] = True
                    _reset_post_node_input_infos(node)
            else:
                # TODO: enable PackedSequence input for LSTM.
                if not (node.type in [nn.LSTM] and len(node.input_tensor_infos) > 2 and node.input_tensor_infos[1].orig_dtype == torch.int64):
                    if node.input_tensor_force_inf_dtype[0] in [torch.qint8, torch.quint8] and not post_node_are_quantized:
                        node.output_tensor_infos[0].inf_dtype = node.input_tensor_force_inf_dtype[0]
                        node.insert_fake_quant_after_outputs[0] = True
                        _reset_post_node_input_infos(node)

qscheme_dict = {
    str(torch.per_tensor_affine): torch.per_tensor_affine,
    str(torch.per_tensor_symmetric): torch.per_tensor_symmetric,
    str(torch.per_channel_affine): torch.per_channel_affine,
    str(torch.per_channel_symmetric): torch.per_channel_symmetric
}

dtype_dict = {
    str(torch.quint8): torch.quint8,
    str(torch.qint8): torch.qint8,
    str(torch.float32): torch.float32,
    str(torch.float64): torch.float64,
    str(torch.float16): torch.float16,
    str(torch.bfloat16): torch.bfloat16,
    str(torch.complex64): torch.complex64,
    str(torch.complex128): torch.complex128,
    str(torch.int16): torch.int16,
    str(torch.int32): torch.int32,
    str(torch.int64): torch.int64,
    str(torch.bool): torch.bool,
    str(torch.uint8): torch.uint8,
    str(torch.int8): torch.int8,
    str(torch.quint4x2): torch.quint4x2,
}

def _get_observer_setting(observer):
    r"""
    Convert torch observer's args to dict for saving to json file.
    Because json only accept number or string, so we will convert some
    class type to string(dtype, qscheme, other class type?).
    """
    observer_setting = OrderedDict()
    observer_setting["name"] =  observer.__class__.__name__
    # get observer arg name
    observer_args = inspect.signature(observer.__init__).parameters
    observer_dict = observer.__dict__
    # Now we only can save number or string to json file.
    for arg_name in observer_args.keys():
        if arg_name in observer_dict:
            if isinstance(observer_dict[arg_name], numbers.Number):
                observer_setting[arg_name] = observer_dict[arg_name]
            elif arg_name == "qscheme" or arg_name == "dtype":
                observer_setting[arg_name] = str(observer_dict[arg_name])
            elif arg_name == "eps" and hasattr(observer, "eps") and \
                isinstance(observer.eps, torch.Tensor) and observer.eps.numel() == 1:
                observer_setting[arg_name] = observer.eps.item()
    return observer_setting

def _create_observer(setting):
    r"""
    Create torch observer according to user's setting.
    """
    if "qscheme" in setting:
        setting["qscheme"] = qscheme_dict[setting["qscheme"]]
    if "dtype" in setting:
        setting["dtype"] = dtype_dict[setting["dtype"]]

    if hasattr(torch.quantization.observer, setting["name"]):
        observer = getattr(torch.quantization.observer, setting["name"])
        setting.pop("name", None)
        return observer.with_args(**setting)
    else:
        raise NameError('torch.quantization.observer %s not found' % setting["name"])

def save_quant_state(quant_state_map, configure_file):
    # save qparam's as json file for tunning
    quant_state_dict = OrderedDict()
    for k, v in quant_state_map.items():
        layer_infos = OrderedDict()
        if len(v.idx_to_seen_q_op_infos) == 0:
            layer_infos["q_op_infos"] = {}
        else:
            q_op_infos = OrderedDict()
            for q_k, op_info in v.idx_to_seen_q_op_infos.items():
                info = OrderedDict()
                info["op_type"] = op_info.type
                info["op_type_is_module"] = op_info.type_is_module
                info["fqn"] = op_info.fqn
                input_tensor_infos = []
                for tensor_info, force_dtype in zip(op_info.input_tensor_infos, op_info.input_tensor_force_inf_dtype):
                    cur_tensor_infos = {}

                    if tensor_info is not None:
                        cur_tensor_infos["id"] = tensor_info.id
                        cur_tensor_infos["orig_dtype"] = str(tensor_info.orig_dtype)
                        cur_tensor_infos["inf_dtype"] = str(tensor_info.inf_dtype)
                        cur_tensor_infos["force_dtype"] = str(force_dtype)
                        if tensor_info.id in v.tensor_id_to_scale_zp:
                            cur_tensor_infos["scale"] = v.tensor_id_to_scale_zp[tensor_info.id][0].tolist()
                            cur_tensor_infos["zero_point"] = v.tensor_id_to_scale_zp[tensor_info.id][1].tolist()
                    input_tensor_infos.append(cur_tensor_infos)
                info["input_tensor_infos"] = input_tensor_infos
                # weight infos
                weight_tensor_infos = []
                for tensor_info in op_info.weight_tensor_infos:
                    cur_tensor_infos = {}
                    if tensor_info is not None:
                        cur_tensor_infos["orig_dtype"] = str(tensor_info.orig_dtype)
                        cur_tensor_infos["inf_dtype"] = str(tensor_info.inf_dtype)
                        weight_idx = str(op_info.idx) + "_" + str(tensor_info.id)
                        if weight_idx in v.weight_tensor_id_to_scale_zp:
                            cur_tensor_infos["scale"] = v.weight_tensor_id_to_scale_zp[weight_idx][0].tolist()
                            cur_tensor_infos["zero_point"] = v.weight_tensor_id_to_scale_zp[weight_idx][1].tolist()
                    weight_tensor_infos.append(cur_tensor_infos)
                info["weight_tensor_infos"] = weight_tensor_infos
                # output infos
                output_tensor_infos = []
                for tensor_info in op_info.output_tensor_infos:
                    cur_tensor_infos = {}
                    if tensor_info is not None:
                        cur_tensor_infos["id"] = tensor_info.id
                        cur_tensor_infos["orig_dtype"] = str(tensor_info.orig_dtype)
                        cur_tensor_infos["inf_dtype"] = str(tensor_info.inf_dtype)
                        if tensor_info.id in v.tensor_id_to_scale_zp:
                            cur_tensor_infos["scale"] = v.tensor_id_to_scale_zp[tensor_info.id][0].tolist()
                            cur_tensor_infos["zero_point"] = v.tensor_id_to_scale_zp[tensor_info.id][1].tolist()
                    output_tensor_infos.append(cur_tensor_infos)
                info["output_tensor_infos"] = output_tensor_infos
                # qconfig
                info["activation_observer"] = _get_observer_setting(op_info.qconfig.activation())
                info["weight_observer"] = _get_observer_setting(op_info.qconfig.weight())
                q_op_infos[q_k] = info
            layer_infos["q_op_infos"] = q_op_infos
        if len(v.seen_nonq_op_infos) == 0:
            layer_infos["nonq_op_infos"] = {}
        else:
            nonq_op_infos = OrderedDict()
            for non_q_k, op_info in enumerate(v.seen_nonq_op_infos):
                info = OrderedDict()
                info["op_type"] = op_info.type
                info["fqn"] = op_info.fqn
                input_tensor_infos = []
                for tensor_info in op_info.input_tensor_infos:
                     cur_tensor_infos = {}
                     if tensor_info is not None:
                        cur_tensor_infos["id"] = tensor_info.id
                        cur_tensor_infos["orig_dtype"] = str(tensor_info.orig_dtype)
                        cur_tensor_infos["inf_dtype"] = str(tensor_info.inf_dtype)
                     input_tensor_infos.append(cur_tensor_infos)
                info["input_tensor_infos"] = input_tensor_infos
                output_tensor_infos = []
                for tensor_info in op_info.output_tensor_infos:
                     cur_tensor_infos = {}
                     if tensor_info is not None:
                        cur_tensor_infos["id"] = tensor_info.id
                        cur_tensor_infos["orig_dtype"] = str(tensor_info.orig_dtype)
                        cur_tensor_infos["inf_dtype"] = str(tensor_info.inf_dtype)
                     output_tensor_infos.append(cur_tensor_infos)
                info["output_tensor_infos"] = output_tensor_infos
                nonq_op_infos[non_q_k] = info
            layer_infos["nonq_op_infos"] = nonq_op_infos
        layer_output_infos = []
        for tensor_info in v.output_qtensor_infos:
            cur_tensor_infos = {}
            if tensor_info is not None:
                cur_tensor_infos["id"] = tensor_info.id
                cur_tensor_infos["orig_dtype"] = str(tensor_info.orig_dtype)
                cur_tensor_infos["inf_dtype"] = str(tensor_info.inf_dtype)
                if tensor_info.id in v.tensor_id_to_scale_zp:
                    cur_tensor_infos["scale"] = v.tensor_id_to_scale_zp[tensor_info.id][0].tolist()
                    cur_tensor_infos["zero_point"] = v.tensor_id_to_scale_zp[tensor_info.id][1].tolist()
            layer_output_infos.append(cur_tensor_infos)
        layer_infos["layer_output_infos"] = layer_output_infos
        quant_state_dict[k] = layer_infos
    # save qparms as json file
    if configure_file is not None:
        with open(configure_file, 'w') as fp:
            json.dump(quant_state_dict, fp, indent = 4)

def load_qconf_summary_to_model(model, qconf_summary):
    """
    This function is about load the user given configure to origin model.
    """
    with open(qconf_summary, 'r') as f:
        quant_state_dict = json.load(f)
    quant_state_map = model._fqn_to_auto_quant_state_map
    for k, v in quant_state_map.items():
        layer_info = quant_state_dict[k]
        user_q_op_infos = layer_info["q_op_infos"]
        for i, q_op_info in user_q_op_infos.items():
            fqn = q_op_info["fqn"]
            cur_fqn = v.idx_to_seen_q_op_infos[int(i)].fqn
            assert int(i) in v.idx_to_seen_q_op_infos and \
                    cur_fqn == fqn, "Loaded quantizable op info doesn't match the current model quantizable op info"
            input_tensor_infos = []
            input_force_dtype_infos = []
            for tensor_info in q_op_info["input_tensor_infos"]:
                if len(tensor_info) > 0:
                    input_tensor_infos.append(QTensorInfo(tensor_info["id"], dtype_dict[tensor_info["orig_dtype"]], dtype_dict[tensor_info["inf_dtype"]]))
                    input_force_dtype_infos.append(dtype_dict[tensor_info["force_dtype"]])
                    if "scale" in tensor_info:
                        scale = torch.FloatTensor(tensor_info["scale"])
                        zp = torch.LongTensor(tensor_info["zero_point"])
                        v.tensor_id_to_scale_zp[tensor_info["id"]] = (scale, zp)
                else:
                    input_tensor_infos.append(None)
                    input_force_dtype_infos.append(None)
            weight_tensor_infos = []
            weight_idx = 0
            for tensor_info in q_op_info["weight_tensor_infos"]:
                if len(tensor_info) > 0:
                    weight_tensor_infos.append(QTensorInfo(weight_idx, dtype_dict[tensor_info["orig_dtype"]], dtype_dict[tensor_info["inf_dtype"]]))
                    if "scale" in tensor_info:
                        scale = torch.FloatTensor(tensor_info["scale"])
                        zp = torch.LongTensor(tensor_info["zero_point"])
                        v.weight_tensor_id_to_scale_zp[str(i) + "_" + str(weight_idx)] = (scale, zp)
                    weight_idx += 1
                else:
                    weight_tensor_infos.append(None)
            output_tensor_infos = []
            insert_fake_quant_after_outputs = []
            for tensor_info in q_op_info["output_tensor_infos"]:
                if len(tensor_info) > 0:
                    output_tensor_infos.append(QTensorInfo(tensor_info["id"], dtype_dict[tensor_info["orig_dtype"]], dtype_dict[tensor_info["inf_dtype"]]))
                    insert_fake_quant_after_outputs.append(False)
                    if "scale" in tensor_info:
                        scale = torch.FloatTensor(tensor_info["scale"])
                        zp = torch.LongTensor(tensor_info["zero_point"])
                        v.tensor_id_to_scale_zp[tensor_info["id"]] = (scale, zp)
                else:
                    output_tensor_infos.append(None)
            activation_observer = q_op_info["activation_observer"]
            weight_observer = q_op_info["weight_observer"]
            qconfig = QConfig(activation=_create_observer(activation_observer), weight = _create_observer(weight_observer))
            # overide the cur model's info
            v.idx_to_seen_q_op_infos[int(i)].input_tensor_infos = input_tensor_infos
            v.idx_to_seen_q_op_infos[int(i)].input_tensor_force_inf_dtype = input_force_dtype_infos
            v.idx_to_seen_q_op_infos[int(i)].output_tensor_infos = output_tensor_infos
            v.idx_to_seen_q_op_infos[int(i)].insert_fake_quant_after_outputs = insert_fake_quant_after_outputs
            v.idx_to_seen_q_op_infos[int(i)].weight_tensor_infos = weight_tensor_infos
            v.idx_to_seen_q_op_infos[int(i)].qconfig = qconfig

        user_nonq_op_infos = layer_info["nonq_op_infos"]
        #v.seen_nonq_op_infos.clear()
        idx = 0
        for _, nonq_op_info in user_nonq_op_infos.items():
            fqn = nonq_op_info["fqn"]
            cur_fqn = v.seen_nonq_op_infos[idx].fqn
            assert cur_fqn == fqn, "Loaded none-quantizable op info doesn't match the current model none-quantizable op info"
            input_tensor_infos = []
            for tensor_info in nonq_op_info["input_tensor_infos"]:
                if len(tensor_info) > 0:
                    input_tensor_infos.append(QTensorInfo(tensor_info["id"], dtype_dict[tensor_info["orig_dtype"]], dtype_dict[tensor_info["inf_dtype"]]))
                else:
                    input_tensor_infos.append(None)
            output_tensor_infos = []
            for tensor_info in nonq_op_info["output_tensor_infos"]:
                if len(tensor_info) > 0:
                    output_tensor_infos.append(QTensorInfo(tensor_info["id"], dtype_dict[tensor_info["orig_dtype"]], dtype_dict[tensor_info["inf_dtype"]]))
                else:
                    output_tensor_infos.append(None)
            v.seen_nonq_op_infos[idx].input_tensor_infos = input_tensor_infos
            v.seen_nonq_op_infos[idx].output_tensor_infos = output_tensor_infos
            idx += 1

        layer_output_info = []
        for tensor_info in layer_info["layer_output_infos"]:
            if len(tensor_info) > 0:
                layer_output_info.append(QTensorInfo(tensor_info["id"], dtype_dict[tensor_info["orig_dtype"]], dtype_dict[tensor_info["inf_dtype"]]))
                if "scale" in tensor_info:
                    scale = torch.FloatTensor(tensor_info["scale"])
                    zp = torch.LongTensor(tensor_info["zero_point"])
                    v.tensor_id_to_scale_zp[tensor_info["id"]] = (scale, zp)
            else:
                layer_output_info.append(None)
        v.output_qtensor_infos = layer_output_info
    # insert observer according to user's setting.
    for _, v in model.named_modules():
        if hasattr(v, '_auto_quant_state'):
            v._auto_quant_state.tensor_id_to_observer.clear()
            v._auto_quant_state.weight_tensor_id_to_observer.clear()
            v._auto_quant_state.insert_observers(v)
    
    # update insert_fake_quant_after_output after load user setting, which for avoiding redundant fake quant.
    nodes = convert_quant_state_map_to_nodes(quant_state_map)
    set_node_output_quantized(nodes)

def _lstm_forward(module, input, hx, weights):
    r"""
    LSTM forward function.
    """
    orig_input = input
    # xxx: isinstance check needs to be in conditional for TorchScript to compile
    # batch_sizes = None
    if isinstance(orig_input, torch.nn.utils.rnn.PackedSequence):
        input, batch_sizes, sorted_indices, unsorted_indices = input
        max_batch_size = batch_sizes[0]
        max_batch_size = int(max_batch_size)
    else:
        batch_sizes = None
        is_batched = input.dim() == 3
        batch_dim = 0 if module.batch_first else 1
        if not is_batched:
            input = input.unsqueeze(batch_dim)
        max_batch_size = input.size(0) if module.batch_first else input.size(1)
        sorted_indices = None
        unsorted_indices = None
    if hx is None:
        num_directions = 2 if module.bidirectional else 1
        real_hidden_size = module.proj_size if module.proj_size > 0 else module.hidden_size
        h_zeros = torch.zeros(module.num_layers * num_directions, max_batch_size, real_hidden_size, \
            dtype=input.dtype, device=input.device)
        c_zeros = torch.zeros(module.num_layers * num_directions, max_batch_size, module.hidden_size, \
            dtype=input.dtype, device=input.device)
        hx = (h_zeros, c_zeros)
    else:
        if batch_sizes is None:  # If not PackedSequence input.
            if is_batched:
                if (hx[0].dim() != 3 or hx[1].dim() != 3):
                    msg = ("For batched 3-D input, hx and cx should "
                            f"also be 3-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors")
                    raise RuntimeError(msg)
            else:
                if hx[0].dim() != 2 or hx[1].dim() != 2:
                    msg = ("For unbatched 2-D input, hx and cx should "
                               f"also be 2-D but got ({hx[0].dim()}-D, {hx[1].dim()}-D) tensors")
                    raise RuntimeError(msg)
                hx = (hx[0].unsqueeze(1), hx[1].unsqueeze(1))

        # Each batch of the hidden state should match the input sequence that
        # the user believes he/she is passing in.
        hx = module.permute_hidden(hx, sorted_indices)

    module.check_forward_args(input, hx, batch_sizes)
    if batch_sizes is None:
        result = _VF.lstm(input, hx, weights, module.bias, module.num_layers, \
            module.dropout, module.training, module.bidirectional, module.batch_first)
    else:
        result = _VF.lstm(input, batch_sizes, hx, weights, module.bias, \
            module.num_layers, module.dropout, module.training, module.bidirectional)
    output = result[0]
    hidden = result[1:]
    # xxx: isinstance check needs to be in conditional for TorchScript to compile
    if isinstance(orig_input, torch.nn.utils.rnn.PackedSequence):
        output_packed = torch.nn.utils.rnn.PackedSequence(output, batch_sizes, sorted_indices, unsorted_indices)
        return output_packed, module.permute_hidden(hidden, unsorted_indices)
    else:
        if not is_batched:
            output = output.squeeze(batch_dim)
            hidden = (hidden[0].squeeze(1), hidden[1].squeeze(1))
        return output, module.permute_hidden(hidden, unsorted_indices)

def module_call_to_function_call(module, args, weights):
    r"""
    This function is a help function which replace nn.module call to funtion call, which implement
    the nn.module's forward function.
    """
    if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Conv3d):
        output = module._conv_forward(args[0], weights[0], module.bias)
    elif isinstance(module, torch.nn.Linear):
        output = F.linear(args[0], weights[0], module.bias)
    elif isinstance(module, torch.nn.EmbeddingBag):
        output = F.embedding_bag(args[0], weights[0], args[1], module.max_norm, \
            module.norm_type, module.scale_grad_by_freq, module.mode, module.sparse,
            args[2] if len(args) == 3 else None, module.include_last_offset, module.padding_idx)
    elif isinstance(module, torch.nn.ConvTranspose2d) or isinstance(module, torch.nn.ConvTranspose3d):
        if module.padding_mode != 'zeros':
            raise ValueError('Only `zeros` padding mode is supported for ConvTranspose2d')
        assert isinstance(module.padding, tuple)
        # One cannot replace List by Tuple or Sequence in "_output_padding" because
        # TorchScript does not support `Sequence[T]` or `Tuple[T, ...]`.
        output_size = args[1] if len(args) == 2 else None
        # master code
        '''
        num_spatial_dims = 2 if isinstance(module, torch.nn.ConvTranspose2d) else 3
        output_padding = module._output_padding(args[0], output_size,
                        module.stride, module.padding, module.kernel_size,
                        num_spatial_dims, module.dilation)
        '''
        output_padding = module._output_padding(args[0], output_size, module.stride, module.padding, module.kernel_size,  module.dilation)
        #output_padding = module._output_padding(*arg_to)
        if isinstance(module, torch.nn.ConvTranspose2d):
            output = F.conv_transpose2d(
                args[0], weights[0], module.bias, module.stride, module.padding,
                output_padding, module.groups, module.dilation)
        else:
            output = F.conv_transpose3d(
                args[0], weights[0], module.bias, module.stride, module.padding,
                output_padding, module.groups, module.dilation)
    elif isinstance(module, torch.nn.LSTM):
        output = _lstm_forward(module, args[0], args[1] if len(args) == 2 else None, weights)
    return output
