import os
import shutil

def split_train_dev():
    """
    Splits the files into 'train' and 'dev' based on a session parameter (01, 02, 03, or 04).

    Parameters:
    - input_folder (str): Path to the folder containing all data with subfolders for each class (1 to 31)
    - output_folder (str): Path to the output folder where 'train' and 'dev' will be created
    """
    sessions = ['01', '02', '03', '04']

    input_folder = "merged_data"
    output_folder = "data_session_"

    for session in sessions:
        output_session_folder = output_folder + session

        # Create the 'train' and 'dev' folders
        train_folder = os.path.join(output_session_folder, 'train')
        dev_folder = os.path.join(output_session_folder, 'dev')
        os.makedirs(train_folder, exist_ok=True)
        os.makedirs(dev_folder, exist_ok=True)


        # Iterate over the subfolders 1 to 31
        for i in range(1, 32):
            subfolder_path = os.path.join(input_folder, str(i))

            if os.path.exists(subfolder_path):
                train_subfolder = os.path.join(train_folder, str(i))
                dev_subfolder = os.path.join(dev_folder, str(i))
                os.makedirs(train_subfolder, exist_ok=True)
                os.makedirs(dev_subfolder, exist_ok=True)

                for item in os.listdir(subfolder_path):
                    item_path = os.path.join(subfolder_path, item)

                    if os.path.isfile(item_path):
                        # Check if the file has the demanded session number
                        if f"_{session}_" in item:
                            # Move to dev folder
                            shutil.copy2(item_path, os.path.join(dev_subfolder, item))
                        else:
                            # Move to train folder
                            shutil.copy2(item_path, os.path.join(train_subfolder, item))


        print(f"Files with session '{session}' moved to 'dev' folder.")
        print("Other files moved to 'train' folder.")


split_train_dev()