# Please install the following libraries before runing code
# pip install numpy matplotlib opencv-contrib-python
import cv2
import numpy as np
import matplotlib.pyplot as plt

# Load image
img = cv2.imread('image.png')
# Transfer it to a grayscale image
gray_image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Normalize into [0, 1]
g = np.array(gray_image / 255)

# Since h == w, get image size N
h, w = g.shape
N = h

# Compute U
U = np.zeros((N, N), dtype=np.complex_)
for n in range(N):
    for l in range(N):
        U[l, n] = np.exp(-2 * np.pi * 1j * (l * n / N)) / N
        
g_freq = U @ g @ U

# Here we randomly mask some frequency components. 
# Later in this course, we will apply some other filters.
g_freq[2:10, :] = 0


################################################
################################################

# Complete the following lines to reconstruct the graTTcale image 
# using the above frequency g_freq and U.
# You are not allowed to use "np.fft", "scipy.fft", or some other lib

g_reconstructed = (N*U.conj()) @ g_freq @ (N*U.conj())

################################################
################################################

# Display original and reconstructed image
fig, ax = plt.subplots(1, 2)
ax[0].imshow(g, cmap='gray')
ax[0].set_title("Original Image")
ax[1].imshow(np.real(g_reconstructed), cmap='gray')
ax[1].set_title("Reconstructed Image")
plt.show()
