import cv2
import numpy as np
import scipy as sp
import scipy.sparse.linalg
import gradio as gr
import os

def get_image(img, mask=False):
    if mask:
        if isinstance(img, str):
            img = cv2.imread(img, cv2.IMREAD_GRAYSCALE)
        elif img.ndim == 3:
            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        return np.where(img > 127, 1, 0)
    else:
        if isinstance(img, str):
            img = cv2.imread(img)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        elif img.ndim == 2:
            img = np.stack((img,)*3, axis=-1)
        return img.astype('double') / 255.0

def neighbours(i, j, max_i, max_j):
    pairs = []
    for n in [-1, 1]:
        if 0 <= i+n <= max_i:
            pairs.append((i+n, j))
        if 0 <= j+n <= max_j:
            pairs.append((i, j+n))
    return pairs

def poisson_blend(img_s, mask, img_t):
    img_s_h, img_s_w = img_s.shape
    nnz = np.sum(mask > 0)
    im2var = np.full(mask.shape, -1, dtype='int32')
    im2var[mask > 0] = np.arange(nnz)
    
    ys, xs = np.where(mask == 1)
    
    # Precompute neighbor indices
    y_n = np.array([ys-1, ys+1, ys, ys])
    x_n = np.array([xs, xs, xs-1, xs+1])
    
    # Clip indices to image boundaries
    y_n = np.clip(y_n, 0, img_s_h-1)
    x_n = np.clip(x_n, 0, img_s_w-1)
    
    # Compute differences
    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
    
    # Construct sparse matrix A and vector b
    rows = np.repeat(np.arange(4*nnz), 2)
    cols = np.column_stack([np.repeat(im2var[ys, xs], 4), im2var[y_n, x_n].ravel()])
    data = np.tile([1, -1], 4*nnz)
    
    mask_n = (im2var[y_n, x_n] != -1).ravel()
    rows = rows[mask_n]
    cols = cols[mask_n]
    data = data[mask_n]
    
    A = sp.sparse.csr_matrix((data, (rows, cols)), shape=(4*nnz, nnz))
    b = d.ravel()
    b[~mask_n] += img_t[y_n.ravel()[~mask_n], x_n.ravel()[~mask_n]]
    
    # Solve the system
    v = sp.sparse.linalg.lsqr(A, b)[0]
    
    # Update the target image
    img_t_out = img_t.copy()
    img_t_out[ys, xs] = v
    
    return np.clip(img_t_out, 0, 1)

def mixed_blend(img_s, mask, img_t):
    img_s_h, img_s_w = img_s.shape
    nnz = np.sum(mask > 0)
    im2var = np.full(mask.shape, -1, dtype='int32')
    im2var[mask > 0] = np.arange(nnz)
    
    ys, xs = np.where(mask == 1)
    
    # Precompute neighbor indices
    y_n = np.array([ys-1, ys+1, ys, ys])
    x_n = np.array([xs, xs, xs-1, xs+1])
    
    # Clip indices to image boundaries
    y_n = np.clip(y_n, 0, img_s_h-1)
    x_n = np.clip(x_n, 0, img_s_w-1)
    
    # Compute differences
    ds = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
    dt = img_t[ys, xs][:, np.newaxis] - img_t[y_n, x_n]
    
    # Choose larger gradient
    d = np.where(np.abs(ds) > np.abs(dt), ds, dt)
    
    # Construct sparse matrix A and vector b
    rows = np.repeat(np.arange(4*nnz), 2)
    cols = np.column_stack([np.repeat(im2var[ys, xs], 4), im2var[y_n, x_n].ravel()])
    data = np.tile([1, -1], 4*nnz)
    
    mask_n = (im2var[y_n, x_n] != -1).ravel()
    rows = rows[mask_n]
    cols = cols[mask_n]
    data = data[mask_n]
    
    A = sp.sparse.csr_matrix((data, (rows, cols)), shape=(4*nnz, nnz))
    b = d.ravel()
    b[~mask_n] += img_t[y_n.ravel()[~mask_n], x_n.ravel()[~mask_n]]
    
    # Solve the system
    v = sp.sparse.linalg.lsqr(A, b)[0]
    
    # Update the target image
    img_t_out = img_t.copy()
    img_t_out[ys, xs] = v
    
    return np.clip(img_t_out, 0, 1)

def laplacian_blend(img1, img2, mask, depth=5, sigma=25):
    def _2d_gaussian(sigma):
        ksize = int(np.ceil(sigma) * 6 + 1)
        gaussian_1d = cv2.getGaussianKernel(ksize, sigma)
        return gaussian_1d @ gaussian_1d.T

    def _low_pass_filter(img, sigma):
        return cv2.filter2D(img, -1, _2d_gaussian(sigma))

    def _high_pass_filter(img, sigma):
        return img - _low_pass_filter(img, sigma)

    def _gaus_pyramid(img, depth, sigma):
        pyramid = [img]
        for _ in range(depth - 1):
            img = _low_pass_filter(cv2.pyrDown(img), sigma)
            pyramid.append(img)
        return pyramid

    def _lap_pyramid(img, depth, sigma):
        pyramid = []
        for d in range(depth - 1):
            next_img = cv2.pyrDown(img)
            lap = img - cv2.pyrUp(next_img, dstsize=img.shape[:2])
            pyramid.append(lap)
            img = next_img
        pyramid.append(img)
        return pyramid

    def _blend(img1, img2, mask):
        return img1 * mask + img2 * (1.0 - mask)

    # Ensure mask is 3D
    if mask.ndim == 2:
        mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)

    # Create Gaussian pyramid for mask
    mask_gaus_pyramid = _gaus_pyramid(mask, depth, sigma)

    # Create Laplacian pyramids for images
    img1_lap_pyramid = _lap_pyramid(img1, depth, sigma)
    img2_lap_pyramid = _lap_pyramid(img2, depth, sigma)

    # Blend pyramids
    blended_pyramid = [_blend(img1_lap, img2_lap, mask_gaus) 
                       for img1_lap, img2_lap, mask_gaus 
                       in zip(img1_lap_pyramid, img2_lap_pyramid, mask_gaus_pyramid)]

    # Reconstruct image
    blended_img = blended_pyramid[-1]
    for lap in reversed(blended_pyramid[:-1]):
        blended_img = cv2.pyrUp(blended_img, dstsize=lap.shape[:2])
        blended_img += lap

    return np.clip(blended_img, 0, 1)

def load_example_images(bg_path, obj_path, mask_path):
    bg_img = cv2.imread(bg_path)
    bg_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2RGB)
    
    obj_img = cv2.imread(obj_path)
    obj_img = cv2.cvtColor(obj_img, cv2.COLOR_BGR2RGB)
    
    mask_img = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    mask_img = np.where(mask_img > 127, 255, 0).astype(np.uint8)
    
    return bg_img, obj_img, mask_img

# Modify the blend_images function to accept numpy arrays directly
def blend_images(bg_img, obj_img, mask_img, blend_method):
    bg_img = get_image(bg_img)
    obj_img = get_image(obj_img)
    mask_img = get_image(mask_img, mask=True)

    # Ensure mask is 2D
    if mask_img.ndim == 3:
        mask_img = mask_img[:,:,0]  # Take the first channel if it's 3D

    # Resize mask to match object image size
    mask_img = cv2.resize(mask_img, (obj_img.shape[1], obj_img.shape[0]))

    if blend_method == "Poisson":
        blend_func = poisson_blend
    elif blend_method == "Mixed Gradient":
        blend_func = mixed_blend
    else:  # Laplacian
        return laplacian_blend(obj_img, bg_img, np.stack((mask_img,)*3, axis=-1), 5, 25)

    blend_img = np.zeros(bg_img.shape)
    for b in range(3):
        blend_img[:,:,b] = blend_func(obj_img[:,:,b], mask_img, bg_img[:,:,b].copy())
    
    return (blend_img * 255).astype(np.uint8)

examples = [
    ["img1.jpg", "img2.jpg", "mask1.jpg", "Poisson"],
    ["img3.jpg", "img4.jpg", "mask2.jpg", "Mixed Gradient"],
    ["img6.jpg", "img9.jpg", "mask3.jpg", "Laplacian"]
]

iface = gr.Interface(
    fn=blend_images,
    inputs=[
        gr.Image(label="Background Image", type="numpy"),
        gr.Image(label="Object Image", type="numpy"),
        gr.Image(label="Mask Image", type="numpy"),
        gr.Radio(["Poisson", "Mixed Gradient", "Laplacian"], label="Blending Method")
    ],
    outputs=gr.Image(label="Blended Image"),
    title="Image Blending with Examples",
    description="Choose from example images or upload your own to blend using different methods.",
    examples=examples,
    cache_examples=True
)

iface.launch()