import onnx
import json
import uuid
import os
import shutil
import numpy as np
from collections import OrderedDict
from onnx import numpy_helper
from .common import (
    update_inp2node_out2node,
    prepare_initializer,
    prepare_data,
    OnnxPreprocess,
)
from .flatbuffer_nn.utils import generate_fake_tensor_info
from mediatools.utils.registry import MNN_MODULE_CONVERT_CLASSES


def default_dump(obj):
    if isinstance(obj, (np.integer, np.floating, np.bool_)):
        return obj.item()
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    else:
        return obj


def update_outensor_quantinfo(node, extraTensorDescribe, name2data):
    assert node.op_type == "DeQuantizePertensorInt"
    name = node.input[0]
    assert name in extraTensorDescribe
    quant_info = extraTensorDescribe[name]['quantInfo']

    qmin, qmax = None, None
    for attr in node.attribute:
        if attr.name == "quant_max":
            qmax = attr.i
        elif attr.name == "quant_min":
            qmin = attr.i

    multiplier_n, shift_n, zero_n = node.input[1:4]
    mul   = name2data[multiplier_n].astype(np.int64)
    shift = name2data[shift_n].astype(np.int64)
    zero = name2data[zero_n].astype(np.int64)
    scale_float = mul / 2**shift
    
    quant_info['scale'] = float(scale_float)
    quant_info['zero'] = float(zero)
    quant_info['min'] = float(qmin)
    quant_info['max'] = float(qmax)


def update_json_for6921(json_info, model_type):
    if model_type in ['det', 'det_cls', 'det_lmk', 'det_lmk_cls']:
        for op in json_info['oplists']:
            if op['main_type'] == 'Convolution2D':
                # update name of output node 
                if 'head0.layer_' in op['main']['weightname'] and 'layer_32_cls_conv' not in op['main']['weightname']:
                    origin_op_name = op['name']
                    new_op_name = op['main']['weightname'].split('.')[1].replace('obj', 'objectness')\
                                    .replace('cls', 'classification')\
                                        .replace('box', 'regression')\
                                        .replace('lmk', 'landmarks')
                    op['name'] = new_op_name
                    out_idx = json_info['outputName'].index(origin_op_name)
                    json_info['outputName'][out_idx] = new_op_name
                    out_tensor_idx = op['outputIndexes'][0]
                    assert json_info['extraTensorDescribe'][out_tensor_idx]['name'] == origin_op_name
                    json_info['extraTensorDescribe'][out_tensor_idx]['name'] = new_op_name
                    tensor_idx = json_info['tensorName'].index(origin_op_name)
                    json_info['tensorName'][tensor_idx] = new_op_name
    elif model_type in ['face_reg']:
        for op in json_info['oplists']:
            if op['main_type'] == 'Convolution2D':
                # update name of output node 
                if 'head0.shared_layer.conv2.0' in op['main']['weightname']:
                    origin_op_name = op['name']
                    new_op_name = 'feature_embedding'
                    op['name'] = new_op_name
                    out_idx = json_info['outputName'].index(origin_op_name)
                    json_info['outputName'][out_idx] = new_op_name
                    out_tensor_idx = op['outputIndexes'][0]
                    assert json_info['extraTensorDescribe'][out_tensor_idx]['name'] == origin_op_name
                    json_info['extraTensorDescribe'][out_tensor_idx]['name'] = new_op_name
                    tensor_idx = json_info['tensorName'].index(origin_op_name)
                    json_info['tensorName'][tensor_idx] = new_op_name
    return json_info

def dump_to_flatbuffer(onnx_model_path, mnn_model_name, **kwargs):
    model = onnx.load(onnx_model_path)
    graph = model.graph
    out2node, inp2node = update_inp2node_out2node(graph)
    name2data = prepare_data(graph)
    named_initializer = prepare_initializer(graph)

    preprocess = OnnxPreprocess()
    preprocess.remove_fake_pad_op(graph, name2data, inp2node, out2node)
    out2node, inp2node = update_inp2node_out2node(graph)

    # json info init
    json_info = {
        "bizCode": "biz",
        "preferForwardType": "CPU",
        "sourceType": "ONNX",
        "tensorNumber": 0,
        "usage": "INFERENCE",
        "mnn_uuid": str(uuid.uuid1()),
    }


    extraTensorDescribe = OrderedDict()
    opinfo = OrderedDict()
    outputName = []
    op_infos = []
    qtensor_info = []

    for node in graph.node:
        try:
            op_infos, qtensor_info = MNN_MODULE_CONVERT_CLASSES[node.op_type]\
                (node, name2data, inp2node, out2node, named_initializer, extraTensorDescribe).dump2json()
        except:
            print('node %s not supported, skip this'%(node.op_type))
        for qinfo in qtensor_info:
            extraTensorDescribe[qinfo['name']] = qinfo
        for op in op_infos:
            opinfo[op['name']] = op
        # add out node
        if node.op_type in ["DeQuantizePertensorInt", "DeQuantizeInt"]:
            output_name = node.input[0]
            outputName.append(output_name)
            update_outensor_quantinfo(node, extraTensorDescribe, name2data)
    
    tensorName = list(extraTensorDescribe.keys())

    json_info.update(
        {
            'extraTensorDescribe': list(extraTensorDescribe.values()),
            'oplists': list(opinfo.values()),
            'outputName': outputName,
            'tensorName': tensorName
        }
    )

    # update for 6921vnne
    if kwargs.get('SA6921vnne', False):
        json_info = update_json_for6921(json_info, kwargs['model_type'])

    with open(mnn_model_name, 'w') as f:
        json.dump(json_info, f, indent=2, default=default_dump)
    
    flatc_path = os.environ.get('FLATC_PATH', None)
    if flatc_path is None:
        print("Flatc address not found, please convert json file manually!!!")
    else:
        schema = 'mediatools/schema/default/MNN.fbs'
        os.system("%s --binary %s %s"%(flatc_path, schema, mnn_model_name))
        base_name = os.path.basename(mnn_model_name)
        shutil.move(base_name.replace('.json', '.bin'), mnn_model_name.replace('.json', '.mnn'))