import pdb
import subprocess
import re
import os




def get_output(input_string, k):
    pattern = r"zero_shot/(\w+)/(.+?)/(\w+)"
    match = re.search(pattern, input_string)
    if match:
        part1 = match.group(1)
        part2 = match.group(3) + f"pass{k}_woproof.jsonl"
        result = os.path.join("compile_result", part1, part2)

        # Create the parent directory if it doesn't exist
        parent_dir = os.path.dirname(result)
        os.makedirs(parent_dir, exist_ok=True)

        print(result)
    else:
        print("No match found.")
        assert True
    return result

# List of input paths
input_path_lists = [
    # "/opt/tiger/formal-align/generate_result/zero_shot/lean4_random_test/generation/lean4_random_15k_rft/5",
    # "/opt/tiger/formal-align/generate_result/zero_shot/lean4_basic_test/generation/lean4_random_15k_rft/5",
    # "/opt/tiger/formal-align/generate_result/zero_shot/wild_test_v2/generation/lean4_random_15k_rft/5"
    # "/opt/tiger/formal-align/generate_result/zero_shot/gsm8k_train/generation/lean4_random_15k/10",
    # "/opt/tiger/formal-align/generate_result/zero_shot/math_train/generation/lean4_random_15k/10",
    # "/opt/tiger/formal-align/generate_result/zero_shot/math_train/generation/lean4_random_15k_rft/10",
    # "/opt/tiger/formal-align/generate_result/zero_shot/lean4_15k_train/generation/lean4_random_15k_rft/10",
    # "/opt/tiger/formal-align/generate_result/zero_shot/lean4_15k_train/generation/lean4_random_15k/10",
    # "/opt/tiger/formal-align/generate_result/zero_shot/lean4_basic_test/generation/lean4_random_15k_rft/5/",
    # "/opt/tiger/formal-align/generate_result/zero_shot/lean4_basic_test/generation/lean4_random_15k/5/"
    # "/opt/tiger/formal-align/generate_result/zero_shot/math_train/generation/lean4_random_15k/10",
    "/opt/tiger/formal-align/generate_result/zero_shot/lean4_random_test/generation/lean4_random_15k_rft/5/",
    "/opt/tiger/formal-align/generate_result/zero_shot/lean4_random_test/generation/lean4_random_15k/5/"
]

# Function to extract k from the input path
def extract_k(input_path):
    return os.path.basename(input_path)


# Iterate through the input paths and run the command
for input_path in input_path_lists:
    k = extract_k(input_path)
    k=1
    print(k)

    output_path = get_output(input_path, k)
    output_log_path = output_path +  'woproof.log'


    if "wild_test" in input_path or "gsm8k_train" in input_path or "math_train" in input_path:
        script_name = "pass_rate_notlean_test_woproof.py"
    else:
        script_name = "pass_rate_new_test_woproof.py"

    command = f"python3 {script_name} --input_path {input_path} --output_path {output_path} --k {k}"
    print(command)

    with open(output_log_path, "w") as output_file:
        subprocess.run(command, shell=True, stdout=output_file, stderr=subprocess.STDOUT)
        print("\n\n", file=output_file)