"""
Create dummy images for training
Generate synthetic infrastructure images
"""

from PIL import Image, ImageDraw
import random
import os


def create_dummy_image(filename, width=640, height=480, num_objects=3):
    """Create a dummy infrastructure image"""
    # Create base image
    img = Image.new('RGB', (width, height), color=(135, 206, 235))  # Sky blue
    draw = ImageDraw.Draw(img)
    
    # Add ground
    draw.rectangle([0, height*0.7, width, height], fill=(100, 80, 60))  # Brown ground
    
    # Add random objects
    colors = [
        (255, 0, 0),    # Red
        (0, 255, 0),    # Green
        (0, 0, 255),    # Blue
        (255, 255, 0),  # Yellow
        (255, 0, 255),  # Magenta
    ]
    
    for i in range(num_objects):
        x = random.randint(50, width-100)
        y = random.randint(50, height-150)
        w = random.randint(30, 100)
        h = random.randint(30, 100)
        color = colors[i % len(colors)]
        
        # Draw object
        draw.rectangle([x, y, x+w, y+h], fill=color, outline=(0, 0, 0), width=2)
        
        # Add label
        draw.text((x, y-15), f"Object {i+1}", fill=(0, 0, 0))
    
    # Save image
    img.save(filename)
    return filename


def create_training_dataset(base_path="/tmp/iom_training_data", num_train=50, num_val=10, num_test=10):
    """Create complete training dataset with images"""
    print("=== Creating Dummy Training Images ===\n")
    
    # Create directories
    splits = {
        "train": num_train,
        "val": num_val,
        "test": num_test
    }
    
    for split, count in splits.items():
        img_dir = os.path.join(base_path, "images", split)
        os.makedirs(img_dir, exist_ok=True)
        
        print(f"Creating {count} {split} images...")
        for i in range(count):
            filename = os.path.join(img_dir, f"{split}_{i:04d}.jpg")
            create_dummy_image(filename, num_objects=random.randint(1, 5))
    
    print(f"\nDataset created:")
    print(f"  Train: {num_train} images")
    print(f"  Val: {num_val} images")
    print(f"  Test: {num_test} images")
    
    return base_path


if __name__ == '__main__':
    create_training_dataset()