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:
        return np.where(img > 127, 1, 0)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB).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 = (mask>0).sum()
    im2var = -np.ones(mask.shape[0:2], dtype='int32')
    im2var[mask>0] = np.arange(nnz)
    
    ys, xs = np.where(mask==1) 
        
    A = sp.sparse.lil_matrix((4*nnz, nnz))
    b = np.zeros(4*nnz)
    
    e = 0
    for n in range(nnz):
        y, x = ys[n], xs[n]  
        
        for n_y, n_x in neighbours(y, x, img_s_h-1, img_s_w-1):
            A[e, im2var[y][x]] = 1
            b[e] = img_s[y][x] - img_s[n_y][n_x]
            
            if im2var[n_y][n_x] != -1:
                A[e, im2var[n_y][n_x]] = -1
            else:
                b[e] += img_t[n_y][n_x]
            e += 1
    
    A = sp.sparse.csr_matrix(A)
    v = sp.sparse.linalg.lsqr(A, b)[0]
    
    img_t_out = img_t.copy()
    
    for n in range(nnz):
        y, x = ys[n], xs[n]
        img_t_out[y][x] = v[im2var[y][x]]
    
    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 = (mask>0).sum()
    im2var = -np.ones(mask.shape[0:2], dtype='int32')
    im2var[mask>0] = np.arange(nnz)
    
    ys, xs = np.where(mask==1) 
        
    A = sp.sparse.lil_matrix((4*nnz, nnz))
    b = np.zeros(4*nnz)
    
    e = 0
    for n in range(nnz):
        y, x = ys[n], xs[n]  
        
        for n_y, n_x in neighbours(y, x, img_s_h-1, img_s_w-1):
            ds = img_s[y][x] - img_s[n_y][n_x]
            dt = img_t[y][x] - img_t[n_y][n_x]
            d = ds if abs(ds) > abs(dt) else dt
            
            A[e, im2var[y][x]] = 1
            b[e] = d
            
            if im2var[n_y][n_x] != -1:
                A[e, im2var[n_y][n_x]] = -1
            else:
                b[e] += img_t[n_y][n_x]
            e += 1
    
    A = sp.sparse.csr_matrix(A)
    v = sp.sparse.linalg.lsqr(A, b)[0]
    
    img_t_out = img_t.copy()
    
    for n in range(nnz):
        y, x = ys[n], xs[n]
        img_t_out[y][x] = v[im2var[y][x]]
    
    return np.clip(img_t_out, 0, 1)

def _2d_gaussian(sigma):
    ksize = np.int(np.ceil(sigma)*6+1)
    gaussian_1d = cv2.getGaussianKernel(ksize, sigma)
    return gaussian_1d * np.transpose(gaussian_1d)

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):
    _im = img.copy()
    pyramid = []
    for d in range(depth-1):
        _im = _low_pass_filter(_im.copy(), sigma)
        pyramid.append(_im)
        _im = cv2.pyrDown(_im)
    return pyramid 

def _lap_pyramid(img, depth, sigma):
    _im = img.copy()
    pyramid = []
    for d in range(depth-1):
        lap = _high_pass_filter(_im.copy(), sigma)
        pyramid.append(lap)
        _im = cv2.pyrDown(_im)
    return pyramid 

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

def laplacian_blend(img1, img2, mask, depth=5, sigma=25):
    mask_gaus_pyramid = _gaus_pyramid(mask, depth, sigma)
    img1_lap_pyramid, img2_lap_pyramid = _lap_pyramid(img1, depth, sigma), _lap_pyramid(img2, depth, sigma)

    blended = [_blend(obj, bg, mask) for obj, bg, mask in zip(img1_lap_pyramid, img2_lap_pyramid, mask_gaus_pyramid)][::-1]
    
    h, w = blended[0].shape[:2]
    
    img1 = cv2.resize(img1, (w, h))
    img2 = cv2.resize(img2, (w, h))
    mask = cv2.resize(mask, (w, h))

    blanded_img = _blend(img1, img2, mask)
    blanded_img = cv2.resize(blanded_img, blended[0].shape[:2])
    
    imgs = []
    for d in range(0, depth-1):
        gaussian_img = _low_pass_filter(blanded_img.copy(), sigma)
        reconstructed_img = cv2.add(blended[d], gaussian_img)
        
        imgs.append(reconstructed_img)
        blanded_img = cv2.pyrUp(reconstructed_img)
        
    return np.clip(imgs[-1], 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)

    # 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()