import numpy as np
import matplotlib.pyplot as plt

"""
fft2 playground.
"""

# Initialise an empty array
field = np.empty([29, 481])

# Set the wave amplitude
amp = 2.5

# Create the synthetic sinusoidal field
for n, i in enumerate(np.linspace(0, 2 * np.pi, field.shape[0])):
    for m, j in enumerate(np.linspace(0, 2 * np.pi, field.shape[1])):
        field[n, m] = amp * np.sin(j)

# Perform 2D fourier and shift the result to centre
f = np.fft.fft2(field)
fshift = np.fft.fftshift(f)

# Calculate the magnitude and phase spectra
magnitude_spectrum = 20*np.log(np.abs(fshift))
phase_spectrum = np.angle(fshift)

# Reconstruct the initial field
f_ishift = np.fft.ifftshift(fshift)
re_field = np.abs(np.fft.ifft2(f_ishift))

# Plot
fig = plt.figure()

fig.add_subplot(411)
plt.imshow(field, cmap='gray')
plt.title('Field'), plt.xticks([]), plt.yticks([])
plt.colorbar()

fig.add_subplot(412)
plt.imshow(magnitude_spectrum, cmap='gray')
plt.title('Magnitude spectrum'), plt.xticks([]), plt.yticks([])

fig.add_subplot(413)
plt.imshow(phase_spectrum, cmap='gray')
plt.title('Phase spectrum'), plt.xticks([]), plt.yticks([])

fig.add_subplot(414)
plt.imshow(re_field, cmap='gray')
plt.title('Reconstructed field'), plt.xticks([]), plt.yticks([])
plt.colorbar()

plt.show()