import json
import os

from Scripts.Prompting import *
from Scripts.Utils import *


def initialize(raw_args=None, baseline=0):
    args = get_args(raw_args)
    args.benchmark = args.dataset.split("_")[0]
    folder = f"config-{args.sample}sample_{args.n_values}n_{args.threshold}threshold_{args.temp}temp"
    files_check = [
        "completeness_checking",
        "dataset",
        "min_spec_candidates",
        "min_spec_dataset",
        "updated_dataset",
    ]
    if not baseline:
        files_check.extend(["struct_check", "struct_scorer"])
        base_folder = os.path.join(
            "Responses", args.llm, args.info_prompt, "all_temp_variation", folder
        )
    else:
        base_folder = os.path.join(
            "Responses", args.llm, args.info_prompt, args.dataset, args.prompt, folder
        )

    # initialize all folders
    create_folder_if_not_exists(base_folder)
    for f in files_check:
        create_folder_if_not_exists(os.path.join(base_folder, f))
    create_folder_if_not_exists(
        os.path.join(base_folder, "completeness_checking", "display_csv")
    )

    # create the subset of dataset
    dataset_folder = os.path.join(
        base_folder, "dataset", args.dataset + "_ambiguous_0.json"
    )
    original_folder = os.path.join(
        "Responses",
        "cache",
        args.dataset + "_ambiguous_0.json",
    )
    if os.path.exists(dataset_folder):
        print(f"Responses already exists at {dataset_folder}")
        return

    data = json.load(open(original_folder, "r"))
    subset = []  
    for i, d in enumerate(data):
        if i == args.sample:
            break
        subset.append(d)

    with open(dataset_folder, "w") as f:
        json.dump(subset, f)
    return


if __name__ == "__main__":
    # try:
        initialize(raw_args=None)
    # except Exception as e:
    #     import pdb
    #     pdb.post_mortem()
