"""
Separated data to image and sound folders.
Author: Martin Tomasovic
"""
import os
import shutil

# original paths
path_train = "./SUR_projekt2024-2025/train"
path_dev = "./SUR_projekt2024-2025/dev"

# new paths with images/class structure
path_new_train = "./Separate_data/train"
path_new_dev = "./Separate_data/dev"

# create base directories
os.makedirs('./Separate_data', exist_ok=True)
os.makedirs(path_new_train, exist_ok=True)
os.makedirs(path_new_dev, exist_ok=True)

# create unified images and sounds directories
os.makedirs(os.path.join(path_new_train, "images"), exist_ok=True)
os.makedirs(os.path.join(path_new_train, "sounds"), exist_ok=True)
os.makedirs(os.path.join(path_new_dev, "images"), exist_ok=True)
os.makedirs(os.path.join(path_new_dev, "sounds"), exist_ok=True)

def process_dataset(source_root, dest_root):
    for root, dirs, files in os.walk(source_root):
        # get class name
        class_name = os.path.basename(root)
        
        # skip the root directory itself
        if root == source_root:
            continue
            
        # create class subdirectories in images and sounds
        img_class_dir = os.path.join(dest_root, "images", class_name)
        sound_class_dir = os.path.join(dest_root, "sounds", class_name)
        
        os.makedirs(img_class_dir, exist_ok=True)
        os.makedirs(sound_class_dir, exist_ok=True)
        
        for file in files:
            src_path = os.path.join(root, file)
            if file.lower().endswith('.png'):
                dst_path = os.path.join(img_class_dir, file)
                shutil.copy(src_path, dst_path)
                print(f"Copied image: {src_path} → {dst_path}")
            elif file.lower().endswith('.wav'):
                dst_path = os.path.join(sound_class_dir, file)
                shutil.copy(src_path, dst_path)
                print(f"Copied sound: {src_path} → {dst_path}")

print("Processing training set...")
process_dataset(path_train, path_new_train)

print("\nProcessing dev set...")
process_dataset(path_dev, path_new_dev)

print("\nFinished!")