import itertools
import random
import argparse
import os 
import sys

from joint_evaluation import DATASET_TO_COMBINATIONS_RESULTS_PATH

sys.path.append("../weaver")
from verification import DATASET_TO_REWARD_MODELS, DATASET_TO_LM_JUDGES

def generate_and_shuffle_tuples(verifiers, k, dataset):
    # Generate all k-tuples
    tuples = list(itertools.combinations(verifiers, k))
    # Shuffle the tuples
    random.shuffle(tuples)
    # Write to the output file

    os.makedirs(DATASET_TO_COMBINATIONS_RESULTS_PATH[dataset], exist_ok=True)
    path = os.path.join(DATASET_TO_COMBINATIONS_RESULTS_PATH[dataset], f"all_{k}_tuples.txt")


    with open(path, 'w') as f:
        for t in tuples:
            f.write(' '.join(t) + '\n')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate and shuffle k-tuples of verifiers.")
    parser.add_argument("--dataset_path", type=str, required=True,
                      help="Path to the dataset")
    parser.add_argument('--k', type=int, required=True, help="Size of each tuple")
    args = parser.parse_args()


    reward_models = DATASET_TO_REWARD_MODELS[args.dataset_path]
    lm_judges = DATASET_TO_LM_JUDGES[args.dataset_path]

    verifiers = reward_models + lm_judges

    generate_and_shuffle_tuples(verifiers, args.k, args.dataset_path)