#!/usr/bin/env python3
"""
Example script demonstrating mosaic generation functionality.
"""

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import os

from src.config import Config
from src.pipeline import MosaicPipeline


def create_sample_image(size=(512, 512)):
    """Create a sample image with gradients and patterns."""
    img_array = np.zeros((size[1], size[0], 3), dtype=np.float32)
    
    # Create gradient patterns
    for y in range(size[1]):
        for x in range(size[0]):
            # Red gradient
            img_array[y, x, 0] = x / size[0]
            
            # Green gradient  
            img_array[y, x, 1] = y / size[1]
            
            # Blue pattern
            img_array[y, x, 2] = (x + y) / (size[0] + size[1])
    
    # Add geometric shapes
    center_x, center_y = size[0] // 2, size[1] // 2
    radius = min(size) // 4
    
    for y in range(size[1]):
        for x in range(size[0]):
            # Circle
            dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
            if dist < radius:
                img_array[y, x] = [1.0, 0.5, 0.2]  # Orange circle
    
    return Image.fromarray((img_array * 255).astype(np.uint8))


def demonstrate_mosaic_generation():
    """Demonstrate mosaic generation with different configurations."""
    
    print("🎨 Mosaic Generator Example")
    print("=" * 40)
    
    # Create sample image
    print("Creating sample image...")
    sample_img = create_sample_image()
    os.makedirs("images", exist_ok=True)
    sample_img.save("images/sample_input.png")
    print("✅ Sample image saved to images/sample_input.png")
    
    # Test different grid sizes
    grid_sizes = [16, 32, 48]
    
    for grid_size in grid_sizes:
        print(f"\nGenerating mosaic with {grid_size}x{grid_size} grid...")
        
        # Create configuration
        config = Config(
            grid=grid_size,
            tile_size=32,
            out_w=512,
            out_h=512
        )
        
        # Create pipeline
        pipeline = MosaicPipeline(config)
        
        # Generate mosaic
        results = pipeline.run_full_pipeline(sample_img)
        
        # Save results
        mosaic_img = results['outputs']['mosaic']
        processed_img = results['outputs']['processed_image']
        
        mosaic_img.save(f"images/mosaic_{grid_size}x{grid_size}.png")
        processed_img.save(f"images/processed_{grid_size}x{grid_size}.png")
        
        # Print metrics
        metrics = results['metrics']
        timing = results['timing']
        
        print(f"✅ Mosaic saved to images/mosaic_{grid_size}x{grid_size}.png")
        print(f"   Processing time: {timing['total']:.3f}s")
        print(f"   MSE: {metrics['mse']:.6f}")
        print(f"   SSIM: {metrics['ssim']:.4f}")
    
    # Test implementation comparison
    print(f"\nComparing implementations...")
    
    config_vect = Config(grid=32, tile_size=32, out_w=512, out_h=512, impl="Vectorised")
    config_loop = Config(grid=32, tile_size=32, out_w=512, out_h=512, impl="Loops")
    
    pipeline_vect = MosaicPipeline(config_vect)
    pipeline_loop = MosaicPipeline(config_loop)
    
    import time
    
    # Time vectorized
    start = time.time()
    results_vect = pipeline_vect.run_full_pipeline(sample_img)
    time_vect = time.time() - start
    
    # Time loop-based
    start = time.time()
    results_loop = pipeline_loop.run_full_pipeline(sample_img)
    time_loop = time.time() - start
    
    speedup = time_loop / time_vect if time_vect > 0 else 0
    
    print(f"✅ Vectorized: {time_vect:.3f}s")
    print(f"✅ Loop-based: {time_loop:.3f}s")
    print(f"✅ Speedup: {speedup:.2f}x")
    
    # Create comparison visualization
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    
    # Original image
    axes[0, 0].imshow(sample_img)
    axes[0, 0].set_title("Original Image")
    axes[0, 0].axis('off')
    
    # 16x16 mosaic
    mosaic_16 = Image.open("images/mosaic_16x16.png")
    axes[0, 1].imshow(mosaic_16)
    axes[0, 1].set_title("16×16 Grid Mosaic")
    axes[0, 1].axis('off')
    
    # 32x32 mosaic
    mosaic_32 = Image.open("images/mosaic_32x32.png")
    axes[0, 2].imshow(mosaic_32)
    axes[0, 2].set_title("32×32 Grid Mosaic")
    axes[0, 2].axis('off')
    
    # 48x48 mosaic
    mosaic_48 = Image.open("images/mosaic_48x48.png")
    axes[1, 0].imshow(mosaic_48)
    axes[1, 0].set_title("48×48 Grid Mosaic")
    axes[1, 0].axis('off')
    
    # Vectorized result
    axes[1, 1].imshow(results_vect['outputs']['mosaic'])
    axes[1, 1].set_title(f"Vectorized ({time_vect:.3f}s)")
    axes[1, 1].axis('off')
    
    # Loop-based result
    axes[1, 2].imshow(results_loop['outputs']['mosaic'])
    axes[1, 2].set_title(f"Loop-based ({time_loop:.3f}s)")
    axes[1, 2].axis('off')
    
    plt.tight_layout()
    plt.savefig("images/mosaic_comparison.png", dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"\n✅ Comparison visualization saved to images/mosaic_comparison.png")
    
    print(f"\n🎉 Example complete! Check the 'images' folder for results.")


if __name__ == "__main__":
    demonstrate_mosaic_generation()