import numpy as np
import scipy as sp
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(parallel=True)
def poisson_sharpening(img, alpha):
    """
    Returns a sharpened image with strength of alpha.
    :param img: the image
    :param alpha: edge threshold and gradient scaler
    """
    img_h, img_w = img.shape
    img_s = img.copy()
    
    im2var = np.arange(img_h * img_w).reshape(img_h, img_w)
    
    A_data = np.zeros(img_h * img_w * 4 * 2, dtype=np.float64)
    A_row = np.zeros(img_h * img_w * 4 * 2, dtype=np.int32)
    A_col = np.zeros(img_h * img_w * 4 * 2, dtype=np.int32)
    b = np.zeros(img_h * img_w * 4 * 2, dtype=np.float64)
    
    e = 0
    for y in prange(img_h):
        for x in prange(img_w):
            A_data[e] = 1
            A_row[e] = e
            A_col[e] = im2var[y, x]
            b[e] = img_s[y, x]
            e += 1
            
            for n_y, n_x in neighbours(y, x, img_h-1, img_w-1):
                A_data[e] = 1
                A_row[e] = e
                A_col[e] = im2var[y, x]
                e += 1
                
                A_data[e] = -1
                A_row[e] = e - 1
                A_col[e] = im2var[n_y, n_x]
                
                b[e-1] = alpha * (img_s[y, x] - img_s[n_y, n_x])
                e += 1
    
    A = sp.sparse.csr_matrix((A_data[:e], (A_row[:e], A_col[:e])), shape=(e, img_h * img_w))
    v = sp.sparse.linalg.lsqr(A, b[:e])[0]

    return clip(v.reshape(img_h, img_w), 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()