import os
import shutil
import random
from tqdm import tqdm

def main():
    # Set the dataset path
    Dataset_Path = 'CVRPDataset'
    img_dir = os.path.join(Dataset_Path, 'img_dir')
    ann_dir = os.path.join(Dataset_Path, 'ann_dir')

    # Set the ratio of training set to test set
    test_frac = 0.1 
    random.seed(123) 

    # Create directories if not exist
    os.makedirs(os.path.join(Dataset_Path, 'train'), exist_ok=True)
    os.makedirs(os.path.join(Dataset_Path, 'val'), exist_ok=True)

    # Get image file paths
    img_paths = os.listdir(img_dir)
    random.shuffle(img_paths) 

    val_number = int(len(img_paths) * test_frac)  
    train_files = img_paths[val_number:]          
    val_files = img_paths[:val_number]            

    print(f"Total images: {len(img_paths)}")
    print(f"Training set images: {len(train_files)}")
    print(f"Test set images: {len(val_files)}")

    # Move the training set images to the train directory
    for each in tqdm(train_files, desc="Move the training set images"):
        src_path = os.path.join(img_dir, each)
        dst_path = os.path.join(Dataset_Path, 'train', each)
        shutil.move(src_path, dst_path)
        
    # Move the test set images to the test directory
    for each in tqdm(val_files, desc="Move the test set images"):
        src_path = os.path.join(img_dir, each)
        dst_path = os.path.join(Dataset_Path, 'val', each)
        shutil.move(src_path, dst_path)

    # Move the train and val directories into img_dir
    shutil.move(os.path.join(Dataset_Path, 'train'), os.path.join(img_dir, 'train'))
    shutil.move(os.path.join(Dataset_Path, 'val'), os.path.join(img_dir, 'val'))

    # Process annotation files
    # Ensure the annotation directories exist
    os.makedirs(os.path.join(Dataset_Path, 'train'), exist_ok=True)
    os.makedirs(os.path.join(Dataset_Path, 'val'), exist_ok=True)

    # Move the training set annotation files to the train directory
    for each in tqdm(train_files, desc="Move the training set annotations"):
        src_path = os.path.join(ann_dir, each.split('.')[0] + '.png')
        dst_path = os.path.join(Dataset_Path, 'train', each.split('.')[0] + '.png')
        shutil.move(src_path, dst_path)

    # Move the test set annotation files to the test directory
    for each in tqdm(val_files, desc="Move the test set annotations"):
        src_path = os.path.join(ann_dir, each.split('.')[0] + '.png')
        dst_path = os.path.join(Dataset_Path, 'val', each.split('.')[0] + '.png')
        shutil.move(src_path, dst_path)

    # Move the train and val annotation directories into ann_dir
    shutil.move(os.path.join(Dataset_Path, 'train'), os.path.join(ann_dir, 'train'))
    shutil.move(os.path.join(Dataset_Path, 'val'), os.path.join(ann_dir, 'val'))

if __name__ == '__main__':
    main()