"""
create_synthetic_dataset.py — Generate synthetic training data locally.

Creates simple image editing pairs for quick model training and testing.
"""

import os
import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
import pathlib


def create_synthetic_pair(idx: int, size: int = 256):
    """Create a synthetic image editing pair (input, edited output)."""
    # Create a base image with random shapes and colors
    img = Image.new('RGB', (size, size), color=(255, 255, 255))
    draw = ImageDraw.Draw(img)
    
    # Draw random shapes
    np.random.seed(idx)  # Deterministic for reproducibility
    for _ in range(5):
        x0, y0 = np.random.randint(0, size, 2)
        x1, y1 = x0 + np.random.randint(20, 100), y0 + np.random.randint(20, 100)
        color = tuple(np.random.randint(0, 256, 3))
        draw.rectangle([x0, y0, x1, y1], fill=color, outline=(0, 0, 0), width=2)
    
    # Add text
    draw.text((10, 10), f"Sample {idx}", fill=(0, 0, 0))
    
    # Create edited version by applying a transformation
    edited = img.copy()
    edited = ImageEnhance.Color(edited).enhance(1.3)  # Increase saturation
    edited = ImageEnhance.Brightness(edited).enhance(0.9)  # Slight darkening
    edited = edited.filter(ImageFilter.SMOOTH)
    
    return img, edited


def create_dataset(num_samples: int = 100, output_dir: str = "./data/synthetic"):
    """Create and save synthetic training pairs in LocalEditDataset format."""
    os.makedirs(output_dir, exist_ok=True)
    
    # Create directory structure expected by LocalEditDataset
    input_dir = os.path.join(output_dir, 'input')
    edited_dir = os.path.join(output_dir, 'edited')
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(edited_dir, exist_ok=True)
    
    prompts = []
    prompt_templates = [
        'enhance image quality',
        'increase saturation',
        'brighten the image',
        'apply artistic filter',
        'improve contrast',
        'warm color tone',
        'cool color tone',
        'add detail',
        'smooth texture',
        'vibrant colors',
    ]
    
    for i in range(num_samples):
        img, edited = create_synthetic_pair(i)
        
        img_path = os.path.join(input_dir, f'{i:06d}.png')
        edited_path = os.path.join(edited_dir, f'{i:06d}.png')
        
        img.save(img_path)
        edited.save(edited_path)
        
        prompt = prompt_templates[i % len(prompt_templates)]
        prompts.append(prompt)
        
        if (i + 1) % 20 == 0:
            print(f'  Created {i+1}/{num_samples} pairs')
    
    # Save prompts file
    prompts_path = os.path.join(output_dir, 'prompts.txt')
    with open(prompts_path, 'w') as f:
        f.write('\n'.join(prompts))
    
    print(f'✓ Dataset created at {output_dir}')
    print(f'  - {num_samples} input images')
    print(f'  - {num_samples} edited images')
    print(f'  - prompts.txt with instructions')


if __name__ == "__main__":
    print("Creating synthetic training dataset...")
    create_dataset(num_samples=100, output_dir="./data/synthetic")