import onnx
import json
import torch
import numpy as np
from collections import OrderedDict
from onnx import numpy_helper
from .utils import SA6920_VNNE_EXPORTOR 
import mediatools.simulator_nn as msn


def dump_to_sa6920_VNNE(fixpoint_model, model_name, **kwargs):
    fixpoint_modules = dict(fixpoint_model.named_modules())
    fixpoint_nodes = list(fixpoint_model.graph.nodes)
    exportor_params = kwargs.get('exportor_params', {})
    Exportor = SA6920_VNNE_EXPORTOR(**exportor_params)

    for node in fixpoint_nodes:
        ## skip nodes
        if node.op == "call_function":
            continue

        if node.target == "x_post_act_fake_quantizer":
            module = fixpoint_modules[node.target]
            zero = module.zero_point
            Exportor.set_input_zero(zero)
            continue

        if node.op == 'call_module':
            module = fixpoint_modules[node.target]
            process_module_params(module, node, Exportor)

    Exportor.write_bin(model_name)
    Exportor.write_txt()



def process_module_params(module, node, export):
    ## weight params, 需要零点, conv2d + relu
    if isinstance(module, msn.modules.conv.Conv2d):
        export.add_weight(module.weight, node.name + '.weight')
        export.add_bias(module.bias, node.name + '.bias')
    elif isinstance(module, msn.addquantizer_int.AddQuantizerInt):
        export.add_binaryadd_scale(module.scale1, node.name + '.scale_1')
        export.add_binaryadd_scale(module.scale2, node.name + '.scale_2')
        shift1 = module.shift1 - module.shift_reduce
        shift2 = module.shift2 - module.shift_reduce
        export.add_shift(shift1, node.name + '.shift_1')
        export.add_shift(shift2, node.name + '.shift_2')
        export.add_binaryadd_bias(module.bias, node.name + '.add_bias')
    elif isinstance(module, msn.mixedquantizer_int.MixedQuantizerInt):
        export.add_scale(module.scale, node.name + '.scale')
        export.add_shift(module.shift, node.name + '.shift')
        if hasattr(module, "is_relu") and module.is_relu:
            export.add_zero_point(module.clip_value_low, node.name + '.zero_point')
    elif isinstance(module, msn.dequantizer_int.DeQuantizerInt):
        scale = float(module.scale / 2**module.shift)
        export.add_dequant(scale, node.name + '.scale')
        zero = float(module.zero_point)
        export.add_dequant(zero, node.name + '.zero_point')
    elif isinstance(module, (torch.nn.Sequential, torch.nn.modules.container.Sequential)):
        for sub_name, sub_module in module.named_children():
            process_module_params(sub_module, node, export)
    elif type(module) == torch.nn.modules.linear.Identity:
        pass
    else:
        raise NotImplementedError("please check the model")