# idea: generate the save folder name with some random generate string, in order to avoid potential name comflicts in repeat experiments

import os
import string
import random
from typing import *
from datetime import datetime
import sys

sys.path.append('../')


def generate_save_folder(
        run_info: Optional[str] = '',
        given_load_file_path: Optional[str] = None,
        recover: Optional[bool] = False,
        all_record_folder_path: str = '../record',
) -> str:
    # idea:  This function helps to generate save path for experiment.
    # if you do not want to set the name, this function will set it for experiment.
    # Note that by using the randomly generate str, replication of experiment may not overwrite the folder of each other

    def inside_generate(
            all_record_folder_path: str,
            startTimeStr: str,
            run_info: str,
    ) -> str:
        random_code = ''.join(
            [random.choice(string.ascii_letters + string.digits) for n in range(4)])
        save_path = all_record_folder_path + '/' + startTimeStr + '_' + os.path.basename(sys.argv[0]).split('.')[
            0] + '_' + run_info + '_' + random_code
        return save_path

    startTimeStr = str(datetime.now().strftime('%Y%m%d_%H%M%S'))

    if given_load_file_path is None:
        # default None
        # no save_folder and no load_path, so can only random generate one

        save_path = inside_generate(
            all_record_folder_path,
            startTimeStr,
            run_info,
        )

        while os.path.isdir(save_path):
            save_path = inside_generate(
                all_record_folder_path,
                startTimeStr,
                run_info,
            )

        # os.mkdir(save_path)

    elif given_load_file_path is not None and recover is True:

        given_load_file_path = given_load_file_path.rstrip('/')

        if os.path.isfile(os.path.abspath(given_load_file_path)):
            load_folder_name = os.path.dirname(given_load_file_path)
        else:
            load_folder_name = given_load_file_path

        print(load_folder_name)

    else:

        given_load_file_path = given_load_file_path.rstrip('/')

        # isfile need abs path otherwise seems to be false anyway.
        if os.path.isfile(os.path.abspath(given_load_file_path)):
            load_folder_name = os.path.basename(
                os.path.dirname(given_load_file_path))
        else:
            load_folder_name = os.path.basename(given_load_file_path)

        generate_base = inside_generate(
            all_record_folder_path,
            startTimeStr,
            run_info,
        )

        save_path = generate_base.split(
            '_baseOn_')[0] + '_baseOn_' + load_folder_name
        # if not contains "_baseOn_", then no split, 0 position is the original
        # else, then keep only the part before the first _baseOn_, rest of str drop

        while os.path.isdir(save_path):

            generate_base = inside_generate(
                all_record_folder_path,
                startTimeStr,
                run_info,
            )

            save_path = generate_base.split(
                '_baseOn_')[0] + '_baseOn_' + load_folder_name

    print('*'*10)
    print(save_path)
    os.mkdir(save_path)

    return save_path
