"""Generate sample images for SIREN super-resolution demo."""
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os


def create_gradient_image(size=(512, 512)):
    """Create a colorful gradient image."""
    width, height = size
    img = np.zeros((height, width, 3), dtype=np.uint8)

    for y in range(height):
        for x in range(width):
            img[y, x, 0] = int(255 * x / width)  # Red gradient
            img[y, x, 1] = int(255 * y / height)  # Green gradient
            img[y, x, 2] = int(255 * (1 - x / width) * (1 - y / height))  # Blue

    return Image.fromarray(img)


def create_pattern_image(size=(512, 512)):
    """Create an image with geometric patterns."""
    width, height = size
    img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(img)

    # Draw concentric circles
    center_x, center_y = width // 2, height // 2
    colors = ['red', 'orange', 'yellow', 'green', 'blue', 'purple']

    for i, color in enumerate(colors):
        radius = (len(colors) - i) * 40
        draw.ellipse(
            [center_x - radius, center_y - radius,
             center_x + radius, center_y + radius],
            outline=color,
            width=3
        )

    # Draw grid
    for i in range(0, width, 50):
        draw.line([(i, 0), (i, height)], fill='lightgray', width=1)
    for i in range(0, height, 50):
        draw.line([(0, i), (width, i)], fill='lightgray', width=1)

    return img


def create_checkerboard_image(size=(512, 512), square_size=32):
    """Create a checkerboard pattern with gradients."""
    width, height = size
    img = Image.new('RGB', (width, height))
    pixels = img.load()

    for y in range(height):
        for x in range(width):
            square_x = x // square_size
            square_y = y // square_size

            # Checkerboard pattern
            if (square_x + square_y) % 2 == 0:
                # Light square with gradient
                intensity = int(200 + 55 * (x % square_size) / square_size)
                pixels[x, y] = (intensity, intensity, intensity)
            else:
                # Dark square with color gradient
                r = int(100 * (x % square_size) / square_size)
                g = int(100 * (y % square_size) / square_size)
                b = 150
                pixels[x, y] = (r, g, b)

    return img


def create_mandala_image(size=(512, 512)):
    """Create a mandala-like pattern."""
    width, height = size
    img = np.zeros((height, width, 3), dtype=np.uint8)

    center_x, center_y = width // 2, height // 2

    for y in range(height):
        for x in range(width):
            dx = x - center_x
            dy = y - center_y

            distance = np.sqrt(dx**2 + dy**2)
            angle = np.arctan2(dy, dx)

            # Create radial pattern
            r = int(127 + 127 * np.sin(distance / 20 + angle * 5))
            g = int(127 + 127 * np.cos(distance / 30 - angle * 3))
            b = int(127 + 127 * np.sin(distance / 40 + angle * 7))

            img[y, x] = [r, g, b]

    return Image.fromarray(img)


def main():
    """Generate all sample images."""
    os.makedirs('samples', exist_ok=True)

    print("Generating sample images...")

    # Generate different sample images
    samples = {
        'cat.jpg': create_mandala_image(),
        'landscape.jpg': create_gradient_image(),
        'portrait.jpg': create_pattern_image(),
        'checkerboard.jpg': create_checkerboard_image(),
    }

    for filename, image in samples.items():
        filepath = os.path.join('samples', filename)
        image.save(filepath, quality=95)
        print(f"Created: {filepath}")

    print("\n✓ All sample images created successfully!")
    print("\nYou can replace these with your own high-resolution images.")


if __name__ == "__main__":
    main()