from .common import *
from .load_config import ModelConfig
from .helper import prepare_result_folder, prepare_tmp_folder, delete_tmp_folder
from .img_preprocess import img2txt_bin
from .generate_fpAnalyser_input_json import generate_knerex_config

import subprocess
import os
import shutil
import numpy as np

def knerex_output_file_name(input_file_path: str) -> str:
    old_file_name = os.path.basename(input_file_path)
    return old_file_name[:-5] + '.quan.wqbi.bie'

def converted_input_folder(input_name: str) -> str:
    if '/' in input_name:
        return input_name.replace('/', '__')
    else:
        return input_name

def run_knerex(config : ModelConfig, thread_num: int =1, hardware: int= 520) -> str:
    """Run knerex with the given config.

    ret: the knerex generated bie file
    """
    logging.info("Preparing FP Analyser inputs...")
    prepare_result_folder('fpAnalyser')
    prepare_tmp_folder()

    # Prepare inputs
    for input_name in config.input_shapes:
        input_folder = config.input_folders[input_name]
        input_shape = config.input_shapes[input_name]
        img2txt_bin(input_folder, True, TMP_FOLDER + "/" + converted_input_folder(input_name),
                    input_shape[2], input_shape[3], config.preprocess_config["img_channel"],
                    config.preprocess_config["img_preprocess_method"], 'img2txt',
                    hardware, config.preprocess_config["radix"],
                    x_pos=config.preprocess_config["p_crop"]["crop_x"],
                    y_pos=config.preprocess_config["p_crop"]["crop_y"],
                    crop_h=config.preprocess_config["p_crop"]["crop_h"],
                    crop_w=config.preprocess_config["p_crop"]["crop_w"],
                    keep_aspect_ratio=config.preprocess_config["keep_aspect_ratio"],
                    enable_crop=config.enable_crop)

    # Prepare knerex config
    logging.info("Preparing FP Analyser config...")
    generate_knerex_config(config.model_file,
                            {input_name: TMP_FOLDER + "/" + converted_input_folder(input_name) for input_name in config.input_shapes},
                            num_of_thread=thread_num,
                            hardware=hardware,
                            outlier=config.outlier,
                            output_updater_json=TMP_FOLDER + '/updater.json')
    # Run Knerex
    logging.info("Running FP Analyzer...")
    os.chdir(TMP_FOLDER)
    commands = [LIBS_FOLDER + '/fpAnalyser/updater/run_updater', '-i', TMP_FOLDER + '/updater.json']
    if config.quantize_mode == "post_sigmoid":
        commands.append("-n")
        commands.append("yolo")
    subprocess.run(commands, check=True)
    output_path = RESULT_FOLDER + "/fpAnalyser/" + knerex_output_file_name(config.model_file)
    shutil.copy2(TMP_FOLDER + '/output.quan.wqbi.bie', output_path)
    return output_path

def run_knerex_without_preprocess(model_path: str,
                                    model_input: Dict,
                                    thread_num: int =1,
                                    hardware: int= 520,
                                    quantize_mode: str="default",
                                    output_path: str=None,
                                    outlier: float=0.999,
                                    datapath_range_method: str='percentage',
                                    percentage: float=0.999,
                                    percentile: float=0.001,
                                    outlier_factor: float=1.0,
                                    bitwidth_mode: str="8") -> str:
    """Run knerex with the given config.

    ret: the knerex generated bie file
    """
    # Save the data to txt
    logging.info("Preparing FP Analyser inputs...")
    prepare_result_folder('fpAnalyser')
    prepare_tmp_folder()
    for input_name in model_input:
        output_folder = TMP_FOLDER + "/" + converted_input_folder(input_name)
        if not os.path.exists(output_folder):
            os.mkdir(output_folder)
        for i in range(len(model_input[input_name])):
            out_file = os.path.join(output_folder, str(i) + ".txt")
            np.savetxt(out_file, model_input[input_name][i].reshape((-1, 1)), fmt="%.8f")

    # Prepare knerex config
    logging.info("Preparing FP Analyser config...")
    generate_knerex_config(model_path,
                            {input_name: TMP_FOLDER + "/" + converted_input_folder(input_name) for input_name in model_input},
                            num_of_thread=thread_num,
                            outlier=outlier,
                            hardware=hardware,
                            output_updater_json=TMP_FOLDER + '/updater.json',
                            datapath_range_method = datapath_range_method,
                            percentage = percentage,
                            percentile = percentile,
                            outlier_factor = outlier_factor,
                            bitwidth_mode = bitwidth_mode)
    # Run Knerex
    logging.info("Running FP Analyzer...")
    os.chdir(TMP_FOLDER)
    commands = [LIBS_FOLDER + '/fpAnalyser/updater/run_updater', '-i', TMP_FOLDER + '/updater.json']
    if quantize_mode == "post_sigmoid":
        commands.append("-n")
        commands.append("yolo")
    subprocess.run(commands, check=True)
    if output_path is None:
        output_path = RESULT_FOLDER + "/fpAnalyser/" + knerex_output_file_name(model_path)
    shutil.copy2(TMP_FOLDER + '/output.quan.wqbi.bie', output_path)
    return output_path