import numpy as np
import matplotlib.pyplot as plt

# 1. Setup the field
imag = np.zeros((512, 512), dtype=np.complex128) + 15

# Add bacteria (phase shift)
bact = 3
while bact > 0:
    X1, Y1 = np.random.randint(0, 400, size=2)
    X2, Y2 = X1 + 50 + np.random.randint(50), Y1 + 50 + np.random.randint(50)
    imag[X1:X2, Y1:Y2] *= (1 - 1j * 0.1 * (1 + np.random.rand()))
    bact -= 1

# Add dirt (amplitude drop)
dirt = 8
while dirt > 0:
    X1, Y1 = np.random.randint(0, 500, size=2)
    X2, Y2 = X1 + np.random.randint(10), Y1 + np.random.randint(10)
    imag[X1:X2, Y1:Y2] = 15 * np.random.rand()
    dirt -= 1

# 2. Fourier Processing
# Original FFT
imag1 = np.fft.fftshift(np.fft.fft2(imag))
mag_orig = np.abs(imag1)**0.5

# Apply Phase Filter
imag1_filtered = imag1.copy()
imag1_filtered[254:260, 254:260] *= (-1j)
mag_filt = np.abs(imag1_filtered)**0.5

# Inverse Transform
ifft_result = np.fft.ifft2(np.fft.ifftshift(imag1_filtered))

# 3. Plotting in a single 2x2 grid
fig, axes = plt.subplots(2, 2, figsize=(12, 10))

# Top Left: Original Intensity
axes[0, 0].imshow(np.abs(imag)**2, cmap='gray', vmin=0, vmax=255)
axes[0, 0].set_title("Original Intensity")

# Top Right: Original FFT Magnitude
axes[0, 1].contour(mag_orig, levels=100, colors='r')
axes[0, 1].axis([150, 360, 150, 360])
axes[0, 1].set_title("Magnitude of FFT")

# Bottom Left: Filtered FFT Magnitude
axes[1, 0].contour(mag_filt, levels=100, colors='b')
axes[1, 0].axis([150, 360, 150, 360])
axes[1, 0].set_title("Magnitude of Filtered FFT")

# Bottom Right: Final Result
axes[1, 1].imshow(np.abs(ifft_result)**2, cmap='gray', vmin=0, vmax=255)
axes[1, 1].set_title("Reconstructed (Filtered) Image")

plt.tight_layout() # Adjusts spacing so titles don't overlap
plt.show()