import cv2
import numpy as np
import scipy as sp
import scipy.sparse.linalg
from numba import njit, prange
import gradio as gr

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

@njit
def build_poisson_matrix(img, alpha, img_h, img_w):
    im2var = np.arange(img_h * img_w).reshape(img_h, img_w)
    A_data = np.zeros(img_h * img_w * 5, dtype=np.float64)
    A_row = np.zeros(img_h * img_w * 5, dtype=np.int32)
    A_col = np.zeros(img_h * img_w * 5, dtype=np.int32)
    b = np.zeros(img_h * img_w * 5, dtype=np.float64)
    
    e = 0
    for y in range(img_h):
        for x in range(img_w):
            A_data[e] = 1
            A_row[e] = e
            A_col[e] = im2var[y, x]
            b[e] = img[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[y, x] - img[n_y, n_x])
                e += 1
    
    return A_data[:e], A_row[:e], A_col[:e], b[:e], e

def poisson_sharpening(img: np.ndarray, alpha: float) -> np.ndarray:
    img_h, img_w = img.shape
    A_data, A_row, A_col, b, e = build_poisson_matrix(img, alpha, img_h, img_w)
    
    A = sp.sparse.csr_matrix((A_data, (A_row, A_col)), shape=(e, img_h * img_w))
    v = sp.sparse.linalg.lsqr(A, b)[0]

    return np.clip(v.reshape(img_h, img_w), 0, 1)

@njit(parallel=True)
def sharpen_image_channels(img, alpha):
    sharpen_img = np.zeros_like(img)
    for b in prange(3):
        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
    return sharpen_img

def get_image(img):
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype('float32') / 255.0

def sharpen_image(input_img, alpha):
    img = get_image(input_img)
    sharpen_img = sharpen_image_channels(img, alpha)
    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()