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, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

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

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

# 4. Define and call the highpass function
def highpass(x, sfty, ax_spatial, ax_freq):
    # Apply filter: Set low-frequency center components to zero
    # (Note: This is technically a "stop-band" or high-pass since center = 0 frequency)
    sfty_filtered = sfty.copy()
    sfty_filtered[403:622] = 0
    
    # Calculate Magnitude for plotting
    YY = np.abs(sfty_filtered) ** 2
    ax_freq.plot(x, YY, '-r', label='Filtered FFT', alpha=0.7)
    ax_freq.legend()
    
    # Calculate Inverse Transform
    yy = np.fft.ifft(np.fft.ifftshift(sfty_filtered))
    
    # Plot reconstructed signal on the spatial axis
    ax_spatial.plot(x, np.abs(yy) ** 2, '-r', label='High-pass Filtered', alpha=0.8)
    ax_spatial.legend()
    
    return YY, yy

# Call the function with the axes references
highpass(x, sfty, ax1, ax2)

plt.tight_layout()
plt.show()