import os

def get_log_dir_index(out_dir):
    dirs = [x[0] for x in os.listdir(out_dir)]
    if '.' in dirs:  # minor change for .ipynb
        dirs.remove('.')
    log_dir_index = '_' + str(len(dirs))

    return log_dir_index

class Config:
    def __init__(self, args):
        # create dirs
        self.base_dir = 'results'
        os.makedirs(self.base_dir, exist_ok=True)

        # common
        self.mode = args.mode
        self.log_dir = args.log_dir
        self.seed = args.seed

        self.dataset = args.dataset
        self.split_name = args.split_name
        self.domain = args.domain
        self.num_shot = args.num_shot
        self.quickrun = args.quickrun
        self.cut_len = args.cut_len

        self.engine = args.engine
        self.temperature = args.temperature
        self.freq_penalty = args.freq_penalty
        self.max_tokens = args.max_tokens
        self.llm_cache_dir = args.llm_cache_dir

        self.use_generic_grammar = args.use_generic_grammar
        self.num_samples = args.num_samples
        self.rule_temperature = args.rule_temperature

        self.prompt_mode = args.prompt_mode
        self.prompt_template = args.prompt_template
        self.add_rule_instruction_flag = args.add_rule_instruction_flag
        self.retrieve_fn = args.retrieve_fn
        self.batch_size = args.batch_size

        self.use_linearized_tree = args.use_linearized_tree

        self.use_oracle_rule_flag = args.use_oracle_rule_flag
        self.constrain_rule_gen_flag = args.constrain_rule_gen_flag
        self.constrain_prog_gen_flag = args.constrain_prog_gen_flag
        self.lazy_constrain_flag = args.lazy_constrain_flag
        self.kg_rule_flag = args.kg_rule_flag
        self.use_action_list_flag = args.use_action_list_flag
        self.add_rule_list_flag = args.add_rule_list_flag

        self.num_iterations = args.num_iterations
        self.iter_prompt_template = args.iter_prompt_template
        self.iter_engine = args.iter_engine
        self.iter_retrieve_fn = args.iter_retrieve_fn

        index = get_log_dir_index(self.base_dir)
        self.cfg_dir = '%s/%s' % (self.base_dir, self.mode + index)
        os.makedirs(self.cfg_dir, exist_ok=True)
        self.result_dir = '%s/results' % self.cfg_dir
        self.log_dir = '%s/log' % self.cfg_dir
        os.makedirs(self.result_dir, exist_ok=True)
        os.makedirs(self.log_dir, exist_ok=True)