import numpy as np
from scipy.signal import convolve2d
from scipy.sparse import lil_matrix
from scipy.sparse.linalg import spsolve
import cv2
import matplotlib.pyplot as plt
from pathlib import Path
from scipy.ndimage import correlate


"""
https://setosa.io/ev/image-kernels/
https://openaccess.thecvf.com/content/CVPR2021/papers/Tran_Explore_Image_Deblurring_via_Encoded_Blur_Kernel_Space_CVPR_2021_paper.pdf
"""

def show(img):
    cv2.imshow('image',img.astype(np.uint8))
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def demo(image_file):
    # Define 2D image and kernel
    image = cv2.imread(image_file, 0)
    image = cv2.resize(image, (200, 200), interpolation= cv2.INTER_LINEAR)

    kernel = np.array([
        [1, 2, 1],
        [2, 4, 2],
        [1, 2, 1]
    ], dtype=np.float32) 
    kernel /= kernel.sum()  # Normalize

    print(kernel)

    # Perform 2D convolution (blurring)
    blurred = convolve2d(image, kernel, mode="same", boundary="fill", fillvalue=0)

    h, w = image.shape
    kh, kw = kernel.shape
    pad_h, pad_w = kh // 2, kw // 2


    show(image)
    show(blurred)

    print("Original image:\n", image)
    print("\nBlurred image:\n", blurred)

    print("\nBuilding linear system for deconvolution...")

    # Step 2: Build sparse matrix A
    N = h * w
    A = lil_matrix((N, N), dtype=np.float32)
    b = blurred.flatten()

    def index(y, x):
        return y * w + x

    for y in range(h):
        for x in range(w):
            row_idx = index(y, x)
            for ky in range(kh):
                for kx in range(kw):
                    iy = y + ky - pad_h
                    ix = x + kx - pad_w
                    if 0 <= iy < h and 0 <= ix < w:
                        col_idx = index(iy, ix)
                        A[row_idx, col_idx] += kernel[ky, kx]

    # Step 3: Solve the sparse system A * x = b
    x = spsolve(A.tocsr(), b)
    deblurred = x.reshape((h, w))

    print("\nDeblurred image:\n", np.round(deblurred, 2))

    show(deblurred)


def get_mask(image_file):
    mask_file = Path(image_file)
    mask_file = mask_file.with_name("mask_" + mask_file.name)

    if mask_file.exists():
        return cv2.imread(str(mask_file), 0)

    drawing = False  # True when mouse is pressed
    brush_size = 5
    image = cv2.imread(image_file)
    mask = np.zeros(image.shape[:2], dtype=np.uint8)
    clone = image.copy()

    def draw_mask(event, x, y, flags, param):
        nonlocal drawing, mask, brush_size

        if event == cv2.EVENT_LBUTTONDOWN:
            drawing = True
        elif event == cv2.EVENT_MOUSEMOVE:
            if drawing:
                cv2.circle(mask, (x, y), brush_size, 255, -1)
                cv2.circle(image, (x, y), brush_size, (0, 0, 255), -1)
        elif event == cv2.EVENT_LBUTTONUP:
            drawing = False


    cv2.namedWindow("Draw Mask")
    cv2.setMouseCallback("Draw Mask", draw_mask)

    while True:
        display = image.copy()
        cv2.putText(display, f'Brush size: {brush_size}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)
        cv2.imshow("Draw Mask", display)
        key = cv2.waitKey(1) & 0xFF
        if key == 13:  # Enter to finish
            break
        elif key == ord('+') or key == ord('='):  # `=` for some keyboard layouts
            brush_size = min(100, brush_size + 1)
        elif key == ord('-') or key == ord('_'):
            brush_size = max(1, brush_size - 1)


    cv2.destroyAllWindows()

    cv2.imwrite(str(mask_file), mask)

    # Apply mask
    masked_image = cv2.bitwise_and(clone, clone, mask=mask)

    cv2.imshow("Masked Image", masked_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()



def color_edge_detection(img, threshold=30):
    # Load image
    img_lab = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)

    # Split Lab channels
    L, A, B = cv2.split(img_lab)

    # Compute gradients using Sobel for each channel
    def gradient_magnitude(channel):
        gx = cv2.Sobel(channel, cv2.CV_64F, 1, 0, ksize=3)
        gy = cv2.Sobel(channel, cv2.CV_64F, 0, 1, ksize=3)
        return gx, gy

    gxL, gyL = gradient_magnitude(L)
    gxA, gyA = gradient_magnitude(A)
    gxB, gyB = gradient_magnitude(B)

    # Combine gradients across channels
    gx_total = gxL**2 + gxA**2 + gxB**2
    gy_total = gyL**2 + gyA**2 + gyB**2
    magnitude = np.sqrt(gx_total + gy_total)

    # Normalize and threshold
    magnitude = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX)
    edges = (magnitude > threshold).astype(np.uint8) * 255

    return edges



def deconvolution(image_file):
    
    blurred = cv2.imread(image_file)
    # blurred = cv2.resize(blurred, (200, 200), interpolation= cv2.INTER_LINEAR)
    mask = get_mask(image_file)

    edges = color_edge_detection(blurred, threshold=8)
    # edges = cv2.bitwise_and(edges, edges, mask=mask)
    show(edges)




if __name__ == "__main__":
    img_file = "assets/real_test.jpg"

    #demo("assets/omas.png")
    deconvolution(img_file)