app.py

"""
app.py
Deployment-only Gradio interface with Numba optimizations and grid visualization.
Uses only pre-built cache files for fast cloud deployment.
Metrics calculated AFTER images are displayed for faster user experience.
"""

import gradio as gr
import numpy as np
import cv2
from pathlib import Path
import time
import pickle
import warnings
import sys
from typing import Tuple, Dict, Optional, List
from dataclasses import dataclass
from sklearn.cluster import MiniBatchKMeans
from sklearn.metrics.pairwise import euclidean_distances

# Add Numba folder to path for imports
SCRIPT_DIR = Path(__file__).parent
NUMBA_DIR = SCRIPT_DIR / "Numba_Scripts"
if NUMBA_DIR.exists() and str(NUMBA_DIR) not in sys.path:
    sys.path.insert(0, str(NUMBA_DIR))

# Try to import Numba optimizations
try:
    from numba_optimizations import (
        NUMBA_AVAILABLE,
        extract_cell_colors_numba,
        compute_squared_distances_numba,
        assemble_mosaic_numba,
        warmup_numba_functions
    )
    print(f"✅ Numba optimizations loaded (Available: {NUMBA_AVAILABLE})")
except ImportError:
    NUMBA_AVAILABLE = False
    print("⚠️  Numba optimizations not available - using NumPy fallback")

# Deployment configuration
TILE_FOLDER = "extracted_images"

# Pre-built cache files that should be uploaded with the app
AVAILABLE_CACHES = {
    16: "cache_16x16_bins8_rot.pkl",
    32: "cache_32x32_bins8_rot.pkl", 
    64: "cache_64x64_bins8_rot.pkl"
}

# Global storage for last generation (for metrics calculation)
_last_generation = {}

@dataclass
class ImageContext:
    """Contextual analysis results."""
    has_faces: bool
    face_regions: List[Tuple[int, int, int, int]]
    is_portrait: bool
    is_landscape: bool
    dominant_colors: np.ndarray
    brightness_map: np.ndarray
    edge_density_map: np.ndarray
    content_complexity: float


class DeploymentMosaicGenerator:
    """Deployment-optimized mosaic generator with Numba support."""
    
    def __init__(self, cache_file: str, use_numba: bool = True):
        """Initialize with existing cache file only."""
        self.cache_file = cache_file
        self.use_numba = use_numba and NUMBA_AVAILABLE
        
        # Load cache data
        try:
            with open(cache_file, 'rb') as f:
                data = pickle.load(f)
            
            self.tile_images = data['tile_images']
            if not isinstance(self.tile_images, np.ndarray):
                self.tile_images = np.array(self.tile_images, dtype=np.uint8)
            
            self.tile_colours = data['tile_colours']
            self.tile_names = data['tile_names']
            self.colour_palette = data['colour_palette']
            self.colour_groups = data['colour_groups']
            
            # Convert colour_indices
            self.colour_indices = {}
            for bin_id, (index, tile_indices) in data['colour_indices'].items():
                if not isinstance(tile_indices, np.ndarray):
                    tile_indices = np.array(tile_indices, dtype=np.int32)
                self.colour_indices[bin_id] = (index, tile_indices)
            
            self.tile_size = data['tile_size']
            
            print(f"Loaded cache: {len(self.tile_images)} tiles, size {self.tile_size}")
            
            # Warmup Numba if available
            if self.use_numba:
                print("Warming up Numba JIT...")
                success = warmup_numba_functions(
                    self.tile_images,
                    self.tile_size[1],
                    self.tile_size[0]
                )
                if success:
                    print("✅ Numba JIT compiled and ready")
                else:
                    print("⚠️  Numba warmup failed, using NumPy")
                    self.use_numba = False
            
        except Exception as e:
            raise RuntimeError(f"Failed to load cache {cache_file}: {e}")
    
    def analyze_context(self, image: np.ndarray) -> ImageContext:
        """Fast context analysis for deployment - Face detection DISABLED."""
        # Face detection disabled for deployment speed
        faces = []
        has_faces = False
        
        # Scene classification (without face detection)
        aspect_ratio = image.shape[1] / image.shape[0]
        is_portrait = aspect_ratio < 1.0  # Just based on aspect ratio
        is_landscape = aspect_ratio > 1.5
        
        # Simplified dominant colors
        pixels = image.reshape(-1, 3)
        sample_size = min(5000, len(pixels))
        sample_indices = np.random.randint(0, len(pixels), size=sample_size)
        sampled_pixels = pixels[sample_indices]
        
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            kmeans = MiniBatchKMeans(
                n_clusters=3, random_state=42, batch_size=500,
                n_init=1, max_iter=30
            )
            kmeans.fit(sampled_pixels)
            dominant_colors = kmeans.cluster_centers_
        
        # Minimal analysis maps
        brightness_map = np.zeros((16, 16), dtype=np.float32)
        edge_density_map = np.zeros((16, 16), dtype=np.float32)
        content_complexity = 0.0
        
        return ImageContext(
            has_faces=False,  # Always False
            face_regions=[],   # Always empty
            is_portrait=is_portrait,
            is_landscape=is_landscape,
            dominant_colors=dominant_colors,
            brightness_map=brightness_map,
            edge_density_map=edge_density_map,
            content_complexity=0.0
        )
    
    def create_mosaic_with_preprocessing(self, image: np.ndarray, grid_size: int, 
                                        diversity_factor: float = 0.1) -> Tuple[np.ndarray, np.ndarray, ImageContext]:
        """Create mosaic with Numba optimization and preprocessing visualization."""
        print(f"Creating {grid_size}x{grid_size} mosaic (Numba: {self.use_numba})...")
        
        # Analyze context
        context = self.analyze_context(image)
        
        # Preprocess image to fit grid
        target_size = grid_size * self.tile_size[0]
        h, w = image.shape[:2]
        scale = max(target_size / w, target_size / h)
        new_w, new_h = int(w * scale), int(h * scale)
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        # Center crop
        start_x = (new_w - target_size) // 2
        start_y = (new_h - target_size) // 2
        processed_image = resized[start_y:start_y + target_size, start_x:start_x + target_size].astype(np.uint8)
        
        # Create grid visualization
        grid_visualization = self.create_grid_overlay(processed_image, grid_size)
        
        # Create mosaic with Numba optimization
        cell_size = target_size // grid_size
        tile_h, tile_w = self.tile_size[1], self.tile_size[0]
        
        # Extract cell colors (Numba-optimized if available)
        if self.use_numba:
            all_cell_colors = extract_cell_colors_numba(processed_image, grid_size, grid_size)
        else:
            # NumPy fallback
            cells = processed_image.reshape(grid_size, cell_size, grid_size, cell_size, 3)
            cells = cells.transpose(0, 2, 1, 3, 4)
            all_cell_colors = cells.mean(axis=(2, 3))
        
        flat_colors = all_cell_colors.reshape(-1, 3)
        
        # Compute distances (Numba-optimized if available)
        if self.use_numba:
            all_distances = compute_squared_distances_numba(flat_colors, self.colour_palette)
        else:
            all_distances = euclidean_distances(flat_colors, self.colour_palette)
        
        best_bins = np.argmin(all_distances, axis=1)
        best_tiles = np.zeros(len(flat_colors), dtype=np.int32)
        
        # Find best tiles for each bin
        for bin_id in np.unique(best_bins):
            bin_id = int(bin_id)
            
            if bin_id not in self.colour_indices:
                continue
            
            mask = best_bins == bin_id
            bin_colors = flat_colors[mask]
            
            if len(bin_colors) == 0:
                continue
            
            index, tile_indices = self.colour_indices[bin_id]
            n_neighbors = min(5, len(tile_indices))
            
            if n_neighbors > 0:
                _, indices = index.kneighbors(bin_colors, n_neighbors=n_neighbors)
                best_tiles[mask] = tile_indices[indices[:, 0]]
        
        best_tiles_grid = best_tiles.reshape(grid_size, grid_size)
        
        # Assemble mosaic (Numba-optimized if available and grid is large enough)
        if self.use_numba and grid_size * grid_size > 256:
            mosaic = assemble_mosaic_numba(
                self.tile_images, best_tiles_grid,
                grid_size, grid_size, tile_h, tile_w
            )
        else:
            # NumPy fancy indexing
            selected_tiles = self.tile_images[best_tiles_grid]
            mosaic = selected_tiles.transpose(0, 2, 1, 3, 4).reshape(
                grid_size * tile_h, grid_size * tile_w, 3
            )
        
        return mosaic, grid_visualization, context
    
    def create_grid_overlay(self, image: np.ndarray, grid_size: int) -> np.ndarray:
        """Create visualization showing grid segmentation with enhanced styling."""
        h, w = image.shape[:2]
        cell_h = h // grid_size
        cell_w = w // grid_size
        
        grid_image = image.copy()
        
        line_color = (255, 255, 255)
        shadow_color = (0, 0, 0)
        line_thickness = max(1, min(w, h) // 500)
        
        # Draw vertical lines
        for i in range(1, grid_size):
            x = i * cell_w
            cv2.line(grid_image, (x-1, 0), (x-1, h), shadow_color, line_thickness)
            cv2.line(grid_image, (x+1, 0), (x+1, h), shadow_color, line_thickness)
            cv2.line(grid_image, (x, 0), (x, h), line_color, line_thickness)
        
        # Draw horizontal lines
        for i in range(1, grid_size):
            y = i * cell_h
            cv2.line(grid_image, (0, y-1), (w, y-1), shadow_color, line_thickness)
            cv2.line(grid_image, (0, y+1), (w, y+1), shadow_color, line_thickness)
            cv2.line(grid_image, (0, y), (w, y), line_color, line_thickness)
        
        # Add border
        border_thickness = max(2, line_thickness * 2)
        cv2.rectangle(grid_image, (0, 0), (w-1, h-1), line_color, border_thickness)
        cv2.rectangle(grid_image, (border_thickness, border_thickness),
                     (w-border_thickness-1, h-border_thickness-1), shadow_color, 1)
        
        # Add text overlay
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = max(0.5, min(w, h) / 800)
        thickness = max(1, int(font_scale * 2))
        
        text_main = f"Grid: {grid_size}x{grid_size}"
        text_sub = f"{grid_size**2} cells total"
        text_cell = f"Cell size: {cell_w}x{cell_h}px"
        
        (main_w, main_h), _ = cv2.getTextSize(text_main, font, font_scale, thickness)
        (sub_w, sub_h), _ = cv2.getTextSize(text_sub, font, font_scale * 0.7, thickness)
        (cell_w_text, cell_h_text), _ = cv2.getTextSize(text_cell, font, font_scale * 0.6, thickness)
        
        padding = 10
        bg_width = max(main_w, sub_w, cell_w_text) + padding * 2
        bg_height = main_h + sub_h + cell_h_text + padding * 4
        
        overlay = grid_image.copy()
        cv2.rectangle(overlay, (10, 10), (10 + bg_width, 10 + bg_height), (0, 0, 0), -1)
        cv2.addWeighted(overlay, 0.7, grid_image, 0.3, 0, grid_image)
        cv2.rectangle(grid_image, (10, 10), (10 + bg_width, 10 + bg_height), line_color, 1)
        
        y_offset = 10 + padding + main_h
        cv2.putText(grid_image, text_main, (10 + padding, y_offset),
                   font, font_scale, line_color, thickness)
        
        y_offset += sub_h + padding // 2
        cv2.putText(grid_image, text_sub, (10 + padding, y_offset),
                   font, font_scale * 0.7, (200, 200, 200), thickness)
        
        y_offset += cell_h_text + padding // 2
        cv2.putText(grid_image, text_cell, (10 + padding, y_offset),
                   font, font_scale * 0.6, (180, 180, 180), thickness)
        
        # Add corner indicators for smaller grids
        if grid_size <= 16:
            for i in range(min(3, grid_size)):
                for j in range(min(3, grid_size)):
                    x = j * cell_w + 5
                    y = i * cell_h + 15
                    cell_num = i * grid_size + j + 1
                    
                    cv2.circle(grid_image, (x + 10, y), 12, (0, 0, 0), -1)
                    cv2.circle(grid_image, (x + 10, y), 12, line_color, 1)
                    cv2.putText(grid_image, str(cell_num), (x + 5, y + 4),
                               font, 0.4, line_color, 1)
        
        return grid_image


def calculate_global_ssim(original: np.ndarray, mosaic: np.ndarray) -> float:
    """Calculate Global SSIM."""
    if original.shape != mosaic.shape:
        original_resized = cv2.resize(original, (mosaic.shape[1], mosaic.shape[0]))
    else:
        original_resized = original
    
    orig_float = original_resized.astype(np.float64)
    mosaic_float = mosaic.astype(np.float64)
    
    C1 = (0.01 * 255) ** 2
    C2 = (0.03 * 255) ** 2
    
    global_ssim_values = []
    
    for channel in range(3):
        orig_channel = orig_float[:, :, channel]
        mosaic_channel = mosaic_float[:, :, channel]
        
        mu1 = np.mean(orig_channel)
        mu2 = np.mean(mosaic_channel)
        
        sigma1_sq = np.var(orig_channel)
        sigma2_sq = np.var(mosaic_channel)
        sigma12 = np.mean((orig_channel - mu1) * (mosaic_channel - mu2))
        
        numerator = (2 * mu1 * mu2 + C1) * (2 * sigma12 + C2)
        denominator = (mu1**2 + mu2**2 + C1) * (sigma1_sq + sigma2_sq + C2)
        
        channel_ssim = numerator / denominator
        global_ssim_values.append(channel_ssim)
    
    return float(np.mean(global_ssim_values))


def calculate_global_color_similarity(original: np.ndarray, mosaic: np.ndarray) -> float:
    """Calculate global color distribution similarity."""
    if original.shape != mosaic.shape:
        original_resized = cv2.resize(original, (mosaic.shape[1], mosaic.shape[0]))
    else:
        original_resized = original
    
    orig_mean = np.mean(original_resized, axis=(0, 1))
    mosaic_mean = np.mean(mosaic, axis=(0, 1))
    
    orig_std = np.std(original_resized, axis=(0, 1))
    mosaic_std = np.std(mosaic, axis=(0, 1))
    
    color_mean_diff = np.linalg.norm(orig_mean - mosaic_mean)
    color_mean_sim = 1.0 / (1.0 + color_mean_diff / 255.0)
    
    color_std_diff = np.linalg.norm(orig_std - mosaic_std)
    color_std_sim = 1.0 / (1.0 + color_std_diff / 255.0)
    
    global_color_sim = 0.6 * color_mean_sim + 0.4 * color_std_sim
    
    return float(global_color_sim)


def calculate_metrics(original: np.ndarray, mosaic: np.ndarray) -> Dict[str, float]:
    """Calculate enhanced quality metrics with global SSIM."""
    if original.shape != mosaic.shape:
        original_resized = cv2.resize(original, (mosaic.shape[1], mosaic.shape[0]))
    else:
        original_resized = original
    
    orig_float = original_resized.astype(np.float64)
    mosaic_float = mosaic.astype(np.float64)
    mse = float(np.mean((orig_float - mosaic_float) ** 2))
    
    psnr = float(20 * np.log10(255.0 / np.sqrt(mse))) if mse > 0 else float('inf')
    
    global_ssim = calculate_global_ssim(original, mosaic)
    global_color_sim = calculate_global_color_similarity(original, mosaic)
    
    correlations = []
    for channel in range(3):
        hist_orig = cv2.calcHist([original_resized], [channel], None, [256], [0, 256])
        hist_mosaic = cv2.calcHist([mosaic], [channel], None, [256], [0, 256])
        corr = cv2.compareHist(hist_orig, hist_mosaic, cv2.HISTCMP_CORREL)
        correlations.append(corr)
    histogram_similarity = float(np.mean(correlations))
    
    ssim_norm = (global_ssim + 1) / 2
    psnr_norm = min(psnr / 50.0, 1.0)
    
    overall = (
        0.4 * ssim_norm +
        0.25 * global_color_sim +
        0.2 * histogram_similarity +
        0.15 * psnr_norm
    ) * 100
    
    return {
        'mse': mse,
        'psnr': psnr,
        'ssim': global_ssim,
        'global_color_similarity': global_color_sim,
        'histogram_similarity': histogram_similarity,
        'overall_quality': float(overall)
    }


def get_best_available_cache(requested_tile_size: int) -> Optional[str]:
    """Get the best available cache for requested tile size."""
    if requested_tile_size in AVAILABLE_CACHES:
        cache_file = AVAILABLE_CACHES[requested_tile_size]
        if Path(cache_file).exists():
            return cache_file
    
    available_sizes = []
    for size, cache_file in AVAILABLE_CACHES.items():
        if Path(cache_file).exists():
            available_sizes.append(size)
    
    if not available_sizes:
        return None
    
    closest_size = min(available_sizes, key=lambda x: abs(x - requested_tile_size))
    return AVAILABLE_CACHES[closest_size]


def create_mosaic_interface(image, grid_size, tile_size, diversity_factor, 
                           enable_rotation, apply_quantization, n_colors):
    """Main interface function - Returns images immediately, metrics calculated after."""
    global _last_generation
    
    if image is None:
        return None, None, None, "⏳ Generating...", "Please upload an image first."
    
    try:
        start_time = time.time()
        
        cache_file = get_best_available_cache(tile_size)
        if cache_file is None:
            error_msg = f"No cache available for tile size {tile_size}x{tile_size}"
            return None, None, None, error_msg, error_msg
        
        # Initialize with Numba support
        generator = DeploymentMosaicGenerator(cache_file, use_numba=True)
        
        # Optional quantization
        if apply_quantization:
            pixels = image.reshape(-1, 3)
            sample_size = min(5000, len(pixels))
            sampled_pixels = pixels[np.random.choice(len(pixels), sample_size, replace=False)]
            
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                kmeans = MiniBatchKMeans(n_clusters=n_colors, batch_size=500, random_state=42)
                kmeans.fit(sampled_pixels)
                labels = kmeans.predict(pixels)
            
            quantized_pixels = kmeans.cluster_centers_[labels]
            image = quantized_pixels.reshape(image.shape).astype(np.uint8)
        
        # Create mosaic with Numba optimization
        mosaic, grid_viz, context = generator.create_mosaic_with_preprocessing(
            image, grid_size, diversity_factor
        )
        
        total_time = time.time() - start_time
        
        # Resize outputs for faster display
        MAX_DISPLAY_SIZE = 1024
        
        if max(mosaic.shape[:2]) > MAX_DISPLAY_SIZE:
            scale = MAX_DISPLAY_SIZE / max(mosaic.shape[:2])
            new_h = int(mosaic.shape[0] * scale)
            new_w = int(mosaic.shape[1] * scale)
            mosaic_display = cv2.resize(mosaic, (new_w, new_h), interpolation=cv2.INTER_AREA)
        else:
            mosaic_display = mosaic
        
        if max(grid_viz.shape[:2]) > MAX_DISPLAY_SIZE:
            scale = MAX_DISPLAY_SIZE / max(grid_viz.shape[:2])
            new_h = int(grid_viz.shape[0] * scale)
            new_w = int(grid_viz.shape[1] * scale)
            grid_viz_display = cv2.resize(grid_viz, (new_w, new_h), interpolation=cv2.INTER_AREA)
        else:
            grid_viz_display = grid_viz
        
        comparison_h = min(mosaic.shape[0], MAX_DISPLAY_SIZE)
        comparison_w_half = min(mosaic.shape[1]//2, MAX_DISPLAY_SIZE//2)
        
        comparison = np.hstack([
            cv2.resize(image, (comparison_w_half, comparison_h), interpolation=cv2.INTER_AREA),
            cv2.resize(mosaic, (comparison_w_half, comparison_h), interpolation=cv2.INTER_AREA)
        ])
        
        # Show placeholder for metrics
        metrics_text = "⏳ Calculating quality metrics..."
        
        status = f"""Generation Successful!

Grid: {grid_size}x{grid_size} = {grid_size**2} tiles
Tile Size: {tile_size}x{tile_size} pixels
Processing Time: {total_time:.2f} seconds
Cache Used: {cache_file}
Numba Optimization: {'✅ ACTIVE' if generator.use_numba else '❌ Inactive'}

Contextual Analysis:
- Faces Detected: {len(context.face_regions)}
- Scene Type: {'Portrait' if context.is_portrait else 'Landscape' if context.is_landscape else 'General'}"""
        
        # Store for metrics calculation
        _last_generation = {
            'original': image,
            'mosaic': mosaic,
            'generated': True
        }
        
        return mosaic_display, comparison, grid_viz_display, metrics_text, status
        
    except Exception as e:
        import traceback
        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
        return None, None, None, error_msg, error_msg


def calculate_and_display_metrics():
    """Calculate metrics after images are displayed."""
    global _last_generation
    
    if not _last_generation.get('generated'):
        return "No mosaic generated yet. Please generate a mosaic first."
    
    try:
        original = _last_generation['original']
        mosaic = _last_generation['mosaic']
        
        # Now calculate metrics
        metrics = calculate_metrics(original, mosaic)
        
        metrics_text = f"""ENHANCED PERFORMANCE METRICS

Mean Squared Error (MSE): {metrics['mse']:.2f}

Peak Signal-to-Noise Ratio (PSNR): {metrics['psnr']:.2f} dB

Global Structural Similarity (SSIM): {metrics['ssim']:.4f}

Global Color Similarity: {metrics['global_color_similarity']:.4f}

Color Histogram Similarity: {metrics['histogram_similarity']:.4f}

Overall Quality Score: {metrics['overall_quality']:.1f}/100"""
        
        return metrics_text
        
    except Exception as e:
        return f"Error calculating metrics: {str(e)}"


def verify_deployment_setup():
    """Check deployment setup."""
    available_caches = {}
    total_size_mb = 0
    
    for size, cache_file in AVAILABLE_CACHES.items():
        if Path(cache_file).exists():
            size_mb = Path(cache_file).stat().st_size / 1024 / 1024
            available_caches[size] = size_mb
            total_size_mb += size_mb
    
    setup_msg = f"Found {len(available_caches)} cache files ({total_size_mb:.1f}MB total)"
    
    return len(available_caches) > 0, setup_msg, available_caches


def get_system_status():
    """System status for deployment."""
    setup_ok, setup_msg, available_caches = verify_deployment_setup()
    
    cache_list = ""
    for size, size_mb in available_caches.items():
        cache_list += f"  {size}x{size}: {size_mb:.1f}MB\n"
    
    numba_status = "✅ ACTIVE" if NUMBA_AVAILABLE else "❌ Not Available (using NumPy)"
    
    status = f"""DEPLOYMENT STATUS
{'='*30}

Cache System: {'✅' if setup_ok else '❌'}
{setup_msg}

Available Caches:
{cache_list if cache_list else "  None found"}

Numba Acceleration: {numba_status}

Smart Selection: System automatically uses the best 
available cache for your chosen tile size.

INNOVATIONS INCLUDED
{'='*30}

Performance: Numba JIT compilation (3-15x speedup)
Contextual Awareness: Scene classification
Color Optimization: Subgrouping, Mini-Batch K-means
Enhanced Metrics: Global SSIM, Color similarity analysis
Grid Visualization: Shows preprocessing segmentation

DEPLOYMENT OPTIMIZED
{'='*30}

- No cache building during startup
- Fast initialization with pre-built caches
- Numba JIT for computational bottlenecks
- Lightweight processing for cloud deployment
- Maintains all core innovations
- Global quality assessment
- Metrics calculated after display for speed"""
    
    return status


def create_interface():
    """Create Gradio interface with Numba support."""
    
    css = """
    .gradio-container { max-width: 100% !important; padding: 0 20px; }
    .left-panel { 
        flex: 0 0 350px; background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
        padding: 25px; border-radius: 15px; box-shadow: 0 4px 15px rgba(0,0,0,0.1);
    }
    .right-panel { 
        flex: 1; background: white; padding: 25px; border-radius: 15px;
        box-shadow: 0 4px 15px rgba(0,0,0,0.1);
    }
    .metrics-display {
        background: linear-gradient(145deg, #667eea 0%, #764ba2 100%);
        color: white; padding: 20px; border-radius: 10px;
        font-family: 'Courier New', monospace; font-size: 14px; line-height: 1.6;
    }
    """
    
    with gr.Blocks(css=css, title="Advanced Mosaic Generator with Numba") as demo:
        
        gr.Markdown("# Advanced Contextual Mosaic Generator (Numba-Optimized)")
        gr.Markdown("AI-powered mosaic creation with Numba JIT compilation, contextual awareness, grid visualization, and performance metrics.")
        
        with gr.Accordion("System Status", open=False):
            status_display = gr.Textbox(value=get_system_status(), lines=26, show_label=False)
            gr.Button("Refresh Status").click(fn=get_system_status, outputs=status_display)
        
        with gr.Row():
            # Left Panel: Controls
            with gr.Column(scale=0, min_width=350, elem_classes=["left-panel"]):
                gr.Markdown("## Configuration")
                
                generate_btn = gr.Button("🚀 Generate Mosaic", variant="primary", size="lg")
                
                gr.Markdown("---")
                
                input_image = gr.Image(type="numpy", label="Upload Image", height=200)
                
                grid_size = gr.Slider(16, 128, 32, step=8, label="Grid Size", info="Number of tiles per side")
                tile_size = gr.Slider(16, 64, 32, step=16, label="Tile Size", info="Must match available cache")
                diversity_factor = gr.Slider(0.0, 0.5, 0.15, step=0.05, label="Diversity", info="Tile variety")
                enable_rotation = gr.Checkbox(label="Enable Rotation", value=False, info="Uses rotation variants if available")
                apply_quantization = gr.Checkbox(label="Color Quantization", value=False)
                n_colors = gr.Slider(4, 24, 12, step=2, label="Colors")

                gr.Markdown("### Or, try with an example:")
                gr.Examples(
                    examples=[
                        ["Images/EmmaPotrait.jpg", 64, 32, 0.15, False, False, 12],
                        ["Images/Batman.jpg", 128, 16, 0.05, False, False, 16],
                        ["Images/Indian_Dog.jpg", 56, 32, 0.2, False, False, 16],
                    ],
                    inputs=[input_image, grid_size, tile_size, diversity_factor, 
                           enable_rotation, apply_quantization, n_colors]
                )
                
                gr.Markdown("### Quick Presets")
                with gr.Row():
                    preset_fast = gr.Button("⚡ Fast", size="sm")
                    preset_quality = gr.Button("💎 Quality", size="sm")
            
            # Right Panel: Results
            with gr.Column(scale=2, elem_classes=["right-panel"]):
                gr.Markdown("## Results & Analysis")
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Generated Mosaic")
                        mosaic_output = gr.Image(height=300, show_label=False)
                    
                    with gr.Column():
                        gr.Markdown("### Comparison (Original | Mosaic)")
                        comparison_output = gr.Image(height=300, show_label=False)
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### Grid Segmentation Visualization")
                        grid_viz_output = gr.Image(height=300, show_label=False)
                        gr.Markdown("*Shows how the image is divided into cells for tile placement*")
                    
                    with gr.Column():
                        gr.Markdown("### Performance Metrics")
                        metrics_output = gr.Textbox(lines=8, elem_classes=["metrics-display"], show_label=False)
                
                status_output = gr.Textbox(label="Generation Status", lines=6)
        
        # Connect functions
        def fast_preset():
            return 24, 32, 0.1, False, False, 8
        
        def quality_preset():
            return 128, 16, 0.0, False, False, 24
        
        # STEP 1: Generate and show images immediately
        # STEP 2: Calculate metrics after images are shown (using .then())
        generate_btn.click(
            fn=create_mosaic_interface,
            inputs=[input_image, grid_size, tile_size, diversity_factor, 
                   enable_rotation, apply_quantization, n_colors],
            outputs=[mosaic_output, comparison_output, grid_viz_output, 
                    metrics_output, status_output]
        ).then(
            fn=calculate_and_display_metrics,
            inputs=[],
            outputs=[metrics_output]
        )
        
        preset_fast.click(fn=fast_preset, outputs=[grid_size, tile_size, diversity_factor, 
                                                   enable_rotation, apply_quantization, n_colors])
        preset_quality.click(fn=quality_preset, outputs=[grid_size, tile_size, diversity_factor, 
                                                        enable_rotation, apply_quantization, n_colors])
    
    return demo


if __name__ == "__main__":
    print("="*80)
    print("Advanced Mosaic Generator - Numba-Optimized Deployment")
    print("="*80)
    print("Checking deployment setup...")
    
    setup_ok, setup_msg, caches = verify_deployment_setup()
    print(f"Cache Status: {setup_msg}")
    print(f"Numba Status: {'✅ Available' if NUMBA_AVAILABLE else '⚠️  Not available (using NumPy)'}")
    
    if setup_ok:
        print("\n✅ Deployment ready!")
        demo = create_interface()
        demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
    else:
        print(f"\n❌ Deployment not ready: {setup_msg}")
        print("Please upload the required cache files")