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

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


# Define 2D image and kernel
image = cv2.imread('assets/omas.png', 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)