"""Gradio app for the Canny pipeline walkthrough.

Live demo of the NumPy implementation. The CUDA version (5-11x faster
depending on image size) lives in the GitHub repo:
    https://github.com/AneeshB20/canny-edge-detection-cuda
"""

from __future__ import annotations

import time
from pathlib import Path

import gradio as gr
import numpy as np

from canny_stages import canny_pipeline_with_stages


GITHUB_URL  = "https://github.com/AneeshB20/canny-edge-detection-cuda"
EXAMPLES_DIR = Path(__file__).parent / "examples"

EXAMPLE_IMAGES = [
    str(EXAMPLES_DIR / "fruits.jpg"),
    str(EXAMPLES_DIR / "shapes.png"),
    str(EXAMPLES_DIR / "coins.png"),
]


# Captions for the 8 panels we display. Keep in sync with canny_stages keys.
STAGE_CAPTIONS = {
    "original":  "Original image (input).",
    "grayscale": "Step 1: Convert to grayscale. Canny operates on single-channel intensity values.",
    "gaussian":  "Step 2: Gaussian blur. Removes noise that would create false edges.",
    "magnitude": "Step 3: Sobel gradient magnitude. Bright = strong intensity change = potential edge.",
    "direction": "Step 4: Gradient direction. Color shows edge orientation - needed for thinning.",
    "nms":       "Step 5: Non-maximum suppression. Thick edges thinned to 1-pixel width by keeping only local maxima along the gradient direction.",
    "threshold": "Step 6: Double threshold. White = strong edges, gray = weak edges, black = suppressed.",
    "edges":     "Step 7: Hysteresis edge tracking. Weak edges kept only if connected to strong edges. Final result.",
}

STAGE_ORDER = ["original", "grayscale", "gaussian", "magnitude",
               "direction", "nms", "threshold", "edges"]


def run_pipeline(image, sigma, low_thresh, high_thresh):
    """Gradio callback. Returns (gallery items, info markdown)."""
    if image is None:
        return [], "**Upload an image first.**"
    if high_thresh < low_thresh:
        return [], "**`high_thresh` must be >= `low_thresh`.**"

    t0 = time.perf_counter()
    stages = canny_pipeline_with_stages(
        image, sigma=float(sigma),
        low_thresh=float(low_thresh), high_thresh=float(high_thresh),
    )
    elapsed_ms = (time.perf_counter() - t0) * 1000.0

    # Substitute sigma into the blur caption so the user sees what they picked.
    captions = dict(STAGE_CAPTIONS)
    captions["gaussian"] = (
        f"Step 2: Gaussian blur (sigma = {float(sigma):.2f}). "
        "Removes noise that would create false edges."
    )

    gallery = [(stages[k], captions[k]) for k in STAGE_ORDER]

    h, w = stages["edges"].shape[:2]
    note = (
        f"**Processing time:** {elapsed_ms:.1f} ms on {h}×{w} pixels (NumPy implementation).\n\n"
        f"The CUDA version of this same pipeline runs **5-11x faster** than NumPy "
        f"depending on image size - see the benchmark plots in the "
        f"[GitHub repo]({GITHUB_URL})."
    )
    return gallery, note


with gr.Blocks(theme=gr.themes.Soft(), title="Canny Edge Detection - Step by Step") as demo:
    gr.Markdown("# Canny Edge Detection - Step by Step")
    gr.Markdown(
        "Upload an image to see every stage of the Canny pipeline. "
        "Built from scratch - no `cv2.Canny()`. "
        f"[View the source on GitHub]({GITHUB_URL})."
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Input image", type="numpy", height=300)
            sigma       = gr.Slider(0.5, 5.0, value=1.4, step=0.1,
                                    label="Gaussian sigma (blur strength)")
            low_thresh  = gr.Slider(5, 100, value=25, step=1,
                                    label="Low threshold")
            high_thresh = gr.Slider(10, 200, value=70, step=1,
                                    label="High threshold")
            run_btn     = gr.Button("Run Canny pipeline", variant="primary")

            gr.Examples(
                examples=[[p] for p in EXAMPLE_IMAGES if Path(p).exists()],
                inputs=[input_image],
                label="Try an example",
            )

        with gr.Column(scale=2):
            output_gallery = gr.Gallery(
                label="Pipeline stages",
                columns=2,
                rows=4,
                object_fit="contain",
                height="auto",
            )
            output_info = gr.Markdown()

    run_btn.click(
        run_pipeline,
        inputs=[input_image, sigma, low_thresh, high_thresh],
        outputs=[output_gallery, output_info],
    )


if __name__ == "__main__":
    demo.launch()