import os
import glob
import shutil
import random

PATCH_DIR = "dataset/extracted_patches"
DATASET_DIR = "dataset/patches_10m"

def main():
    random.seed(42)
    
    features = sorted(glob.glob(os.path.join(PATCH_DIR, "features", "*.npy")))
    if not features:
        print(f"No patches found in {PATCH_DIR}")
        return
        
    random.shuffle(features)
    
    total = len(features)
    train_end = int(total * 0.70)
    val_end = int(total * 0.85)
    
    train_files = features[:train_end]
    val_files = features[train_end:val_end]
    test_files = features[val_end:]
    
    for split in ['train', 'val', 'test']:
        os.makedirs(os.path.join(DATASET_DIR, split, 'features'), exist_ok=True)
        os.makedirs(os.path.join(DATASET_DIR, split, 'masks'), exist_ok=True)
        
    print(f"Total patches: {total}")
    print(f"Train set: {len(train_files)} patches")
    print(f"Validation set: {len(val_files)} patches")
    print(f"Test set: {len(test_files)} patches")
    
    print("Copying to split directories...")
    
    def copy_split(files, split_name):
        for f in files:
            name = os.path.basename(f)
            m = os.path.join(PATCH_DIR, "masks", name)
            
            shutil.copy(f, os.path.join(DATASET_DIR, split_name, 'features', name))
            if os.path.exists(m):
                shutil.copy(m, os.path.join(DATASET_DIR, split_name, 'masks', name))
                
    copy_split(train_files, 'train')
    copy_split(val_files, 'val')
    copy_split(test_files, 'test')
    
    print(f"Split successful. Dataset is at {DATASET_DIR}")

if __name__ == '__main__':
    main()