import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt


# =========================================================
# KERNELS
# =========================================================
kernels = {

    "Blur": np.ones((3, 3), np.float32) / 9,

    "Sharpen": np.array([
        [0, -1, 0],
        [-1, 5, -1],
        [0, -1, 0]
    ]),

    "Edge Detection": np.array([
        [-1, -1, -1],
        [-1,  8, -1],
        [-1, -1, -1]
    ]),

    "Emboss": np.array([
        [-2, -1, 0],
        [-1,  1, 1],
        [0,   1, 2]
    ])
}


# =========================================================
# IMAGE PROCESSING
# =========================================================
def process_image(image, kernel_name):

    if image is None:
        return None

    # RGB → Gray
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Selected Kernel
    kernel = kernels[kernel_name]

    # Apply Kernel
    filtered = cv2.filter2D(gray, -1, kernel)

    # Edge Detection
    canny = cv2.Canny(gray, 100, 200)

    # Histogram
    hist = cv2.calcHist([gray], [0], None, [256], [0, 256])

    # =========================================================
    # VISUALIZATION
    # =========================================================
    fig, axs = plt.subplots(2, 3, figsize=(14, 8))

    fig.suptitle(
        "Kernel Matrix Visualization",
        fontsize=18,
        fontweight="bold"
    )

    # ---------------------------------------------------------
    axs[0, 0].imshow(image)
    axs[0, 0].set_title("Original Image")
    axs[0, 0].axis("off")

    # ---------------------------------------------------------
    axs[0, 1].imshow(gray, cmap="gray")
    axs[0, 1].set_title("Grayscale")
    axs[0, 1].axis("off")

    # ---------------------------------------------------------
    axs[0, 2].imshow(filtered, cmap="gray")
    axs[0, 2].set_title(f"{kernel_name} Output")
    axs[0, 2].axis("off")

    # ---------------------------------------------------------
    axs[1, 0].imshow(canny, cmap="gray")
    axs[1, 0].set_title("Canny Edge Detection")
    axs[1, 0].axis("off")

    # ---------------------------------------------------------
    axs[1, 1].plot(hist)
    axs[1, 1].set_title("Pixel Histogram")
    axs[1, 1].set_xlim([0, 256])

    # ---------------------------------------------------------
    axs[1, 2].imshow(kernel, cmap="coolwarm")

    axs[1, 2].set_title(f"{kernel_name} Matrix")

    for i in range(kernel.shape[0]):
        for j in range(kernel.shape[1]):

            axs[1, 2].text(
                j,
                i,
                str(kernel[i, j]),
                ha="center",
                va="center",
                fontsize=12,
                fontweight="bold",
                color="black"
            )

    axs[1, 2].set_xticks([])
    axs[1, 2].set_yticks([])

    plt.tight_layout()

    return fig


# =========================================================
# UI
# =========================================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown(
        """
        # 🧠 Kernel Matrix Visualization
        
        Upload image and visualize:
        - Grayscale Conversion
        - Convolution Kernels
        - Edge Detection
        - Histogram
        - Kernel Matrix
        """
    )

    with gr.Row():

        input_image = gr.Image(
            type="numpy",
            label="Upload Image"
        )

        kernel_dropdown = gr.Dropdown(
            choices=list(kernels.keys()),
            value="Edge Detection",
            label="Select Kernel"
        )

    output_plot = gr.Plot()

    # AUTO RUN
    input_image.change(
        fn=process_image,
        inputs=[input_image, kernel_dropdown],
        outputs=output_plot
    )

    kernel_dropdown.change(
        fn=process_image,
        inputs=[input_image, kernel_dropdown],
        outputs=output_plot
    )

demo.launch()