src/gradio_interface.py

"""
Gradio interface functions for the Mosaic Generator.
"""

import gradio as gr
import numpy as np
from PIL import Image
import time
from typing import Tuple, Dict, List

from .config import Config, Implementation, MatchSpace
from .pipeline import MosaicPipeline
from .metrics import calculate_comprehensive_metrics, interpret_metrics


def create_default_config(
    grid_size: int = 32,
    tile_size: int = 32,
    output_width: int = 768,
    output_height: int = 768,
    color_matching: str = "Lab (perceptual)",
    use_uniform_quantization: bool = False,
    quantization_levels: int = 8,
    use_kmeans_quantization: bool = False,
    kmeans_colors: int = 8,
    normalize_tile_brightness: bool = False
) -> Config:
    """Create configuration from Gradio interface parameters."""
    
    # Convert string parameters to enums
    match_space = MatchSpace.LAB if color_matching == "Lab (perceptual)" else MatchSpace.RGB
    
    return Config(
        grid=grid_size,
        tile_size=tile_size,
        out_w=output_width,
        out_h=output_height,
        impl=Implementation.VECT,  # Always use vectorized
        match_space=match_space,
        use_uniform_q=use_uniform_quantization,
        q_levels=quantization_levels,
        use_kmeans_q=use_kmeans_quantization,
        k_colors=kmeans_colors,
        tile_norm_brightness=normalize_tile_brightness
    )


def generate_mosaic(
    image: Image.Image,
    grid_size: int,
    tile_size: int,
    output_width: int,
    output_height: int,
    color_matching: str,
    use_uniform_quantization: bool,
    quantization_levels: int,
    use_kmeans_quantization: bool,
    kmeans_colors: int,
    normalize_tile_brightness: bool,
    progress=gr.Progress()
) -> Tuple[Image.Image, Image.Image, str, str]:
    """
    Generate mosaic from input image with given parameters.
    
    Returns:
        Tuple of (mosaic_image, processed_image, metrics_text, timing_text)
    """
    if image is None:
        return None, None, "Please upload an image.", ""
    
    try:
        # Create configuration
        config = create_default_config(
            grid_size, tile_size, output_width, output_height,
            color_matching, use_uniform_quantization,
            quantization_levels, use_kmeans_quantization, kmeans_colors,
            normalize_tile_brightness
        )
        
        # Create pipeline
        pipeline = MosaicPipeline(config)
        
        # Update progress
        progress(0.1, desc="Initializing pipeline...")
        
        # Run pipeline
        progress(0.2, desc="Loading tiles (first time only)...")
        progress(0.4, desc="Generating mosaic...")
        results = pipeline.run_full_pipeline(image)
        
        progress(0.7, desc="Calculating metrics...")
        
        # Extract results
        mosaic_img = results['outputs']['mosaic']
        processed_img = results['outputs']['processed_image']
        
        # Format metrics
        metrics = results['metrics']
        interpretations = results['metrics_interpretation']
        
        metrics_text = f"""
**Quality Metrics:**
- **MSE (Mean Squared Error):** {metrics['mse']:.6f} - {interpretations['mse']}
- **PSNR (Peak Signal-to-Noise Ratio):** {metrics['psnr']:.2f} dB - {interpretations['psnr']}
- **SSIM (Structural Similarity):** {metrics['ssim']:.4f} - {interpretations['ssim']}
- **RMSE (Root Mean Squared Error):** {metrics['rmse']:.6f}
- **MAE (Mean Absolute Error):** {metrics['mae']:.6f}

**Color Analysis:**
- **Color MSE:** {metrics['color_mse']:.6f}
- **Histogram Correlation:** {metrics['histogram_correlation']:.4f}
        """
        
        # Format timing information
        timing = results['timing']
        timing_text = f"""
**Processing Times:**
- **Preprocessing:** {timing['preprocessing']:.3f} seconds
- **Grid Analysis:** {timing['grid_analysis']:.3f} seconds  
- **Tile Mapping:** {timing['tile_mapping']:.3f} seconds
- **Total Time:** {timing['total']:.3f} seconds

**Configuration:**
- **Grid Size:** {config.grid}x{config.grid} ({config.grid**2} tiles total)
- **Tile Size:** {config.tile_size}x{config.tile_size} pixels
- **Output Resolution:** {mosaic_img.width}x{mosaic_img.height}
- **Implementation:** {config.impl.value}
- **Color Matching:** {config.match_space.value}
        """
        
        progress(1.0, desc="Complete!")
        
        return mosaic_img, processed_img, metrics_text, timing_text
        
    except Exception as e:
        error_msg = f"Error generating mosaic: {str(e)}"
        print(error_msg)
        return None, None, error_msg, ""




def benchmark_grid_sizes(
    image: Image.Image,
    grid_sizes: str,
    progress=gr.Progress()
) -> str:
    """Benchmark different grid sizes."""
    if image is None:
        return "Please upload an image for benchmarking."
    
    try:
        # Parse grid sizes
        sizes = [int(x.strip()) for x in grid_sizes.split(',')]
        
        results = []
        total_tests = len(sizes)
        
        for i, grid_size in enumerate(sizes):
            progress((i + 1) / total_tests, desc=f"Testing grid size {grid_size}x{grid_size}...")
            
            config = create_default_config(grid_size, 32, 768, 768)
            pipeline = MosaicPipeline(config)
            
            start_time = time.time()
            pipeline_results = pipeline.run_full_pipeline(image)
            processing_time = time.time() - start_time
            
            results.append({
                'grid_size': grid_size,
                'processing_time': processing_time,
                'total_tiles': grid_size * grid_size,
                'tiles_per_second': (grid_size * grid_size) / processing_time,
                'mse': pipeline_results['metrics']['mse'],
                'ssim': pipeline_results['metrics']['ssim']
            })
        
        # Generate report
        report = "**Grid Size Performance Analysis:**\n\n"
        
        for result in results:
            report += f"**Grid {result['grid_size']}x{result['grid_size']}:**\n"
            report += f"- Processing Time: {result['processing_time']:.3f}s\n"
            report += f"- Total Tiles: {result['total_tiles']}\n"
            report += f"- Tiles per Second: {result['tiles_per_second']:.1f}\n"
            report += f"- MSE: {result['mse']:.6f}\n"
            report += f"- SSIM: {result['ssim']:.4f}\n\n"
        
        # Scaling analysis
        if len(results) >= 2:
            first = results[0]
            last = results[-1]
            tile_ratio = last['total_tiles'] / first['total_tiles']
            time_ratio = last['processing_time'] / first['processing_time']
            
            report += "**Scaling Analysis:**\n"
            report += f"- Tile increase ratio: {tile_ratio:.2f}x\n"
            report += f"- Time increase ratio: {time_ratio:.2f}x\n"
            report += f"- Scaling efficiency: {tile_ratio/time_ratio:.2f}\n"
            report += f"- Linear scaling: {'Yes' if abs(time_ratio - tile_ratio) / tile_ratio < 0.1 else 'No'}\n"
        
        return report
        
    except Exception as e:
        return f"Error during grid size benchmarking: {str(e)}"


def create_interface():
    """Create the Gradio interface."""
    
    with gr.Blocks(title="Mosaic Generator", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🎨 Mosaic Generator")
        gr.Markdown("Generate beautiful mosaic-style images from your photos using advanced image processing techniques.")
        
        with gr.Tab("Generate Mosaic"):
            with gr.Row():
                with gr.Column(scale=1):
                    # Input controls
                    gr.Markdown("## Upload & Configure")
                    
                    input_image = gr.Image(
                        type="pil",
                        label="Upload Image",
                        height=300
                    )
                    
                    with gr.Accordion("Basic Settings", open=True):
                        grid_size = gr.Slider(
                            minimum=8, maximum=128, step=8, value=32,
                            label="Grid Size (N×N tiles)"
                        )
                        tile_size = gr.Slider(
                            minimum=4, maximum=64, step=4, value=32,
                            label="Tile Size (pixels)"
                        )
                        output_width = gr.Slider(
                            minimum=256, maximum=1024, step=64, value=768,
                            label="Output Width"
                        )
                        output_height = gr.Slider(
                            minimum=256, maximum=1024, step=64, value=768,
                            label="Output Height"
                        )
                    
                    with gr.Accordion("Advanced Settings", open=False):
                        color_matching = gr.Radio(
                            choices=["Lab (perceptual)", "RGB (euclidean)"],
                            value="Lab (perceptual)",
                            label="Color Matching Space"
                        )
                        
                        gr.Markdown("**Color Quantization:**")
                        use_uniform_quantization = gr.Checkbox(
                            label="Use Uniform Quantization",
                            value=False
                        )
                        quantization_levels = gr.Slider(
                            minimum=4, maximum=16, step=2, value=8,
                            label="Quantization Levels",
                            visible=True
                        )
                        
                        use_kmeans_quantization = gr.Checkbox(
                            label="Use K-means Quantization",
                            value=False
                        )
                        kmeans_colors = gr.Slider(
                            minimum=4, maximum=32, step=2, value=8,
                            label="K-means Colors"
                        )
                        
                        normalize_tile_brightness = gr.Checkbox(
                            label="Normalize Tile Brightness",
                            value=False
                        )
                    
                    generate_btn = gr.Button("Generate Mosaic", variant="primary", size="lg")
                
                with gr.Column(scale=2):
                    # Output display
                    gr.Markdown("## Results")
                    
                    with gr.Row():
                        mosaic_output = gr.Image(
                            label="Generated Mosaic",
                            height=400
                        )
                        processed_output = gr.Image(
                            label="Processed Input",
                            height=400
                        )
                    
                    with gr.Row():
                        metrics_output = gr.Markdown(label="Quality Metrics")
                        timing_output = gr.Markdown(label="Processing Information")
        
        with gr.Tab("Performance Analysis"):
            gr.Markdown("## Performance Benchmarking")
            
            with gr.Row():
                with gr.Column():
                    benchmark_image = gr.Image(
                        type="pil",
                        label="Image for Benchmarking",
                        height=200
                    )
                    
                    gr.Markdown("### Grid Size Benchmarking")
                    grid_sizes_input = gr.Textbox(
                        value="16,32,48,64",
                        label="Grid Sizes (comma-separated)",
                        placeholder="16,32,48,64"
                    )
                    benchmark_grid_btn = gr.Button("Benchmark Grid Sizes", variant="secondary")
                
                with gr.Column():
                    benchmark_output = gr.Markdown(label="Benchmark Results")
        
        with gr.Tab("About"):
            gr.Markdown("""
            ## About the Mosaic Generator
            
            This application implements a complete mosaic generation pipeline with the following features:
            
            **Note**: The first time you generate a mosaic, it will load tiles from the Hugging Face dataset. This may take a few moments, but subsequent generations will be much faster as tiles are cached.
            
            ### Core Functionality
            - **Image Preprocessing**: Resize and crop images to fit grid requirements
            - **Color Quantization**: Optional uniform and K-means quantization
            - **Grid Analysis**: Vectorized operations for efficient processing
            - **Tile Mapping**: Replace grid cells with matching image tiles
            - **Quality Metrics**: MSE, PSNR, SSIM, and color similarity analysis
            
            ### Performance Features
            - **Vectorized Operations**: NumPy-based efficient processing
            - **Grid Size Benchmarking**: Performance analysis across different resolutions
            - **Real-time Metrics**: Processing time and quality measurements
            
            ### Technical Details
            - Uses Hugging Face datasets for tile sources
            - Supports LAB and RGB color space matching
            - Configurable grid sizes from 8×8 to 128×128
            - Adjustable tile sizes and output resolutions
            
            ### Assignment Requirements Met
            ✅ Image selection and preprocessing  
            ✅ Grid division and thresholding  
            ✅ Vectorized NumPy operations  
            ✅ Tile mapping and replacement  
            ✅ Gradio interface with parameter controls  
            ✅ Similarity metrics (MSE, SSIM)  
            ✅ Performance analysis and benchmarking  
            """)
        
        # Event handlers
        generate_btn.click(
            fn=generate_mosaic,
            inputs=[
                input_image, grid_size, tile_size, output_width, output_height,
                color_matching, use_uniform_quantization,
                quantization_levels, use_kmeans_quantization, kmeans_colors,
                normalize_tile_brightness
            ],
            outputs=[mosaic_output, processed_output, metrics_output, timing_output]
        )
        
        benchmark_grid_btn.click(
            fn=benchmark_grid_sizes,
            inputs=[benchmark_image, grid_sizes_input],
            outputs=[benchmark_output]
        )
        
        # Update visibility of quantization controls
        use_uniform_quantization.change(
            fn=lambda x: gr.Slider(visible=x),
            inputs=[use_uniform_quantization],
            outputs=[quantization_levels]
        )
        
        use_kmeans_quantization.change(
            fn=lambda x: gr.Slider(visible=x),
            inputs=[use_kmeans_quantization],
            outputs=[kmeans_colors]
        )
    
    return demo