"""
Enhanced Gradio Interface for Image Mosaic Generator

Features:
- Modern, intuitive UI with tabs
- Real-time preview
- Performance analysis tools
- Optimization comparisons
"""

import gradio as gr
import numpy as np
from typing import Tuple, Optional

# Import modular components
from mosaic_generator import (
    TileManager,
    ImageProcessor,
    MosaicBuilder,
    MetricsCalculator,
    PerformanceBenchmark,
    DEFAULT_TILE_SIZE,
    MIN_GRID_SIZE,
    MAX_GRID_SIZE,
    DEFAULT_GRID_SIZE
)

# Initialize components globally
print("🔧 Initializing Mosaic Generator...")
tile_manager = TileManager(tile_size=DEFAULT_TILE_SIZE)
# Pre-generate tiles
print("🎨 Generating tile set...")
tile_manager.generate_procedural_tiles()
print(f"✅ Generated {tile_manager.get_tile_count()} tiles")
image_processor = ImageProcessor(tile_size=DEFAULT_TILE_SIZE)
mosaic_builder = MosaicBuilder(tile_manager, image_processor)
benchmark = PerformanceBenchmark(mosaic_builder)


# ===================================
# Main Generation Functions
# ===================================

def generate_mosaic_fast(image: np.ndarray, grid_size: int,
                         use_kdtree: bool) -> Tuple[Optional[np.ndarray], str]:
    """
    Fast mosaic generation with quality metrics.
    """
    if image is None:
        return None, "⚠️ Please upload an image first."

    # Generate mosaic
    mosaic, duration = mosaic_builder.create_mosaic_optimized(
        image, grid_size, grid_size, use_kdtree=use_kdtree
    )

    if mosaic is None:
        return None, "❌ Failed to generate mosaic."

    # Calculate metrics
    metrics = MetricsCalculator.calculate_all_metrics(image, mosaic)
    stats_text = MetricsCalculator.format_metrics(metrics, duration, grid_size)

    # Add optimization info
    opt_method = "KD-Tree" if use_kdtree else "Brute Force"
    stats_text += f"\n\n**Matching Algorithm:** {opt_method}"

    return mosaic, stats_text


def run_grid_benchmark(image: np.ndarray) -> str:
    """
    Run comprehensive grid size benchmark.
    """
    if image is None:
        return "⚠️ Please upload an image first."

    return benchmark.run_grid_size_benchmark(image)


def run_optimization_comparison(image: np.ndarray, grid_size: int) -> str:
    """
    Compare optimization strategies.
    """
    if image is None:
        return "⚠️ Please upload an image first."

    return benchmark.run_optimization_comparison(image, grid_size)


def generate_side_by_side(image: np.ndarray,
                          grid_size: int) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], str]:
    """
    Generate side-by-side comparison with original.
    """
    if image is None:
        return None, None, "⚠️ Please upload an image first."

    # Generate mosaic
    mosaic, duration = mosaic_builder.create_mosaic_optimized(
        image, grid_size, grid_size
    )

    if mosaic is None:
        return None, None, "❌ Failed to generate mosaic."

    # Calculate metrics
    metrics = MetricsCalculator.calculate_all_metrics(image, mosaic)

    comparison_text = f"### Side-by-Side Comparison\n\n"
    comparison_text += f"**Grid Size:** {grid_size}×{grid_size}\n"
    comparison_text += f"**Processing Time:** {duration:.4f}s\n"
    comparison_text += f"**SSIM Score:** {metrics['ssim']:.4f}\n"
    comparison_text += f"**MSE:** {metrics['mse']:.2f}"

    return image, mosaic, comparison_text


# ===================================
# Gradio Interface
# ===================================

# Custom CSS for better styling
custom_css = """
.gradio-container {
    font-family: 'Inter', sans-serif;
}

.main-title {
    text-align: center;
    color: #2563eb;
    margin-bottom: 1rem;
}

.stat-box {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 1rem;
    border-radius: 8px;
    color: white;
    margin: 0.5rem 0;
}

.metric-good {
    color: #10b981;
    font-weight: bold;
}

.metric-bad {
    color: #ef4444;
    font-weight: bold;
}
"""


def get_example_images():
    return [
        ["examples/cat.png"],
    ]



with gr.Blocks(title="🎨 Image Mosaic Generator") as demo:
    # Header
    gr.Markdown(
        """
        # 🎨 Image Mosaic Generator
        ### Transform images into stunning mosaics with optimized algorithms

        Upload an image and explore different mosaic configurations. This tool uses vectorized NumPy operations
        and optional KD-Tree optimization for fast, high-quality results.
        """,
        elem_classes="main-title"
    )

    # Main tabs
    with gr.Tabs():
        # ===== TAB 1: Quick Generation =====
        with gr.Tab("🚀 Quick Generate"):
            gr.Markdown("Generate a mosaic quickly with your preferred settings.")

            with gr.Row():
                with gr.Column(scale=1):
                    input_image_quick = gr.Image(
                        label="📤 Upload Image",
                        type="numpy",
                        height=400,
                    )
                    gr.Examples(
                        examples=get_example_images(),
                        inputs=input_image_quick,
                        label="Click to load example image"
                    )

                    grid_slider_quick = gr.Slider(
                        minimum=MIN_GRID_SIZE,
                        maximum=MAX_GRID_SIZE,
                        value=DEFAULT_GRID_SIZE,
                        step=4,
                        label="🎚️ Grid Size (tiles per dimension)",
                        info="Higher values = more detail but slower"
                    )

                    use_kdtree_check = gr.Checkbox(
                        label="⚡ Use KD-Tree Optimization",
                        value=False,
                        info="Faster for large grids (64+)"
                    )

                    generate_btn_quick = gr.Button(
                        "🎨 Generate Mosaic",
                        variant="primary",
                        size="lg"
                    )

                with gr.Column(scale=1):
                    output_image_quick = gr.Image(
                        label="🖼️ Mosaic Result",
                        height=400
                    )

                    stats_quick = gr.Markdown(
                        "Upload an image and click 'Generate Mosaic' to begin.",
                        label="📊 Statistics"
                    )

            generate_btn_quick.click(
                fn=generate_mosaic_fast,
                inputs=[input_image_quick, grid_slider_quick, use_kdtree_check],
                outputs=[output_image_quick, stats_quick]
            )
            input_image_quick.change(
                fn=generate_mosaic_fast,
                inputs=[input_image_quick, grid_slider_quick, use_kdtree_check],
                outputs=[output_image_quick, stats_quick]
            )

        # ===== TAB 2: Benchmarks =====
        with gr.Tab("📊 Benchmarks"):
            gr.Markdown(
                """
                ### Performance Analysis Tools

                Run comprehensive benchmarks to analyze performance across different configurations.
                """
            )

            with gr.Row():
                with gr.Column():
                    input_image_bench = gr.Image(
                        label="📤 Upload Image",
                        type="numpy"
                    )
                    gr.Examples(
                        examples=get_example_images(),
                        inputs=input_image_bench,
                        label="Click to load example image"
                    )

                    gr.Markdown("#### Grid Size Benchmark")
                    gr.Markdown("Compare performance across multiple grid sizes (16, 32, 64, 128)")

                    bench_grid_btn = gr.Button(
                        "🏃 Run Grid Benchmark",
                        variant="secondary"
                    )

                    gr.Markdown("#### Optimization Strategy Comparison")
                    gr.Markdown("Compare different algorithmic optimizations")

                    grid_size_opt = gr.Slider(
                        minimum=MIN_GRID_SIZE,
                        maximum=MAX_GRID_SIZE,
                        value=64,
                        step=8,
                        label="Grid Size for Comparison"
                    )

                    bench_opt_btn = gr.Button(
                        "⚡ Compare Optimizations",
                        variant="secondary"
                    )

                with gr.Column():
                    bench_results = gr.Markdown(
                        "Upload an image and run benchmarks to see results.",
                        label="📈 Results"
                    )

            bench_grid_btn.click(
                fn=run_grid_benchmark,
                inputs=[input_image_bench],
                outputs=[bench_results]
            )

            bench_opt_btn.click(
                fn=run_optimization_comparison,
                inputs=[input_image_bench, grid_size_opt],
                outputs=[bench_results]
            )
            input_image_quick.change(
                fn=run_grid_benchmark,
                inputs=[input_image_bench],
                outputs=[bench_results]
            )
            input_image_bench.change(
                fn=run_optimization_comparison,
                inputs=[input_image_bench, grid_size_opt],
                outputs=[bench_results]
            )

        # ===== TAB 4: Info =====
        with gr.Tab("ℹ️ Info"):
            gr.Markdown(
                """
                ## About Image Mosaic Generator

                This tool creates artistic mosaics by reconstructing input images using small colored tiles.

                ### Features

                - **🎨 Procedural Tile Generation**: Creates 216 unique tiles spanning the RGB color spectrum
                - **⚡ Optimized Algorithms**: Multiple optimization strategies including vectorized NumPy operations and KD-Tree matching
                - **📊 Quality Metrics**: MSE (Mean Squared Error) and SSIM (Structural Similarity Index) 
                - **🔬 Performance Analysis**: Comprehensive benchmarking tools

                ### How It Works

                1. **Image Preprocessing**: Input image is resized to match the grid dimensions
                2. **Color Extraction**: Average color is computed for each grid cell
                3. **Tile Matching**: Best-matching tile is found for each cell using distance metrics
                4. **Reconstruction**: Tiles are assembled to create the final mosaic

                ### Optimization Strategies

                - **Vectorized Operations**: NumPy array operations replace nested loops
                - **KD-Tree Matching**: Spatial data structure for O(log n) nearest neighbor queries
                - **Efficient Resizing**: PIL's BOX filter for optimal color averaging
                - **Caching**: Pre-computed tile features reduce redundant calculations

                ### Quality Metrics

                - **MSE (Mean Squared Error)**: Measures pixel-level differences (lower is better)
                - **SSIM (Structural Similarity)**: Measures perceptual similarity (higher is better, 0-1 range)

                ### Tips for Best Results

                - Start with **32×32 grid** for balanced quality and speed
                - Use **64×64 or higher** for detailed images
                - Enable **KD-Tree optimization** for grids larger than 64×64
                - Images with **clear colors** work best

                ### Performance

                Typical processing times (1024×1024 image):
                - 32×32 grid: ~0.05s
                - 64×64 grid: ~0.15s
                - 128×128 grid: ~0.5s

                ---

                **Version 2.0.0** | Built with NumPy, Pillow, scikit-image, and Gradio
                """
            )

    # Footer
    gr.Markdown(
        """
        ---
        💡 **Tip**: For best results, try different grid sizes and compare the quality metrics!
        """,
        elem_classes="footer"
    )

if __name__ == "__main__":
    print("\n🚀 Launching Image Mosaic Generator...")
    print(f"📦 Tile Bank: {tile_manager.get_tile_count()} tiles")
    print(f"🎯 Default Grid Size: {DEFAULT_GRID_SIZE}×{DEFAULT_GRID_SIZE}")
    print(f"⚙️ Tile Size: {DEFAULT_TILE_SIZE}×{DEFAULT_TILE_SIZE} pixels\n")

    demo.launch(
        share=True,
        show_error=True,
        server_name="127.0.0.1",
        server_port=7860
    )