import torch
import operator
import copy
import torch.nn as nn
import torch.nn.qat as nnqat
import torch.nn.intrinsic.qat as nniqat

import mqbench.nn.intrinsic.qat as qnniqat
import mqbench.nn.qat as qnnqat
import mediatools.simulator_nn as simnn
import mqbench.fake_quantize as mquant
import mediatools.fake_quantize as mtquant
from torch.fx import GraphModule
# from torch.quantization.fx.utils import _parent_name
from mediatools.utils.registry import register_simulator_convert_function
from mediatools.simulator_nn.utils import (
    process_fixpoint_weight_bias,
    flatten_args,
    merge_quant_params,
    merge_elementwise_quant_params,
    to_quantizer_params,
    merge_add_quant_params
)

# copy from torch1.8.1 torch.quantization.fx.utils
# turn foo.bar -> ['foo', 'bar']
def _parent_name(target):
    r = target.rsplit('.', 1)
    if len(r) == 1:
        return '', r[0]
    else:
        return r[0], r[1]

_exclude_module_type = (
            torch.nn.PixelShuffle,
            torch.nn.PixelUnshuffle,
            torch.nn.Identity,
            torch.nn.MaxPool2d,
        )

_shared_func_type = ()

@register_simulator_convert_function(qnnqat.Conv2d)
@register_simulator_convert_function(nnqat.Conv2d)
def convert_qnnqat_conv2d(model, convert_node, cached_quant_info, fakequantizer, fixpoint_params):
    modules = dict(model.named_modules())
    convert_module = modules[convert_node.target]
    cached_quant_info[convert_node.target] = to_quantizer_params(convert_module.weight_fake_quant)
    if convert_node.all_input_nodes[0].op == "call_module" and \
        isinstance(modules[convert_node.all_input_nodes[0].target], _exclude_module_type):
        input_module = modules[convert_node.all_input_nodes[0].all_input_nodes[0].target]
    else:
        input_module = modules[convert_node.all_input_nodes[0].target]
    assert isinstance(input_module, torch.quantization.FakeQuantizeBase)

    output_node = list(convert_node.users)[0]
    output_module = modules[output_node.target]
    assert isinstance(output_module, torch.quantization.FakeQuantizeBase)

    cached_quant_info[convert_node.all_input_nodes[0].target] = to_quantizer_params(input_module)
    cached_quant_info[output_node.target] = to_quantizer_params(output_module)
    
    conv = simnn.Conv2d(convert_module.in_channels, convert_module.out_channels, convert_module.kernel_size, 
                        convert_module.stride, convert_module.padding, convert_module.dilation,
                        convert_module.groups, convert_module.bias is not None, convert_module.padding_mode,
                        padding_value=to_quantizer_params(input_module)["zero_point"], weight_bitwidth=convert_module.weight_fake_quant.bitwidth.item())
    conv.weight = convert_module.weight
    if convert_module.bias is not None:
        conv.bias = convert_module.bias
    convert_conv = process_fixpoint_weight_bias(conv.eval(), convert_module.weight_fake_quant, input_module, output_module, **fixpoint_params.get('bias_params', {}))
    conv_parent_name, conv_name = _parent_name(convert_node.target)
    setattr(modules[conv_parent_name], conv_name, convert_conv)


@register_simulator_convert_function(nniqat.ConvReLU2d)
@register_simulator_convert_function(qnniqat.ConvReLU2d)
def convert_qnnqat_conv2dReLU(model, convert_node, cached_quant_info, fakequantizer, fixpoint_params):
    modules = dict(model.named_modules())
    convert_module = modules[convert_node.target]
    cached_quant_info[convert_node.target] = to_quantizer_params(convert_module.weight_fake_quant)
    if convert_node.all_input_nodes[0].op == "call_module" and \
        isinstance(modules[convert_node.all_input_nodes[0].target], _exclude_module_type):
        input_module = modules[convert_node.all_input_nodes[0].all_input_nodes[0].target]
    else:
        input_module = modules[convert_node.all_input_nodes[0].target]
    assert isinstance(input_module, torch.quantization.FakeQuantizeBase)

    output_node = list(convert_node.users)[0]
    output_module = modules[output_node.target]
    assert isinstance(output_module, torch.quantization.FakeQuantizeBase)
    cached_quant_info[output_node.target] = to_quantizer_params(output_module)

    cached_quant_info[convert_node.all_input_nodes[0].target] = to_quantizer_params(input_module)
    
    conv = simnn.Conv2d(convert_module.in_channels, convert_module.out_channels, convert_module.kernel_size, 
                        convert_module.stride, convert_module.padding, convert_module.dilation,
                        convert_module.groups, convert_module.bias is not None, convert_module.padding_mode,
                        padding_value=to_quantizer_params(input_module)["zero_point"], is_relu=1,
                        weight_bitwidth=convert_module.weight_fake_quant.bitwidth.item())
    conv.weight = convert_module.weight
    if convert_module.bias is not None:
        conv.bias = convert_module.bias
    convert_conv = process_fixpoint_weight_bias(conv.eval(), convert_module.weight_fake_quant, input_module, output_module, **fixpoint_params.get('bias_params', {}))
    # convert_conv = simnn.ConvReLU2d(convert_conv, torch.nn.ReLU(inplace=True))
    conv_parent_name, conv_name = _parent_name(convert_node.target)
    setattr(modules[conv_parent_name], conv_name, convert_conv)


@register_simulator_convert_function(qnnqat.ConvTranspose2d)
def convert_qnnqat_ConvTranspose2d(model, convert_node, cached_quant_info, fakequantizer, fixpoint_params):
    modules = dict(model.named_modules())
    convert_module = modules[convert_node.target]
    cached_quant_info[convert_node.target] = to_quantizer_params(convert_module.weight_fake_quant)
    if convert_node.all_input_nodes[0].op == "call_module" and \
        isinstance(modules[convert_node.all_input_nodes[0].target], _exclude_module_type):
        input_module = modules[convert_node.all_input_nodes[0].all_input_nodes[0].target]
    else:
        input_module = modules[convert_node.all_input_nodes[0].target]
    cached_quant_info[convert_node.all_input_nodes[0].target] = to_quantizer_params(input_module)
    assert isinstance(input_module, torch.quantization.FakeQuantizeBase)
    deconv = simnn.ConvTranspose2d(convert_module.in_channels, convert_module.out_channels, convert_module.kernel_size, 
                                  stride=convert_module.stride, padding=convert_module.padding, output_padding=convert_module.output_padding,
                                  groups=convert_module.groups, bias=convert_module.bias is not None, 
                                  dilation=convert_module.dilation,
                                  padding_mode=convert_module.padding_mode,
                                  weight_bitwidth=convert_module.weight_fake_quant.bitwidth.item()) 
    deconv.weight = convert_module.weight
    if convert_module.bias is not None:
        deconv.bias = convert_module.bias
    convert_deconv = process_fixpoint_weight_bias(deconv.eval(), convert_module.weight_fake_quant, input_module, **fixpoint_params.get('bias_params', {}))
    deconv_parent_name, deconv_name = _parent_name(convert_node.target)
    setattr(modules[deconv_parent_name], deconv_name, convert_deconv)


@register_simulator_convert_function(qnniqat.ConvTransposeReLU2d)
def convert_qnnqat_ConvTransposeReLU2d(model, convert_node, cached_quant_info, fakequantizer, fixpoint_params):
    modules = dict(model.named_modules())
    convert_module = modules[convert_node.target]
    cached_quant_info[convert_node.target] = to_quantizer_params(convert_module.weight_fake_quant)
    if convert_node.all_input_nodes[0].op == "call_module" and \
        isinstance(modules[convert_node.all_input_nodes[0].target], _exclude_module_type):
        input_module = modules[convert_node.all_input_nodes[0].all_input_nodes[0].target]
    else:
        input_module = modules[convert_node.all_input_nodes[0].target]
    cached_quant_info[convert_node.all_input_nodes[0].target] = to_quantizer_params(input_module)
    assert isinstance(input_module, torch.quantization.FakeQuantizeBase)
    deconv = simnn.ConvTranspose2d(convert_module.in_channels, convert_module.out_channels, convert_module.kernel_size, 
                                  stride=convert_module.stride, padding=convert_module.padding, output_padding=convert_module.output_padding,
                                  groups=convert_module.groups, bias=convert_module.bias is not None, 
                                  dilation=convert_module.dilation, padding_mode=convert_module.padding_mode, 
                                  padding_value=to_quantizer_params(input_module)["zero_point"], is_relu=1,
                                  weight_bitwidth=convert_module.weight_fake_quant.bitwidth.item()) 
    deconv.weight = convert_module.weight
    if convert_module.bias is not None:
        deconv.bias = convert_module.bias
    convert_deconv = process_fixpoint_weight_bias(deconv.eval(), convert_module.weight_fake_quant, input_module, **fixpoint_params.get('bias_params', {}))
    # convert_deconv = simnn.ConvTransposeReLU2d(convert_deconv, torch.nn.ReLU(inplace=True))
    deconv_parent_name, deconv_name = _parent_name(convert_node.target)
    setattr(modules[deconv_parent_name], deconv_name, convert_deconv)

@register_simulator_convert_function(qnnqat.Linear)
@register_simulator_convert_function(nnqat.Linear)
def convert_nnqat_linear(model, convert_node, cached_quant_info, fakequantizer, fixpoint_params):
    modules = dict(model.named_modules())
    convert_module = modules[convert_node.target]
    cached_quant_info[convert_node.target] = to_quantizer_params(convert_module.weight_fake_quant)
    if (convert_node.all_input_nodes[0].op == "call_module" and isinstance(modules[convert_node.all_input_nodes[0].target], _exclude_module_type)) or \
        (convert_node.all_input_nodes[0].op == "call_function" and convert_node.all_input_nodes[0].target in fakequantizer._exclude_func_type):
        input_module = modules[convert_node.all_input_nodes[0].all_input_nodes[0].target]
    else:
        input_module = modules[convert_node.all_input_nodes[0].target]
    cached_quant_info[convert_node.all_input_nodes[0].target] = to_quantizer_params(input_module)
    assert isinstance(input_module, torch.quantization.FakeQuantizeBase)
    linear = simnn.Linear(convert_module.in_features, convert_module.out_features, convert_module.bias is not None)
    linear.weight = convert_module.weight
    if convert_module.bias is not None:
        linear.bias = convert_module.bias
    convert_linear = process_fixpoint_weight_bias(linear.eval(), convert_module.weight_fake_quant, input_module, **fixpoint_params.get('bias_params', {}))
    conv_parent_name, conv_name = _parent_name(convert_node.target)
    setattr(modules[conv_parent_name], conv_name, convert_linear)


@register_simulator_convert_function(mquant.FixedFakeQuantize)
@register_simulator_convert_function(mquant.LearnableFakeQuantize)
@register_simulator_convert_function(mtquant.FixedFakeQuantize)
@register_simulator_convert_function(mtquant.LearnableFakeQuantize)
@register_simulator_convert_function(mtquant.PseudoFakeQuantize)
def convert_fakequantizer_quantizer(model: GraphModule, 
                                   convert_node, 
                                   cached_quant_info, 
                                   output_node_info, 
                                   fakequantizer,
                                   fixpoint_only,
                                   fixpoint_params,
                                   **kwargs):
    modules = dict(model.named_modules())
    convert_module = modules[convert_node.target]
    cached_quant_info[convert_node.target] = to_quantizer_params(convert_module)
    # len(convert_node.all_input_nodes)==1, cause we process act_fakequantize here 
    if len(convert_node.all_input_nodes[0].all_input_nodes) == 0 or \
        (convert_node.all_input_nodes[0].target == operator.getitem and len(convert_node.all_input_nodes[0].all_input_nodes[0].all_input_nodes) == 0):
        # process as input node 
        if fixpoint_only:
            quant_params = to_quantizer_params(convert_module)
            quant_params.update(fixpoint_params.get('quantizer_params', {}))
            quantizer = simnn.QuantizerInt(**quant_params)
        else:
            quantizer = simnn.QuantizerFloat(**to_quantizer_params(convert_module))
        module_parent_name, module_name = _parent_name(convert_node.target)
        setattr(modules[module_parent_name], module_name, quantizer)
    
        model.recompile()
        model.graph.lint()
    else:
        users_names = [x.target for x in list(convert_node.users.keys())]
        if 'output' in users_names:
            output_node_info[convert_node.target] = convert_node
        if ((convert_node.all_input_nodes[0].op == 'call_function' and convert_node.all_input_nodes[0].target in fakequantizer._passed_func_type) or
            (convert_node.all_input_nodes[0].op == 'call_module' and isinstance(modules[convert_node.all_input_nodes[0].target], fakequantizer._passed_module_type))):
            # process as passed node
            # copy fakequantizer to passed node
            assert convert_node.all_input_nodes[0].target not in cached_quant_info, "passed node should not have a fakequantizer"
            cached_quant_info[convert_node.all_input_nodes[0].target] = cached_quant_info[convert_node.target]
            quantizer = process_mixquantizer_node(model, modules, convert_node.all_input_nodes[0], cached_quant_info, fakequantizer, fixpoint_only, fixpoint_params)
            operator_node=convert_node.all_input_nodes[0].all_input_nodes[0]
        else:
            quantizer = process_mixquantizer_node(model, modules, convert_node, cached_quant_info, fakequantizer, fixpoint_only, fixpoint_params)
            operator_node=convert_node.all_input_nodes[0]
            
    # if isinstance(quantizer, dict):
    #     module_parent_name, module_name = _parent_name(convert_node.target)
    #     setattr(modules[module_parent_name], module_name, quantizer[convert_node])
    #     for node, mixedquantizer in quantizer.items():
    #         if node != convert_node:
    #             new_node_name =  '{}_{}_act_fake_quantizer'.format(node.name, operator_node.name)
    #             setattr(model, new_node_name, mixedquantizer)
    #             with model.graph.inserting_before(operator_node):
    #                 new_node = model.graph.call_module(new_node_name, (node,), {})
    #                 operator_node.replace_input_with(node, new_node)
    # else:
    #     module_parent_name, module_name = _parent_name(convert_node.target)
    #     setattr(modules[module_parent_name], module_name, quantizer)
    
    # model.recompile()
    # model.graph.lint()



def insert_output_quantizer(model: GraphModule, 
                            cached_quant_info, 
                            output_node_info, 
                            fixpoint_only,
                            fixpoint_params,
                            **kwargs):
    # Update modules.
    graph = model.graph
    nodes = list(model.graph.nodes)
    quantizer_name = 'output_quantizer'
    for node_name, out_node in output_node_info.items():
        insert_module_name = "{}_{}".format(out_node.name, quantizer_name)
        if not hasattr(model, insert_module_name):
            if fixpoint_only:
                quant_params = cached_quant_info[out_node.target]
                quant_params.update(fixpoint_params.get('quantizer_params', {}))
                quantizer = simnn.DeQuantizerInt(**quant_params)
            else:
                quantizer = simnn.DeQuantizerFloat(**cached_quant_info[out_node.target])
            setattr(model, insert_module_name, quantizer)
        with graph.inserting_after(out_node):
            insert_node = graph.create_node("call_module", insert_module_name, (), {})
            out_node.replace_all_uses_with(insert_node)
            insert_node.args = (out_node,)
            out_node.users = dict(insert_module_name=None)
            # for _node in nodes:
            #     if _node.target == 'output':
            #         for i, _arg in enumerate(_node.args):
            #             if isinstance(_arg, tuple):
            #                 for j, sub_arg in enumerate(_arg):
            #                     if sub_arg == out_node:
            #                         all_args = list(_node.args)
            #                         _tmp = list(all_args[i])
            #                         _tmp[j] = insert_node
            #                         all_args[i] = tuple(_tmp)
            #                         _node.args = tuple(all_args)
            #             else:
            #                 if _arg == out_node:
            #                     _tmp = list(_node.args)
            #                     _tmp[i] = insert_node
            #                     _node.args = tuple(_tmp)
    model.recompile()
    model.graph.lint()


def process_mixquantizer_node(model, modules, cur_node, cached_quant_info, fakequantizer, fixpoint_only, fixpoint_params):
    input_nodes = cur_node.all_input_nodes[0].all_input_nodes
    convert_node_quant_params = cached_quant_info[cur_node.target]
    if len(input_nodes) > 1:
        # process as elementwise node
        input_quant_params=[]
        for input_node in input_nodes:
            if input_node.target in cached_quant_info:
                input_quant_params.append(cached_quant_info[input_node.target])
            else:
                input_quant_params.append(cached_quant_info[input_node.all_input_nodes[0].target])

        if cur_node.all_input_nodes[0].target == operator.add:
            quantizer_params = merge_add_quant_params(input_quant_params[0], input_quant_params[1], convert_node_quant_params)

            if fixpoint_only:
                quantizer_params.update(fixpoint_params.get('quantizer_params', {}))
                quantizer = simnn.AddQuantizerInt(**quantizer_params)
            else:
                quantizer = simnn.AddQuantizerFloat(**quantizer_params)

            setattr(model, cur_node.all_input_nodes[0].name, quantizer)
            setattr(cur_node.all_input_nodes[0], 'op', 'call_module')
            setattr(cur_node.all_input_nodes[0], 'target', cur_node.all_input_nodes[0].name)
            setattr(model, cur_node.name, nn.Identity())
            model.recompile()
            model.graph.lint()
            cached_quant_info[cur_node.all_input_nodes[0].target] = convert_node_quant_params
        else:
            raise NotImplementedError("Convert module type not supported")
    else:
        if (cur_node.all_input_nodes[0].op == 'call_function' and cur_node.all_input_nodes[0].target in [torch.nn.functional.adaptive_avg_pool2d]) or \
            (cur_node.all_input_nodes[0].op == 'call_module' and isinstance(modules[cur_node.all_input_nodes[0].target], torch.nn.AdaptiveAvgPool2d)):
            # process as elementwise node
            input_node = cur_node.all_input_nodes[0].all_input_nodes[0]
            convert_node_quant_params = cached_quant_info[cur_node.target]
            quantizer_params = merge_elementwise_quant_params(cached_quant_info[input_node.target], convert_node_quant_params)
        else:
            # process as module with weight and bias
            preone_node = cur_node.all_input_nodes[0]
            if (preone_node.op == 'call_module' and isinstance(modules[preone_node.target], _exclude_module_type)) or \
                (preone_node.op == 'call_function' and preone_node.target in fakequantizer._exclude_func_type):
                preone_node = preone_node.all_input_nodes[0]
            pretwo_node = preone_node.all_input_nodes[0]
            if (pretwo_node.op == 'call_module' and isinstance(modules[pretwo_node.target], _exclude_module_type)) or \
                (pretwo_node.op == 'call_function' and pretwo_node.target in fakequantizer._exclude_func_type):
                pretwo_node = pretwo_node.all_input_nodes[0]
            preone_quant_params = cached_quant_info[preone_node.target]
            pretwo_quant_params = cached_quant_info[pretwo_node.target]
            convert_node_quant_params = cached_quant_info[cur_node.target]
            quantizer_params = merge_quant_params(pretwo_quant_params, preone_quant_params, convert_node_quant_params)
        # create quantizer from quantizer_params
        if fixpoint_only:
            quantizer_params.update(fixpoint_params.get('quantizer_params', {}))
            quantizer = simnn.MixedQuantizerInt(**quantizer_params)
            pre_module = modules[cur_node.all_input_nodes[0].target]
            if hasattr(pre_module, 'is_relu') and pre_module.is_relu == 1:
                quantizer.is_relu = True
                quantizer.clip_value_low = torch.nn.Parameter(convert_node_quant_params['zero_point'].round().to(torch.float32))
        else:
            quantizer = simnn.MixedQuantizerFloat(**quantizer_params)

        module_parent_name, module_name = _parent_name(cur_node.target)
        setattr(modules[module_parent_name], module_name, quantizer)
    
        model.recompile()
        model.graph.lint()
