import os
import shutil
import random

random.seed(42)

def prepare_data(source_dir, target_dir, classes_to_include):
    """
    Restructures the dataset into binary classes (fresh/spoiled).
    """
    for label in ['fresh', 'spoiled']:
        os.makedirs(os.path.join(target_dir, label), exist_ok=True)

    for category in os.listdir(source_dir):
        # Determine if this category belongs to 'fresh' or 'spoiled'
        if any(incl in category for incl in classes_to_include):
            label = 'fresh' if 'fresh' in category else 'spoiled'
            source_path = os.path.join(source_dir, category)
            
            for img in os.listdir(source_path):
                # Ensure we only copy image files (basic check)
                if img.lower().endswith(('.png', '.jpg', '.jpeg')):
                    shutil.copy(
                        os.path.join(source_path, img),
                        os.path.join(target_dir, label, f"{category}_{img}")
                    )

def split_data(processed_dir, split_dir, train_p=0.7, val_p=0.15):
    """
    Splits the processed binary data into train, val, and test sets.
    """
    for split in ['train', 'val', 'test']:
        for label in ['fresh', 'spoiled']:
            os.makedirs(os.path.join(split_dir, split, label), exist_ok=True)

    for label in ['fresh', 'spoiled']:
        label_dir = os.path.join(processed_dir, label)
        images = os.listdir(label_dir)
        random.shuffle(images)

        n = len(images)
        train_end = int(n * train_p)
        val_end = int(n * (train_p + val_p))

        splits = {
            'train': images[:train_end],
            'val': images[train_end:val_end],
            'test': images[val_end:]
        }

        for split_name, split_images in splits.items():
            for img in split_images:
                shutil.copy(
                    os.path.join(label_dir, img),
                    os.path.join(split_dir, split_name, label, img)
                )

if __name__ == "__main__":
    SOURCE = "./data/raw/dataset"
    PROCESSED = "./data/processed"
    FINAL = "./data/split"
    
    # Strictly fruits and vegetables as requested "for now"
    CLASSES = ['fruits', 'vegetables']
    
    print(f"Restructuring data for: {CLASSES}...")
    prepare_data(SOURCE, PROCESSED, CLASSES)
    
    print("Splitting data into train/val/test...")
    split_data(PROCESSED, FINAL)
    
    print("Data preparation complete.")