import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft2, fftshift, ifft2, ifftshift
from imageio.v2 import imread

# 1. Load and Pre-process
imag = imread('grid.gif')

if len(imag.shape) == 3:
    imag = np.dot(imag[...,:3], [0.2989, 0.5870, 0.1140])

# 2. Fourier Calculations
imag_fft = fftshift(fft2(imag))
mag_spectrum = np.abs(imag_fft)

# Inverse Transform
imag_recon = ifft2(ifftshift(imag_fft))
recon_intensity = np.abs(imag_recon)**2

# 3. Create the Subplot Grid (2 rows, 2 columns)
fig, axes = plt.subplots(2, 2, figsize=(12, 10))

# Top Left: Original Image
axes[0, 0].imshow(imag, cmap='gray')
axes[0, 0].set_title("Original Image")

# Top Right: FFT Magnitude (Contour)
# Using a log scale or sqrt can sometimes make FFT contours easier to see
axes[0, 1].contour(mag_spectrum, 100, colors='r')
axes[0, 1].axis([150, 360, 150, 360])
axes[0, 1].set_title("FFT Magnitude (Zoomed)")

# Bottom Left: FFT Magnitude (Full view)
# Adding a 3rd plot since we have a 2x2 grid anyway
axes[1, 0].imshow(np.log(1 + mag_spectrum), cmap='magma')
axes[1, 0].set_title("FFT Log-Magnitude (Full)")

# Bottom Right: Reconstruction
axes[1, 1].imshow(recon_intensity, cmap='gray', vmin=0, vmax=255)
axes[1, 1].set_title("Inverse Fourier Transform")

# Adjust layout so titles don't overlap
plt.tight_layout()
plt.show()