import numpy as np
import cv2
import matplotlib.pyplot as plt
from numba import jit
import gradio as gr

@jit(nopython=True, parallel=True)
def poisson_sharpening_rgb(image, alpha):
    height, width, channels = image.shape
    sharpened = np.zeros_like(image, dtype=np.float32)

    for c in range(channels):
        for i in range(height):
            for j in range(width):
                # Compute indices for neighboring pixels
                left = max(0, j - 1)
                right = min(width - 1, j + 1)
                top = max(0, i - 1)
                bottom = min(height - 1, i + 1)

                # Compute differences with neighboring pixels
                diff_left = float(image[i, j, c]) - float(image[i, left, c])
                diff_right = float(image[i, j, c]) - float(image[i, right, c])
                diff_top = float(image[i, j, c]) - float(image[top, j, c])
                diff_bottom = float(image[i, j, c]) - float(image[bottom, j, c])

                # Compute sharpened pixel value
                sharpened[i, j, c] = min(max(
                    float(image[i, j, c]) + alpha * (diff_left + diff_right + diff_top + diff_bottom),
                    0.0), 255.0)

    return sharpened.astype(np.uint8)

def sharpen_image(image, alpha):
    # Ensure the image is in RGB format
    if image.shape[2] == 4:  # If RGBA, convert to RGB
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
    elif len(image.shape) == 2:  # If grayscale, convert to RGB
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    
    # Apply sharpening
    sharpened = poisson_sharpening_rgb(image, alpha)
    
    return sharpened

# Create examples list
examples = [
    ["img1.jpg", 2.0],
    ["img2.PNG", 2.0],
]

# Create the Gradio interface
iface = gr.Interface(
    fn=sharpen_image,
    inputs=[
        gr.Image(label="Input Image", type="numpy"),
        gr.Slider(minimum=1.0, maximum=15.0, step=0.01, value=2.0, label="Sharpening Strength (alpha)")
    ],
    outputs=gr.Image(label="Sharpened Image"),
    title="Poisson Image Sharpening",
    description="Upload an image or choose from the examples, then adjust the sharpening strength to enhance edges and details.",
    theme='bethecloud/storj_theme',
    examples=examples,
    cache_examples=True
)

iface.launch()