import gradio as gr
import numpy as np
import cv2
from skimage.metrics import structural_similarity as ssim
from PIL import Image

def preprocess_image(image, grid_size):
    """Resize image to ensure consistent processing."""
    if isinstance(image, str):  # If image is a file path
        image = cv2.imread(image)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    height, width = image.shape[:2]
    max_dim = 512
    scale = max_dim / max(height, width)
    new_width = int(width * scale)
    new_height = int(height * scale)
    return cv2.resize(image, (new_width, new_height))

def quantize_colors(image, n_colors=8):
    """Quantize colors using optimized K-means clustering."""
    pixels = image.reshape(-1, 3).astype(np.float32)  # Flatten image for K-means
    
    # Define K-means criteria and apply clustering
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
    _, labels, centers = cv2.kmeans(pixels, n_colors, None, criteria, 10, cv2.KMEANS_PP_CENTERS)
    
    # Convert back to uint8 and reshape
    centers = np.uint8(centers)
    quantized = centers[labels.flatten()]
    return quantized.reshape(image.shape)

def create_grid(image, grid_size):
    """Divide image into grid cells."""
    height, width = image.shape[:2]
    cell_height = height // grid_size
    cell_width = width // grid_size
    
    grid = []
    for i in range(grid_size):
        row = []
        for j in range(grid_size):
            y_start = i * cell_height
            x_start = j * cell_width
            y_end = height if i == grid_size - 1 else (i + 1) * cell_height
            x_end = width if j == grid_size - 1 else (j + 1) * cell_width
            cell = image[y_start:y_end, x_start:x_end]
            row.append(cell)
        grid.append(row)
    return grid

def get_average_color(cell):
    """Calculate average color of a grid cell."""
    return np.mean(cell, axis=(0, 1)).astype(np.uint8)

def create_mosaic(image, grid_size, use_color_quantization, n_colors=8):
    """Create mosaic from input image using optional color quantization."""
    processed_image = preprocess_image(image, grid_size)
    
    if use_color_quantization:
        processed_image = quantize_colors(processed_image, n_colors)
    
    grid = create_grid(processed_image, grid_size)
    
    height, width = processed_image.shape[:2]
    cell_height = height // grid_size
    cell_width = width // grid_size
    
    mosaic = np.zeros_like(processed_image)
    
    for i in range(grid_size):
        for j in range(grid_size):
            y_start = i * cell_height
            x_start = j * cell_width
            y_end = height if i == grid_size - 1 else (i + 1) * cell_height
            x_end = width if j == grid_size - 1 else (j + 1) * cell_width

            avg_color = get_average_color(grid[i][j])
            mosaic[y_start:y_end, x_start:x_end] = avg_color

    
    return processed_image, mosaic

def calculate_similarity(original, mosaic):
    """Calculate similarity between original and mosaic images."""
    original_gray = cv2.cvtColor(original, cv2.COLOR_RGB2GRAY)
    mosaic_gray = cv2.cvtColor(mosaic, cv2.COLOR_RGB2GRAY)
    
    similarity = ssim(original_gray, mosaic_gray)
    mse = np.mean((original - mosaic) ** 2)
    
    return similarity, mse

def process_image(input_image, grid_size, use_quantization, n_colors):
    """Process image with K-means color quantization instead of a fixed palette."""
    original, mosaic = create_mosaic(input_image, grid_size, use_quantization, n_colors)
    similarity, mse = calculate_similarity(original, mosaic)
    
    return (mosaic, 
            f"Structural Similarity: {similarity:.4f}",
            f"Mean Squared Error: {mse:.4f}")

iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Slider(minimum=8, maximum=64, step=8, value=16, label="Grid Size"),
        gr.Checkbox(label="Use Color Quantization"),
        gr.Slider(minimum=2, maximum=16, step=1, value=8, label="Number of Colors (K-Means)")
    ],
    outputs=[
        gr.Image(type="numpy", label="Mosaic Result"),
        gr.Textbox(label="Structural Similarity (SSIM)"),
        gr.Textbox(label="Mean Squared Error (MSE)")
    ],
    title="Interactive Image Mosaic Generator (Optimized K-Means)",
    description="Upload an image to create a mosaic-style reconstruction. Adjust the grid size and color quantization settings.",
    examples=[["https://res.cloudinary.com/daigovpbf/image/upload/c_crop,g_auto,h_800,w_800/samples/cup-on-a-table", 8, False, 0],
              ["https://res.cloudinary.com/daigovpbf/image/upload/c_crop,g_auto,h_800,w_800/samples/dessert-on-a-plate", 16, True, 8],
              ["https://res.cloudinary.com/daigovpbf/image/upload/c_crop,g_auto,h_800,w_800/samples/breakfast", 32, True, 4],
              ["https://res.cloudinary.com/daigovpbf/image/upload/c_crop,g_auto,h_800,w_800/samples/balloons", 64, True, 8]],
    theme="default"
)

if __name__ == "__main__":
    iface.launch(share=True)