import numpy as np
from scipy.signal import convolve2d
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)

show(image)
show(blurred)

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

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

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

# Build matrix A and vector b for Ax = b
A = []
b = []

for y in range(h):
    for x in range(w):
        row = np.zeros((h, w), dtype=np.float32)

        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:
                    row[iy, ix] += kernel[ky, kx]

        A.append(row.flatten())
        b.append(blurred[y, x])

A = np.array(A)
b = np.array(b)

# Solve for the deblurred image
deblurred_flat = np.linalg.solve(A, b)
deblurred = deblurred_flat.reshape((h, w))

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

show(deblurred)