from utils.common import *
from utils.helper import *
from utils.load_config import BatchConfig
from utils.generate_batch_config import generate_batch_bconf, generate_batch_conf

import argparse
import os
import shutil
import subprocess

# Parse arguments
argparser = argparse.ArgumentParser(
    description="Run batch compiler for KDP 520")

argparser.add_argument(
    '-c',
    '--config',
    type=str,
    default='/data1/batch_input_params.json',
    help="path of input_params.json")

args = argparser.parse_args()

# Load config
batch_config = BatchConfig(args.config)
prepare_tmp_folder()
output_path = prepare_result_folder('batch_compile')

# check config
bie_models = dict()
for batch_model in batch_config.model_list:
    if batch_model.path.endswith(".bie"):
        continue
    elif batch_model.path.endswith(".onnx") and batch_model.radix is not None:
        continue
    else:
        raise ValueError("Currently, batch compile only support models after fix point analysis.")

# Generate batch compile config
generate_batch_conf(batch_config.encryption,
                    '520',
                    output_path + '/batch_compile_config.json',
                    weight_compress=batch_config.weight_compress)

# Generate batch compile b_conf
generate_batch_bconf(batch_config,
                    output_path + '/batch_compile_config.json',
                    output_path + '/batch_compile_bconfig.json')

# Run batch compile
os.chdir(output_path)
commands = [LIBS_FOLDER + '/compiler/batch_compile',
                '-T', '520',
                output_path + '/batch_compile_bconfig.json',
                '-s', '-t', get_toolchain_version()]
if batch_config.dedicated_buffer:
    commands.append('-o')
subprocess.run(commands, check=True)

delete_tmp_folder()