# Create a script to split the dataset into train, validation, and test sets
# The filenames are from BASE_PATH/cfs/signals_MAE/EEG/C3-M2/*.npy
import os
import random
import numpy as np

BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'datasets')
SEED = 42
random.seed(SEED)
np.random.seed(SEED)

def split_data(data_dir, output_dir, train_ratio=0.7, val_num=50, test_ratio=0.3):
    # Get all the filenames in the data directory
    all_files = []
    for root, dirs, files in os.walk(data_dir):
        for file in files:
            if file.endswith('.npy'):
                pid = file.split('.')[0]  # Assuming filename format is PID_XXX.npy
                all_files.append(pid)
    
    # Shuffle the filenames
    random.shuffle(all_files)
    
    # Split the data into train, validation, and test sets
    total_files = len(all_files)
    train_end = int(total_files * train_ratio) - val_num
    val_end = train_end + val_num
    
    train_files = all_files[:train_end]
    val_files = all_files[train_end:val_end]
    test_files = all_files[val_end:]
    
    # Save the splits to output directory
    os.makedirs(output_dir, exist_ok=True)
    
    with open(os.path.join(output_dir, 'train.txt'), 'w') as f:
        for file in train_files:
            f.write(file + '\n')
    
    with open(os.path.join(output_dir, 'val.txt'), 'w') as f:
        for file in val_files:
            f.write(file + '\n')
    
    with open(os.path.join(output_dir, 'test.txt'), 'w') as f:
        for file in test_files:
            f.write(file + '\n')

if __name__ == "__main__":
    data_dir = os.path.join(BASE_PATH, 'cfs', 'signals_MAE', 'EEG', 'C3-M2')
    output_dir = '.'
    split_data(data_dir, output_dir)