import numpy as np
import matplotlib.pyplot as plt

# 1. Setup Data
x = np.arange(1024)
y = np.zeros(1024)
ymask = np.zeros(100)
ymask[16:26] = 1
ymask[46:56] = 1
ymask[76:86] = 1
y[463:563] = ymask

# 2. Fourier Calculations
fty = np.fft.fft(y)
sfty = np.fft.fftshift(fty)
Y_orig = np.abs(sfty) ** 2

# 3. Create a single figure with 2 subplots (1 row, 2 columns)
fig, (ax_spatial, ax_freq) = plt.subplots(1, 2, figsize=(14, 5))

# Plot initial y(x) on the spatial axis
ax_spatial.plot(x, y, '-k', label='Original Signal')
ax_spatial.axis([400, 600, -0.2, 1.2])
ax_spatial.set_xlabel('x')
ax_spatial.set_ylabel('y')
ax_spatial.set_title('Spatial Domain: y(x)')

# Plot initial FFT on the frequency axis
Ymax = 1.2 * np.max(Y_orig)
ax_freq.plot(x, Y_orig, '-k', label='Original Spectrum')
ax_freq.axis([200, 800, 0, Ymax])
ax_freq.set_xlabel('Frequency')
ax_freq.set_ylabel('|Y(f)|^2')
ax_freq.set_title('Frequency Domain: Fourier Transform')

# 4. Define and call the lowpass function
def lowpass(x, sfty_in, ax_s, ax_f):
    # Work on a copy to keep the original spectrum intact
    sfty_filtered = sfty_in.copy()
    
    # Apply low-pass: Zero out the "outer" high frequencies
    sfty_filtered[:462] = 0
    sfty_filtered[563:] = 0
    
    # Calculate Magnitude for plotting
    YY = np.abs(sfty_filtered) ** 2
    ax_f.plot(x, YY, '-r', label='Low-pass Filtered', alpha=0.7)
    ax_f.legend()
    
    # Calculate Inverse Transform
    # Remember to shift back before applying ifft
    yy = np.fft.ifft(np.fft.ifftshift(sfty_filtered))
    
    # Plot reconstructed signal (the smoothing effect)
    ax_s.plot(x, np.abs(yy) ** 2, '-r', label='Smoothed Signal', alpha=0.8)
    ax_s.legend()
    
    return YY, yy

# Call the function
lowpass(x, sfty, ax_spatial, ax_freq)

plt.tight_layout()
plt.show()