
import os
import sys
import json
import argparse
import pandas as pd
import logging
logging.basicConfig(format='%(asctime)s [%(name)s %(filename)s:%(funcName)s:%(lineno)d] %(levelname)s: %(message)s', level=logging.INFO)
import overall_utils
from game24.run_game24 import run_game24
from sudoku.run_sudoku import run_sudoku
from gsm8k.run_math import run_math
from gsm8k.run_math_gang import run_math_gang
from QA.run_qa import run_qa

# Silence other loggers
logger_blocklist = ["httpx"]
for module in logger_blocklist:
    logging.getLogger(module).setLevel(logging.WARNING)

qa_datasets = ["hotpotQA", "CSQA", "StrategyQA"]

if __name__ == '__main__':
    # load arg
    args = argparse.ArgumentParser()
    args.add_argument('--config_file', type=str, default='configs/game_of_24_bfs.yaml')
    args.add_argument('--key', type=str, default="")
    args.add_argument('--max_workers', type=int, default=1)
    args.add_argument('--server', type=str, default="")
    args = args.parse_args()

    # load model
    config = overall_utils._load_yaml(args.config_file)

    dataset_name = config["dataset"]
    config["key"] = args.key
    config["max_workers"] = args.max_workers
    config["server"] = args.server
    logging.info(f"running experiement with dataset_name: {dataset_name}, model_name: {config['llms']}, strategy: {config['strategy']}")
    if dataset_name in qa_datasets:
        # go to run_hotpotqa
        run_qa(config)
    elif dataset_name == "game24":
        run_game24(config)
    elif dataset_name == "sudoku":
        run_sudoku(config)
    elif dataset_name == "gsm8k" or dataset_name == "gsm8k_eval":
        run_math(config)
    elif dataset_name == "counterintuitive_AR":
        run_math(config)
    elif dataset_name == "MATH":
        run_math(config)
    elif dataset_name == "ASDiv":
        run_math(config)
    elif dataset_name == "SVAMP":
        run_math(config)
    elif dataset_name == "TheoremQA":
        run_math(config)