import os
import shutil

# Paths to the original dataset folder and the new folders
original_folder = '../data/sim_transfer_cube_human'
train_folder = original_folder + '/train'
val_folder = original_folder + '/val'

# Create the train and val folders if they don't exist
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)

# Get a list of all episode files
episode_files = sorted([f for f in os.listdir(original_folder) if f.startswith('episode_') and f.endswith('.hdf5')])

# Separate the files into train and val based on the episode number
train_files = [f for f in episode_files if not f.split('_')[1].split('.')[0].endswith(('8', '9'))]
val_files = [f for f in episode_files if f.split('_')[1].split('.')[0].endswith(('8', '9'))]

# Copy and rename the train files
for i, filename in enumerate(train_files):
    src_path = os.path.join(original_folder, filename)
    dst_path = os.path.join(train_folder, f'episode_{i}.hdf5')
    shutil.move(src_path, dst_path)

# Copy and rename the val files
for i, filename in enumerate(val_files):
    src_path = os.path.join(original_folder, filename)
    dst_path = os.path.join(val_folder, f'episode_{i}.hdf5')
    shutil.move(src_path, dst_path)

print("Dataset split and files renamed successfully.")