
import argparse  
import utils.utils as utils
  
def parse_args():  
    parser = argparse.ArgumentParser(description='Generate answer for problem')
    #### general args
    parser.add_argument('--data', default="train", choices=['train', 'test'], type=str, help='train or test set to be used')
    parser.add_argument('--cache_dir', default="../cache/dataset", type=str, help='cache directory for dataset')
    parser.add_argument('--num_process', default=1, type=int, help='number of processes')
    parser.add_argument('--ref_iter', default=0, type=int, help='reference iteration id of data generation')
    parser.add_argument('--save_iter', default=0, type=int, help='save iteration id of data generation')
    parser.add_argument('--index', default="0-10", type=utils.parse_range, help='range of indices')  
    parser.add_argument('--num_retrive', default=8, type=int, help='number of retrived samples')
    parser.add_argument('--overwrite', default=True, type=utils.parse_bool, help='whether to cover existing')
    parser.add_argument('--type', default="solution", choices=['check', 'mutation', 'refresh', 'simplification', 'fusion', 'solution', 'formalization', 'informalization'], \
                            type=str, help='type of generation')
    #### informalization args
    parser.add_argument('--informal_prompt', choices=['origin', 'modify', 'math-word'], default='math-word', type=str, help='which prompt used for informalization')
    #### mutation args
    parser.add_argument('--mutation_type', default="mcmc", choices=['mcmc', 'enum'], type=str, help='type of mutation')
    parser.add_argument('--remove_aux', default=True, type=utils.parse_bool, help='whether to remove auxillary vars after mutation')
    parser.add_argument('--allow_real', default=False, type=utils.parse_bool, help='whether to try mutation in real type if failed')
    parser.add_argument('--simplify', default=True, type=utils.parse_bool, help='whether to simplify the expression')
    parser.add_argument('--comment', default=True, type=utils.parse_bool, help='whether to add comment')
    parser.add_argument('--rename', default=False, type=utils.parse_bool, help='whether to rename the vars when mutating or refreshing')
    args = parser.parse_args()

    return args
