import pdb
import subprocess
import re

# Output file
output_file = "pass_rate_output.txt"

# Clearing the output file before appending new content
with open(output_file, "w") as file:
    file.write("")

def get_output(input_path, k):
    try:
        # Split the input path based on '/'
        parts = input_path.split('/')
        
        # Find the index of 'zero_shot' and get the 4th part after it
        zero_shot_index = parts.index('zero_shot')

        part1 = parts[zero_shot_index + 1]
        part3 = parts[zero_shot_index + 3]
        part4 = parts[zero_shot_index + 4]
        result = f"openllm_result/{part3}/{part1}_pass{part4}.json"
        print(result)
        return result
    
    except:

        print("No match found.")
        return None


# List of input paths
input_path_lists = [
    # "../auto-info/generate_result/zero_shot/wild_test/generation/deepseek-math-7b-base/1/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/deepseek-math-7b-instruct/1/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/deepseek-math-7b-instruct/1/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/llemma_7b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/llemma_7b/1/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/llemma_7b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/llemma_34b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/llemma_34b/1/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/llemma_34b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/internlm2-math-7b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/internlm2-math-7b/1/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/internlm2-math-7b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/Mistral-7B-Instruct-v0.2/1/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/Mistral-7B-Instruct-v0.2/1/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/Mistral-7B-Instruct-v0.2/1/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/internlm2-math-20b/1/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/internlm2-math-20b/1/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/internlm2-math-20b/1/",

    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/deepseek-math-7b-base/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/deepseek-math-7b-base/5/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/deepseek-math-7b-base/5/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/deepseek-math-7b-instruct/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/deepseek-math-7b-instruct/5/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/deepseek-math-7b-instruct/5/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/llemma_7b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/llemma_7b/5/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/llemma_7b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/llemma_34b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/llemma_34b/5/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/llemma_34b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/internlm2-math-7b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/internlm2-math-7b/5/",
    "../auto-info/generate_result/zero_shot/wild_test/generation/internlm2-math-7b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/Mistral-7B-Instruct-v0.2/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/Mistral-7B-Instruct-v0.2/5/",
    # "../auto-info/generate_result/zero_shot/wild_test/generation/Mistral-7B-Instruct-v0.2/5/",
    # "../auto-info/generate_result/zero_shot/lean4_basic_test/generation/internlm2-math-20b/5/",
    # "../auto-info/generate_result/zero_shot/lean4_random_test/generation/internlm2-math-20b/5/",
    "../auto-info/generate_result/zero_shot/wild_test/generation/internlm2-math-20b/5/",
]


def extract_group(input_path):
    try:
        # Split the input path based on '/'
        parts = input_path.split('/')
        
        # Find the index of 'zero_shot' and get the 4th part after it
        zero_shot_index = parts.index('zero_shot')
        k = parts[zero_shot_index + 4]
        return k
    except (IndexError, ValueError):
        # Handle cases where 'zero_shot' is not found or there are not enough parts
        return None

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

    if "wild" in input_path or "gsm8k_train" in input_path or "math_train" in input_path:
        print(f"wild")
        print(f"Running for input path: {input_path}", file=open(output_file, "a"))
        command = f"python3 openllm_pass_rate_new_notlean_test.py --input_path {input_path} --output_path {get_output(input_path,k)}  --k {k}"
        subprocess.run(command, shell=True, stdout=open(output_file, "a"), stderr=subprocess.STDOUT)
        print("\n\n",file=open(output_file, "a"))
    else:
        print(f"lean")
        print(f"Running for input path: {input_path}", file=open(output_file, "a"))
        command = f"python3 openllm_pass_rate_new_test.py --input_path {input_path} --output_path {get_output(input_path, k)} --k {k}"
        subprocess.run(command, shell=True, stdout=open(output_file, "a"), stderr=subprocess.STDOUT)
        print("\n\n",file=open(output_file, "a"))