import numpy as np
from numba import njit, prange, vectorize
import gradio as gr

@vectorize(['float64(float64, float64)'], nopython=True)
def clip(a, max_value):
    return min(max(a, 0), max_value)

@njit
def neighbours(y, x, max_y, max_x):
    n = []
    if y > 0:
        n.append((y-1, x))
    if y < max_y:
        n.append((y+1, x))
    if x > 0:
        n.append((y, x-1))
    if x < max_x:
        n.append((y, x+1))
    return n

@njit
def poisson_sharpening(img, alpha, num_iterations=50):
    img_h, img_w = img.shape
    v = img.copy()
    
    for _ in range(num_iterations):
        for y in range(img_h):
            for x in range(img_w):
                neighbors = neighbours(y, x, img_h-1, img_w-1)
                num_neighbors = len(neighbors)
                neighbor_sum = sum(v[ny, nx] for ny, nx in neighbors)
                laplacian = neighbor_sum - num_neighbors * v[y, x]
                v[y, x] += (laplacian + alpha * (img[y, x] - v[y, x])) / (num_neighbors + alpha)
    
    return clip(v, 1.0)

@njit(parallel=True)
def sharpen_image(img, alpha):
    # Convert alpha to float
    alpha = float(alpha)
    
    # Ensure the image is in the correct format
    img = img.astype(np.float64) / 255.0
    
    sharpen_img = np.zeros_like(img)
    for b in prange(3):
        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
    
    # Convert back to uint8 for Gradio output
    return (sharpen_img * 255).astype(np.uint8)

# Create examples list
examples = [
    ["img1.jpg", 9.0],
    ["img2.PNG", 7.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=9.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()