import gradio as gr
import cv2
import numpy as np
import os
import pickle
import math
from skimage.metrics import structural_similarity as ssim

# ----------------- Constants -----------------
DATASET_FOLDER = "Dataset"  # (Not used directly now)
KD_TREE_PATH = "kdtree_dataset.pkl"  # Path to the precomputed KDTree file
KD_TILE_SIZE = (50, 50)  # Must match the tile size used when building the KDTree

# ----------------- Feature Extraction Functions -----------------
def compute_features(image):
    """
    Compute a set of features for an image:
      - Average Lab color (using a Gaussian-blurred version)
      - Edge density using Canny edge detection (normalized)
      - Texture measure using the standard deviation of the grayscale image (normalized)
      - Average gradient magnitude computed via Sobel operators (normalized)
    Returns: (avg_lab, avg_edge, avg_texture, avg_grad)
    """
    blurred = cv2.GaussianBlur(image, (5, 5), 0)
    img_lab = cv2.cvtColor(blurred, cv2.COLOR_RGB2LAB)
    avg_lab = np.mean(img_lab, axis=(0, 1))
    
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 100, 200)
    avg_edge = np.mean(edges) / 255.0
    avg_texture = np.std(gray) / 255.0
    grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
    grad_mag = np.sqrt(grad_x**2 + grad_y**2)
    avg_grad = np.mean(grad_mag) / 255.0
    
    return avg_lab, avg_edge, avg_texture, avg_grad

def compute_weighted_features(image):
    """
    Compute the weighted feature vector for KDTree search.
    The image should be resized to KD_TILE_SIZE before feature extraction.
    Weights: for Lab channels use 1.0; for edge, texture, and gradient use 0.5
    (implemented as multiplying by sqrt(0.5)).
    """
    scale = np.array([1.0, 1.0, 1.0, math.sqrt(0.5), math.sqrt(0.5), math.sqrt(0.5)])
    avg_lab, avg_edge, avg_texture, avg_grad = compute_features(image)
    raw_feature = np.concatenate([avg_lab, [avg_edge, avg_texture, avg_grad]])
    weighted_feature = raw_feature * scale
    return weighted_feature

# ----------------- Utility Function -----------------
def ensure_min_size(image, min_size=7):
    """
    Ensure that the image has at least a minimum size; if not, resize it.
    """
    h, w = image.shape[:2]
    if h < min_size or w < min_size:
        new_w = max(min_size, w)
        new_h = max(min_size, h)
        image = cv2.resize(image, (new_w, new_h))
    return image

# ----------------- Mosaic Generation Functions -----------------
def create_photo_mosaic(input_image, kdtree_path, num_tiles_y, progress=None):
    """
    Create an image mosaic using a precomputed KDTree.
    For each mosaic tile in the input image, the tile is resized to KD_TILE_SIZE,
    its weighted features are computed, and the KDTree is queried to find the best match.
    The matched dataset image is then resized to the tile’s actual size before placement.
    """
    # Load the precomputed KDTree and dataset images
    with open(kdtree_path, "rb") as f:
        tree_data = pickle.load(f)
    tree = tree_data['tree']
    dataset_images = tree_data['images']

    original_image = input_image.copy()
    height, width, _ = original_image.shape

    # Determine mosaic grid dimensions
    tile_height = height // num_tiles_y
    aspect_ratio = width / height
    num_tiles_x = int(num_tiles_y * aspect_ratio)
    tile_width = width // num_tiles_x
    print(f"Adjusted number of tiles: {num_tiles_x} (width) x {num_tiles_y} (height)")

    mosaic = np.zeros_like(original_image)
    rows = list(range(0, height, tile_height))
    cols = list(range(0, width, tile_width))
    total_tiles = len(rows) * len(cols)
    tile_count = 0

    for y in rows:
        for x in cols:
            y_end = min(y + tile_height, height)
            x_end = min(x + tile_width, width)
            tile = original_image[y:y_end, x:x_end]

            # Resize the tile to the KDTree tile size for feature extraction
            tile_resized = cv2.resize(tile, KD_TILE_SIZE)
            query_feature = compute_weighted_features(tile_resized)

            # Query the KDTree for the best match (returns index)
            dist, ind = tree.query([query_feature], k=1)
            best_index = ind[0][0]
            best_match = dataset_images[best_index]

            # Resize the best match image to the current tile size and place it into the mosaic
            best_match_resized = cv2.resize(best_match, (x_end - x, y_end - y))
            mosaic[y:y_end, x:x_end] = best_match_resized

            tile_count += 1
            if progress is not None:
                progress(tile_count / total_tiles)

    # Save the final mosaic (convert from RGB to BGR for saving with cv2)
    output_path = "mosaic_output.jpg"
    cv2.imwrite(output_path, cv2.cvtColor(mosaic, cv2.COLOR_RGB2BGR))
    return output_path

def create_color_mosaic(input_image, num_tiles_y, progress=None):
    """
    Create a simple color mosaic by dividing the image into grid cells and filling
    each cell with its average RGB color.
    """
    original_image = input_image.copy()
    height, width, _ = original_image.shape
    tile_height = height // num_tiles_y
    aspect_ratio = width / height
    num_tiles_x = int(num_tiles_y * aspect_ratio)
    tile_width = width // num_tiles_x
    print(f"Adjusted number of tiles: {num_tiles_x} (width) x {num_tiles_y} (height)")

    mosaic = np.zeros_like(original_image)
    rows = list(range(0, height, tile_height))
    cols = list(range(0, width, tile_width))
    total_tiles = len(rows) * len(cols)
    tile_count = 0

    for y in rows:
        for x in cols:
            y_end = min(y + tile_height, height)
            x_end = min(x + tile_width, width)
            tile = original_image[y:y_end, x:x_end]
            avg_color = np.mean(tile, axis=(0, 1)).astype(np.uint8)
            mosaic[y:y_end, x:x_end] = avg_color
            tile_count += 1
            if progress is not None:
                progress(tile_count / total_tiles)

    output_path = "color_mosaic_output.jpg"
    cv2.imwrite(output_path, cv2.cvtColor(mosaic, cv2.COLOR_RGB2BGR))
    return output_path

# ----------------- Performance Metrics -----------------
def compute_mse(original, mosaic):
    """
    Compute Mean Squared Error (MSE) between two images.
    """
    original = original.astype("float")
    mosaic = mosaic.astype("float")
    err = np.sum((original - mosaic) ** 2)
    mse = err / float(original.shape[0] * original.shape[1] * original.shape[2])
    return mse

def compute_ssim_metric(original, mosaic):
    """
    Compute Structural Similarity Index (SSIM) between two images.
    """
    min_dim = min(original.shape[0], original.shape[1])
    win_size = 7 if min_dim >= 7 else (min_dim if min_dim % 2 == 1 else min_dim - 1)
    ssim_value, _ = ssim(original, mosaic, win_size=win_size, channel_axis=-1, full=True)
    return ssim_value

# ----------------- Gradio Interface -----------------
def mosaic_gradio(input_image, num_tiles_y, mosaic_type, progress=gr.Progress()):
    """
    Gradio interface function to generate and return the mosaic image along with performance metrics.
    mosaic_type: "Color Mosaic" or "Image Mosaic"
    Returns: (mosaic_image_file, performance_metrics_string)
    """
    if mosaic_type == "Color Mosaic":
        mosaic_path = create_color_mosaic(input_image, num_tiles_y, progress)
    else:
        mosaic_path = create_photo_mosaic(input_image, KD_TREE_PATH, num_tiles_y, progress)

    mosaic_image = cv2.imread(mosaic_path)
    if mosaic_image is None:
        return None, "Error: Mosaic image could not be loaded."
    mosaic_image = cv2.cvtColor(mosaic_image, cv2.COLOR_BGR2RGB)

    input_for_metrics = ensure_min_size(input_image.copy())
    mosaic_for_metrics = ensure_min_size(mosaic_image.copy())

    mse_value = compute_mse(input_for_metrics, mosaic_for_metrics)
    ssim_value = compute_ssim_metric(input_for_metrics, mosaic_for_metrics)

    metrics_text = f"MSE: {mse_value:.2f}\nSSIM: {ssim_value:.4f}"
    return mosaic_path, metrics_text

# ----------------- Gradio App Setup -----------------
examples = [
    ["input_images/1.jpg", 90, "Image Mosaic"],
    ["input_images/2.jpg", 90, "Image Mosaic"],
    ["input_images/3.jpg", 90, "Image Mosaic"],
    ["input_images/6.jpg", 90, "Image Mosaic"],
    ["input_images/7.jpg", 90, "Image Mosaic"],
    ["input_images/8.jpg", 90, "Image Mosaic"],
    ["input_images/9.jpg", 90, "Image Mosaic"],
    ["input_images/10.jpg", 90, "Image Mosaic"]
]

iface = gr.Interface(
    fn=mosaic_gradio,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Slider(10, 200, value=90, step=5, label="Number of Tiles (Height)"),
        gr.Radio(choices=["Color Mosaic", "Image Mosaic"], label="Mosaic Type", value="Image Mosaic")
    ],
    outputs=[
        gr.Image(type="filepath", label="Generated Mosaic"),
        gr.Textbox(label="Performance Metrics")
    ],
    title="Photo Mosaic Generator",
    description=("Upload an image, choose the number of tiles (height) and mosaic type. "
                 "Select 'Color Mosaic' for a mosaic using average colors, or 'Image Mosaic' to use dataset images "
                 "matched by color, edge density, texture, and gradient features. "
                 "After mosaic generation, performance metrics (MSE and SSIM) will be displayed."),
    examples=examples
)

iface.launch()