#!/usr/bin/env python3
"""
app.py — Optimised Mosaic Generator (Lab 5)

This powers the Gradio Space using:
- 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 (Load CIFAR tiles ONCE per Space launch)
# -------------------------------------------------------------
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
):
    """Full mosaic generation pipeline with error handling."""

    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 image
    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}"

    # LAB conversion + mean colors
    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 conversion failed: {e}"

    w, h, cell_w, cell_h = dims

    # Prepare cached scaled tiles
    TM.prepare_scaled_tiles(cell_w, cell_h)

    # Find best matching tiles (FAISS)
    try:
        idxs = TM.lookup_tiles(cell_means)
    except Exception as e:
        return None, None, None, f"Tile lookup failed: {e}"

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

    # Metrics
    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

    # Build 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 + OpenCV + LAB mosaic generator.\n")

        with gr.Row():

            # ----------------------------------------------------
            # LEFT COLUMN (INPUTS)
            # ----------------------------------------------------
            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 (cells per side)"
                )
                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 Overlay")

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

                # ----------------------------------------------------
                # EXAMPLE IMAGES (LOADED FROM REPO ROOT)
                # ----------------------------------------------------
                gr.Markdown("### Example Images")

                example_files = [
                    "725px-Mona_Lisa_by_Leonardo_da_Vinci_from_C2RMF_retouched-e1660680153902.webp",
                    "WhatsApp Image 2025-11-08 at 01.39.58_cddcf540.jpg",
                ]

                example_list = [[f] for f in example_files]

                gr.Examples(
                    examples=example_list,
                    inputs=[img_in],
                    label="",
                    cache_examples=False,  # required for HF Spaces
                )

            # ----------------------------------------------------
            # RIGHT COLUMN (OUTPUTS)
            # ----------------------------------------------------
            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 Output"):
                    img_mosaic = gr.Image()

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

        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 APP
# -------------------------------------------------------------
if __name__ == "__main__":
    demo = build_demo()
    demo.launch()