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 730")

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,
    "730",
    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_V2_FOLDER + "/compiler/batch_compile",
    "-T",
    "730",
    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()