#!/usr/bin/env python3
"""
app.py

Gradio interface for the Optimised Mosaic Generator (Lab 5).
Loads CIFAR tiles once, then performs fast LAB-based matching using FAISS.

Connects UI to:
- crop_to_multiple()
- compute_cell_means_lab()
- TileManager
- MosaicBuilder
- MSE / SSIM metrics
"""

import gradio as gr
import numpy as np
import time
import os
from PIL import Image, ImageDraw

from mosaic_generator.image_processor import crop_to_multiple, compute_cell_means_lab
from mosaic_generator.tile_manager import TileManager
from mosaic_generator.mosaic_builder import MosaicBuilder
from mosaic_generator.metrics import mse, ssim_rgb


# ---------------------------------------------------------------
# GLOBAL TILE MANAGER — Loaded ONCE for entire HuggingFace Space
# ---------------------------------------------------------------
TM = TileManager()
TM.load(sample_size=20000)


# ---------------------------------------------------------------
# MAIN PIPELINE
# ---------------------------------------------------------------
def run_pipeline(
    img, grid_size, tile_px, tile_sample,
    quantize_on, quantize_colors, show_grid
):
    """
    Runs full mosaic generation pipeline.

    Returns:
        original image        (PIL)
        segmented grid image  (PIL)
        mosaic image          (PIL)
        report string
    """

    if img is None:
        return None, None, None, "Upload an image first."

    img_np = np.array(img.convert("RGB"))
    grid_n = int(grid_size)

    # --- Crop to perfect grid ---
    base = crop_to_multiple(img_np, grid_n)

    # --- Optional quantization ---
    if quantize_on:
        try:
            q = Image.fromarray(base).quantize(
                colors=int(quantize_colors),
                method=Image.MEDIANCUT,
                dither=Image.Dither.NONE
            ).convert("RGB")
            base = np.array(q)
        except Exception as e:
            return None, None, None, f"Quantization failed: {e}"

    # --- Compute LAB means ---
    try:
        t0 = time.perf_counter()
        cell_means, dims = compute_cell_means_lab(base, grid_n)
        t1 = time.perf_counter()
    except Exception as e:
        return None, None, None, f"LAB computation failed: {e}"

    w, h, cell_w, cell_h = dims

    # --- Prepare FAISS tiles ---
    TM.prepare_scaled_tiles(cell_w, cell_h)

    # --- Lookup nearest tiles ---
    try:
        idxs = TM.lookup_tiles(cell_means)
    except Exception as e:
        return None, None, None, f"Tile lookup failed: {e}"

    # --- Build final mosaic ---
    builder = MosaicBuilder(TM)
    try:
        mosaic_np = builder.build(idxs, dims, grid_n)
    except Exception as e:
        return None, None, None, f"Mosaic build failed: {e}"

    t2 = time.perf_counter()

    # --- Compute metrics (safe fallback) ---
    try:
        mse_val = mse(base, mosaic_np)
        ssim_val = ssim_rgb(base, mosaic_np)
    except:
        mse_val, ssim_val = -1, -1

    # --- Optional grid overlay ---
    segmented = Image.fromarray(base)
    if show_grid:
        seg = segmented.copy()
        draw = ImageDraw.Draw(seg)
        for x in range(0, w, cell_w):
            draw.line([(x, 0), (x, h)], fill="red", width=1)
        for y in range(0, h, cell_h):
            draw.line([(0, y), (w, y)], fill="red", width=1)
        segmented = seg

    # --- Text report ---
    report = (
        f"MSE: {mse_val:.2f}\n"
        f"SSIM: {ssim_val:.4f}\n\n"
        f"Preprocessing Time: {t1 - t0:.3f}s\n"
        f"Mosaic Build Time:  {t2 - t1:.3f}s\n"
        f"Total Time:          {t2 - t0:.3f}s\n"
    )

    return (
        Image.fromarray(base),
        segmented,
        Image.fromarray(mosaic_np),
        report
    )


# ---------------------------------------------------------------
# GRADIO UI
# ---------------------------------------------------------------
def build_demo():
    with gr.Blocks(title="High-Performance Mosaic Generator") as demo:

        gr.Markdown("# ⚡ High-Performance Mosaic Generator (Lab 5)")
        gr.Markdown("Ultra-fast FAISS-powered image mosaic generator.\n")

        with gr.Row():

            # ---------------- LEFT PANEL ----------------
            with gr.Column(scale=1):

                img_in = gr.Image(type="pil", label="Upload Image")

                grid_size = gr.Radio(
                    ["16", "32", "64", "128"],
                    value="32",
                    label="Grid Size"
                )

                tile_px = gr.Radio(
                    ["8", "16", "24", "32"],
                    value="16",
                    label="Tile Resolution (px)"
                )

                tile_sample = gr.Slider(
                    512, 20000, step=256, value=2048,
                    label="Tile Sample Size"
                )

                quantize_on = gr.Checkbox(True, label="Enable Color Quantization")

                quantize_colors = gr.Slider(
                    8, 128, value=32, step=8,
                    label="Quantization Palette Size"
                )

                show_grid = gr.Checkbox(True, label="Show Grid")

                run_btn = gr.Button("Generate Mosaic", variant="primary")

                # ---------------------------------------------------
                # EXAMPLES REMOVED (HF blocks external URLs)
                # ---------------------------------------------------
                gr.Markdown(
                    "### Example Images\n"
                    "⚠️ Disabled on Hugging Face due to security restrictions.\n"
                    "Please upload your own image."
                )

            # ---------------- RIGHT PANEL ----------------
            with gr.Column(scale=2):

                with gr.Tab("Original"):
                    img_orig = gr.Image()

                with gr.Tab("Grid View"):
                    img_seg = gr.Image()

                with gr.Tab("Mosaic"):
                    img_mosaic = gr.Image()

                report = gr.Textbox(label="Timing & Metrics", lines=10)

        # Connect button
        run_btn.click(
            fn=run_pipeline,
            inputs=[img_in, grid_size, tile_px, tile_sample,
                    quantize_on, quantize_colors, show_grid],
            outputs=[img_orig, img_seg, img_mosaic, report]
        )

    return demo


# ---------------------------------------------------------------
# LAUNCH
# ---------------------------------------------------------------
if __name__ == "__main__":
    demo = build_demo()
    demo.launch()