import os
import shutil

from loguru import logger


def get_filename() -> str:
    current_file_name = os.path.basename(__file__)
    log_name = "{}.log"
    return log_name.format(current_file_name.split('.')[0])


def set_log() -> None:
    filename = get_filename()
    logger.add(f'../log/{filename}')


set_log()


# Input and output folder paths
input_folder = '/mnt/nfs/data/geneformer/output/v2/looms'
output_folder = '/mnt/nfs/data/geneformer/output/v2/looms_sub'
per_group_num = 100


if __name__ == '__main__':
    # Create output folder if it does not exist
    os.makedirs(output_folder, exist_ok=True)

    # Iterate through files in the input folder
    files = os.listdir(input_folder)
    group_count = 0

    for i, file in enumerate(files, start=1):
        # Build output subfolder path
        group_folder = os.path.join(output_folder, f'group_{group_count + 1}')

        # Create output subfolder if it does not exist
        os.makedirs(group_folder, exist_ok=True)

        # Build input and output file paths
        input_file = os.path.join(input_folder, file)
        output_file = os.path.join(group_folder, file)

        # Copy file into the output subfolder
        shutil.copy(input_file, output_file)

        if i % per_group_num == 0:
            group_count += 1