from .common import *
import onnx
import json
import argparse
import os


def generate_knerex_config(input_onnx_file: str, input_txt_folders: dict,
            output_updater_json: str = TMP_FOLDER + "/updater.json",
            outlier: float=0.999,
            num_of_thread: int=8,
            output_model: str=TMP_FOLDER + "/output.bie",
            hardware: int=520,
            datapath_range_method: str='percentage',
            percentage: float=0.999,
            percentile: float=0.001,
            outlier_factor: float=1.0,
            bitwidth_mode: str="8"):

    # Get onnx input name
    input_txts = [{"image_folder": input_txt_folders[key], "operation_name": key}
                    for key in input_txt_folders]

    if hardware == 520:
        hardware_type = 7
    elif hardware == 720:
        hardware_type = 6
    elif hardware == 530:
        hardware_type = 8
    else:
        hardware_type = 0
    # Construct output updater config
    dict_updater = {
        "inmodel": input_onnx_file,
        "datapath": " ",
        "weight": " ",
        "testconfig": " ",
        "type": hardware_type,
        "outmodel": output_model,
        "verbose": 0,
        "samescale": 0,
        "outputscale": 1,
        "cpuscale": 1,
        "num_threads": int(num_of_thread),
        "dump_level": 7,
        "max_scale": 0.984375,
        "bn_weight_pct ": 1,
        "conv_weight_pct ": 1,
        "model_input_txts": input_txts,
        "inferencer_type": "CPU",
        "outlier": float(outlier),
        "datapath_range_method": datapath_range_method,
        "percentage": percentage,
        "percentile": percentile,
        "outlier_factor": outlier_factor,
        "bitwidth_mode": bitwidth_mode
    }

    base_dir = os.path.dirname(os.path.realpath(input_onnx_file))
    user_config_path = os.path.join(base_dir, "user_config.json")
    if os.path.isfile(user_config_path):
        dict_updater["inputconfig"] = user_config_path
    elif os.path.isfile("/tmp/user_config.json"):
        dict_updater["inputconfig"] = "/tmp/user_config.json"

    # Write to file
    if os.path.exists(output_updater_json):
        os.remove(output_updater_json)

    fp = open(output_updater_json, 'w+')
    json.dump(dict_updater, fp)
    fp.close()


    # if os.path.exists(output_bias_json):
    #     os.remove(output_bias_json)

    # fp = open(output_bias_json, 'w+')
    # json.dump(dict_datapath, fp)
    # fp.close()

if __name__ == "__main__":
    argparser = argparse.ArgumentParser(
        description="generate fpAnalyser input json file"
    )

    argparser.add_argument(
        '-i',
        '--input_onnx_file',
        help="path of input onnx model file"
    )

    argparser.add_argument(
        '-im',
        '--input_img_folder',
        default="/workspace/.tmp/fpAnalyserInput",
        help="path of input img folder"
    )

    argparser.add_argument(
        '-odata',
        '--output_datapath_json',
        default="/workspace/.tmp/datapath_analysis.json",
        help="output datapath json file"
    )

    argparser.add_argument(
        '-oupdate',
        '--output_updater_json',
        default="/workspace/.tmp/updater_520.json",
        help="output updater json file"
    )

    argparser.add_argument(
        '-obias',
        '--output_bias_json',
        default="/workspace/.tmp/bias_adjust.json",
        help="output bias adjust json file"
    )

    argparser.add_argument(
        '-ot',
        '--outlier',
        default=0.999,
        help="outlier"
    )
    argparser.add_argument(
        '-thread',
        '--num_of_thread',
        default=8,
        help="num of thread"
    )

    args = argparser.parse_args()

    # Get input names
    onnx_model = onnx.load(args.input_onnx_file)
    assert len(onnx_model.graph.input)  == 1, "currently cannot support model with multily inputs"
    onnx_input_name = onnx_model.graph.input[0].name

    generate_knerex_config(args.input_onnx_file, {onnx_input_name, args.input_img_folder}, args.output_updater_json, args.outlier, args.num_of_thread)