import json
import yaml
import sys
import os


def load_config(config_path):
    with open(config_path, "r") as file:
        config = yaml.safe_load(file)

    if "quant" in config:
        quant_info = config["quant"]
    if "output_path" in quant_info:
        output_path = os.path.join(quant_info["output_path"], config["task_id"])
    if "march" in quant_info:
        march = "nash" if "nash" in quant_info["march"] else "bayes"
        convert_yaml_path = f"pyq_yaml/{march}/"
    # prepare the nash and bayes bpu
    ## first prepare the VisionEncoder yaml
    VisionEncoder_yaml_path = os.path.join(convert_yaml_path, "VisionEncoder.yaml")
    with open(VisionEncoder_yaml_path, "r") as file:
        VisionEncoder_yaml = yaml.safe_load(file)
    VisionEncoder_yaml["model_parameters"]["onnx_model"] = quant_info["VisionEncoder"]["onnx_model"]
    VisionEncoder_yaml["calibration_parameters"]["cal_data_dir"] = quant_info["VisionEncoder"]["calibration_data"]
    VisionEncoder_yaml["model_parameters"]["march"] = quant_info["march"]

    # Make sure output ptq_yaml directory exists
    output_ptq_yaml_dir = os.path.join(output_path, "ptq_yaml")
    os.makedirs(output_ptq_yaml_dir, exist_ok=True)

    # Save VisionEncoder yaml to output/ptq_yaml
    VisionEncoder_yaml_save_path = os.path.join(output_ptq_yaml_dir, "VisionEncoder.yaml")
    with open(VisionEncoder_yaml_save_path, "w") as file:
        yaml.safe_dump(VisionEncoder_yaml, file)

    ## second prepare the TransformerLayers yaml
    TransformerLayers_yaml_path = os.path.join(convert_yaml_path, "TransformerLayers.yaml")
    with open(TransformerLayers_yaml_path, "r") as file:
        TransformerLayers_yaml = yaml.safe_load(file)
    TransformerLayers_yaml["model_parameters"]["onnx_model"] = quant_info["TransformerLayers"]["onnx_model"]
    TransformerLayers_yaml["model_parameters"]["march"] = quant_info["march"]
    
    TransformerLayers_Cal_dir = quant_info["TransformerLayers"]["calibration_data"]
    # (Fix cal_data_dir variable)
    cal_data_dir = TransformerLayers_Cal_dir
    sub_dirs = [d for d in os.listdir(cal_data_dir) if os.path.isdir(os.path.join(cal_data_dir, d))]
    input_names = []
    for name in sub_dirs:
        if name == "state":
            input_names.append("states")
        else:
            input_names.append(f"{name}_features")
    input_name_str = ";".join(input_names) + ";"

    TransformerLayers_yaml["input_parameters"]["input_name"] = input_name_str
    TransformerLayers_yaml["input_parameters"]["input_type_rt"] = "featuremap;" * len(input_names)
    TransformerLayers_yaml["input_parameters"]["input_layout_rt"] = "NCHW;" * len(input_names)
    TransformerLayers_yaml["input_parameters"]["input_type_train"] = "featuremap;" * len(input_names)
    TransformerLayers_yaml["input_parameters"]["input_layout_train"] = "NCHW;" * len(input_names)
    TransformerLayers_yaml["input_parameters"]["norm_type"] = "no_preprocess;" * len(input_names)
    TransformerLayers_yaml["calibration_parameters"]["cal_data_dir"] = ";".join([os.path.join(TransformerLayers_Cal_dir, name) for name in sub_dirs]) + ";"
    TransformerLayers_yaml["calibration_parameters"]["cal_data_type"] = "float32;" * len(input_names)

    # Save TransformerLayers yaml to output/ptq_yaml
    TransformerLayers_yaml_save_path = os.path.join(output_ptq_yaml_dir, "TransformerLayers.yaml")
    with open(TransformerLayers_yaml_save_path, "w") as file:
        yaml.safe_dump(TransformerLayers_yaml, file)
            
            




if __name__ == "__main__":
    config_path = sys.argv[1]
    config = load_config(config_path)