import os
import torch
import torch.nn.intrinsic as nni
import torch.nn.intrinsic.qat as nniqat
import torch.distributed as dist
import os.path as osp
import torch._C._onnx as _C_onnx
from torch.fx import GraphModule
import mqbench.fusion_method 
from mediatools.simulator_nn import *
from mqbench.fusion_method import convert_nniqat_convbn, convert_qnniqat_convbn
import mqbench.nn.intrinsic as qnni
import mqbench.nn.intrinsic.qat as qnniqat
from mqbench.utils import deepcopy_graphmodule
from mqbench.utils.logger import logger
from mqbench.utils.registry import (
    register_convert_function,
    FUSED_MODULE_CONVERT_FUNCTION,
    DEFAULT_MODEL_QUANTIZER
)
from torch.quantization import FakeQuantizeBase
from mqbench.quantization.default_bias_fake_quant import bias_fake_quantizer
from mqbench.deploy import remove_fakequantize_and_collect_params
from mqbench.custom_symbolic_opset import *

from mediatools.deploy.deploy_linear import remove_fakequantize_and_collect_params
from mediatools.prepare_by_platform import BackendType
from mediatools.convert_method import insert_output_quantizer, _parent_name
from mediatools.utils.registry import (
    register_deploy_mnn_function,
    register_simulator_function,
    register_deploy_onnx_function,
    BACKEND_SIMULATOR_FUNCTION,
    SIMULATOR_MODULE_CONVERT_FUNCTION,
    BACKEND_DEPLOY_MNN_FUNCTION,
    BACKEND_DEPLOY_ONNX_FUNCTION
)


@register_convert_function(nniqat.ConvBnReLU2d)
def convert_nniqat_convbnrelu(model, fused_node):
    convert_nniqat_convbn(model, fused_node)
    modules = dict(model.named_modules())
    fused_module = modules[fused_node.target]
    conv = torch.nn.Conv2d(fused_module.in_channels, fused_module.out_channels, fused_module.kernel_size, 
                      fused_module.stride, fused_module.padding, fused_module.dilation,
                      fused_module.groups, fused_module.bias is not None, fused_module.padding_mode)
    conv.weight = fused_module.weight
    if fused_module.bias is not None:
        conv.bias = fused_module.bias
    # We need nn.qat.conv here to export weight quantize node.
    conv.qconfig = fused_module.qconfig
    fused_conv = nni.ConvReLU2d(conv, torch.nn.ReLU(inplace=True))
    fused_conv.qconfig = fused_module.qconfig
    fused_conv = nniqat.ConvReLU2d.from_float(fused_conv)
    # Attach weight fake quantize params.
    fused_conv.weight_fake_quant = fused_module.weight_fake_quant
    conv_parent_name, conv_name = _parent_name(fused_node.target)
    setattr(modules[conv_parent_name], conv_name, fused_conv)


@register_convert_function(qnniqat.ConvBnReLU2d)
def convert_qnniqat_convbnrelu(model, fused_node):
    convert_qnniqat_convbn(model, fused_node)
    modules = dict(model.named_modules())
    fused_module = modules[fused_node.target]
    # Create a Conv2d from FusedModule.
    conv = torch.nn.Conv2d(fused_module.in_channels, fused_module.out_channels, fused_module.kernel_size, 
                           fused_module.stride, fused_module.padding, fused_module.dilation,
                           fused_module.groups, fused_module.bias is not None, fused_module.padding_mode)
    conv.weight = fused_module.weight
    if fused_module.bias is not None:
        conv.bias = fused_module.bias
    conv.qconfig = fused_module.qconfig
    fused_conv = nni.ConvReLU2d(conv, torch.nn.ReLU(inplace=True))
    # We need nn.qat.conv here to export weight quantize node.
    fused_conv.qconfig = fused_module.qconfig
    fused_conv = qnniqat.ConvReLU2d.from_float(fused_conv)
    # Attach weight fake quantize params.
    fused_conv.weight_fake_quant = fused_module.weight_fake_quant
    if hasattr(fused_module, 'bias_fake_quant'):
        fused_conv.bias_fake_quant = fused_module.bias_fake_quant
    else:
        fused_conv.bias_fake_quant = bias_fake_quantizer()
        fused_conv.bias_fake_quant.set_quant_type('param')
    conv_parent_name, conv_name = _parent_name(fused_node.target)
    setattr(modules[conv_parent_name], conv_name, fused_conv)


@register_simulator_function(BackendType.AOFullA8W4)
@register_simulator_function(BackendType.AOFullA8W6)
@register_simulator_function(BackendType.PPLW8A8SYM)
@register_deploy_mnn_function(BackendType.AOFullA8W4)
@register_deploy_mnn_function(BackendType.AOFullA8W6)
@register_deploy_mnn_function(BackendType.PPLW8A8SYM)
@register_deploy_onnx_function(BackendType.PPLW8A8)
@register_deploy_onnx_function(BackendType.PPLW8A8SYM)
def convert_merge_bn(model: GraphModule, **kwargs):
    logger.info("Folding BN into weight and bias.")
    nodes = list(model.graph.nodes)
    modules = dict(model.named_modules())
    for node in nodes:
        if node.op == 'call_module':
            if type(modules[node.target]) in FUSED_MODULE_CONVERT_FUNCTION:
                FUSED_MODULE_CONVERT_FUNCTION[type(modules[node.target])](model, node)


@register_simulator_function(BackendType.AOFullA8W4)
@register_deploy_mnn_function(BackendType.AOFullA8W4)
@register_simulator_function(BackendType.AOFullA8W6)
@register_simulator_function(BackendType.PPLW8A8SYM)
@register_deploy_mnn_function(BackendType.AOFullA8W6)
@register_deploy_mnn_function(BackendType.PPLW8A8SYM)
def convert_fixpoint(model: GraphModule, **kwargs):
    logger.info("Convert qat model to fix point inference model")
    fixpoint_only = kwargs.get('fixpoint_only', False)
    fixpoint_params = kwargs.get('fixpoint_params', {})
    nodes = list(model.graph.nodes)
    modules = dict(model.named_modules())
    cached_quant_info = {}
    fakequantizer = DEFAULT_MODEL_QUANTIZER[kwargs['backend_type']]({}, {})
    with torch.no_grad():
        # process module with weights and bias
        for node in nodes:
            if node.op == 'call_module' and not isinstance(modules[node.target], FakeQuantizeBase):
                if type(modules[node.target]) in SIMULATOR_MODULE_CONVERT_FUNCTION:
                    SIMULATOR_MODULE_CONVERT_FUNCTION[type(modules[node.target])](model, node, cached_quant_info, fakequantizer, fixpoint_params)
        output_node_info = {}
        # process fake quant node
        for node in nodes:
            if node.op == 'call_module' and isinstance(modules[node.target], FakeQuantizeBase):
                if type(modules[node.target]) in SIMULATOR_MODULE_CONVERT_FUNCTION:
                    SIMULATOR_MODULE_CONVERT_FUNCTION[type(modules[node.target])](
                        model, node, cached_quant_info, output_node_info, fakequantizer, 
                        fixpoint_only=fixpoint_only,
                        fixpoint_params=fixpoint_params
                    )
        # process output node.
        insert_output_quantizer(model, cached_quant_info, output_node_info,
                                fixpoint_only, fixpoint_params)


@register_deploy_mnn_function(BackendType.AOFullA8W4)
@register_deploy_mnn_function(BackendType.PPLW8A8SYM)
@register_deploy_onnx_function(BackendType.PPLW8A8)
@register_deploy_onnx_function(BackendType.PPLW8A8SYM)
def deploy_onnx(model: GraphModule, input_shape_dict, dummy_input, onnx_model_path, **kwargs):
    if not dist.is_initialized() or dist.get_rank() == 0:
        logger.info("Export to onnx.")
    device = kwargs.get('device', next(model.parameters()).device)
    dir_path = '/'.join(onnx_model_path.split('/')[:-1])
    os.makedirs(dir_path, exist_ok=True)
    model.to(device)
    input_names = None
    if dummy_input is None:
        dummy_input = {name: torch.rand(shape).to(device) for name, shape in input_shape_dict.items()}
        input_names = list(dummy_input.keys())

    with torch.no_grad():
        torch.onnx.export(model, dummy_input, onnx_model_path,
                          input_names=input_names,
                          opset_version=11,
                          do_constant_folding=True,
                          custom_opsets={'' : 11},
                          enable_onnx_checker=False,
                          verbose=True)


@register_deploy_mnn_function(BackendType.AOFullA8W4)
@register_deploy_mnn_function(BackendType.PPLW8A8SYM)
def deploy_MNN(model: GraphModule, onnx_model_path=None, mnn_model_name=None, **kwargs):
    try:
        from mediatools.deploy.deploy_flatbuffer import dump_to_flatbuffer
    except:
        if not dist.is_initialized() or dist.get_rank() == 0:
            logger.info("don't have exporter for mnn")
        return 0
    logger.info("Extract qparams for mnn.")
    dump_to_flatbuffer(onnx_model_path, mnn_model_name, **kwargs)


@register_deploy_onnx_function(BackendType.PPLW8A8)
@register_deploy_onnx_function(BackendType.PPLW8A8SYM)
def deploy_qparams_pplw8a16(model: GraphModule, onnx_model_path, model_name, **kwargs):
    logger.info("Extract qparams for PPLW8A16.")
    remove_fakequantize_and_collect_params(onnx_model_path, model_name, backend='ppl')


def convert_fixpoint_simulator(model: GraphModule, backend_type: BackendType,
                               fixpoint_only=True, **kwargs):
    r"""Convert model to fixpoint inference model.

    Args:
        model (GraphModule): GraphModule prepared qat module.
        backend_type (BackendType): specific which backend should be converted to.
    """
    simulator_model = deepcopy_graphmodule(model)
    for convert_function in BACKEND_SIMULATOR_FUNCTION[backend_type]:
        convert_function(simulator_model, backend_type=backend_type, fixpoint_only=fixpoint_only, **kwargs)
    return simulator_model


def export_mnn(model, backend_type: BackendType,
                   input_shape_dict=None, dummy_input=None, output_path='./',
                   model_name='onnx_fixpoint_qmodel', device='cuda', **kwargs):
    device = model.device if device is None else device
    assert isinstance(model, GraphModule), 'converted model must have type GraphModule'
    kwargs.update({
        'backend_type': backend_type,
        'input_shape_dict': input_shape_dict,
        'dummy_input': dummy_input,
        'output_path': output_path,
        'model_name': model_name,
        'onnx_model_path': osp.join(output_path, '{}.onnx'.format(model_name)),
        'mnn_model_name': osp.join(output_path, '{}.json'.format(model_name)),
        'fixpoint_only': True,
        'device': device,
    })
    deploy_model = deepcopy_graphmodule(model)
    for convert_function in BACKEND_DEPLOY_MNN_FUNCTION[backend_type]:
        deploy_model.to(device)
        convert_function(deploy_model, **kwargs)

    return deploy_model


def convert_onnx_deploy(model: GraphModule, backend_type: BackendType,
                   input_shape_dict=None, dummy_input=None, output_path='./',
                   model_name='mqbench_qmodel', deploy_to_qlinear=False, device='cuda', **extra_kwargs):
    r"""Convert model to onnx model and quantization params depends on backend.

    Args:
        model (GraphModule): GraphModule prepared qat module.
        backend_type (BackendType): specific which backend should be converted to.
        input_shape_dict (dict): keys are model input name(should be forward function
                                 params name, values are list of tensor dims)
        output_path (str, optional): path to save convert results. Defaults to './'.
        model_name (str, optional): name of converted onnx model. Defaults to 'mqbench_qmodel'.

    >>> note on input_shape_dict:
        example: {'input_0': [1, 3, 224, 224]
                'input_1': [1, 3, 112, 112]
                }
        while forward function signature is like:
                def forward(self, input_0, input_1):
                    pass
    """
    device = model.device if device is None else device
    kwargs = {
        'input_shape_dict': input_shape_dict,
        'dummy_input': dummy_input,
        'output_path': output_path,
        'model_name': model_name,
        'onnx_model_path': osp.join(output_path, '{}.onnx'.format(model_name)),
        'deploy_to_qlinear': deploy_to_qlinear,
    }
    kwargs.update(extra_kwargs)
    deploy_model = deepcopy_graphmodule(model)
    deploy_model.to(device)
    for convert_function in BACKEND_DEPLOY_ONNX_FUNCTION[backend_type]:
        convert_function(deploy_model, **kwargs)