# Generated by pandoc-plot 1.8.0
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(2019)

import numpy as np
import matplotlib.pyplot as plt
from random import random, seed

seed(2022)

t = np.linspace(-1, 1, 1024)
f = np.exp(1j * 2 * 2 * np.pi * t) + np.exp(1j * 2 * np.pi * 5 * t)

noise = np.zeros_like(f)
for freq in range(20, 50):
    phase = random()
    amplitude = random()
    noise += amplitude * np.sin(2*np.pi*freq*t + phase)

# For technical reasons, the frequency components of the Fourier transform
# are not arranged as one would expect. In order for the plot to look OK, we
# need to use `np.fft.fftshift`
fhat = np.fft.fftshift(np.fft.fft(f + noise))
w = np.fft.fftshift(np.fft.fftfreq(f.shape[0], d=t[1] - t[0]))

window = np.ones_like(w)
window[np.logical_or(w > 10, w < -10)] = 0
filtered = np.fft.ifft(np.fft.ifftshift(window*fhat))

fig = plt.figure(figsize=(8,9))
ax_t = plt.subplot(3,1,1)
ax_w = plt.subplot(3,1,2)
ax_filtered = plt.subplot(3,1,3, sharex=ax_t)

ax_t.plot(t, np.real(f)+noise, color='r', label="Noisy signal")
ax_t.plot(t, np.real(f), color="k", label="Pure signal")
ax_t.set_xlim([t.min(), t.max()])
ax_t.set_xlabel("Time [s]")
ax_t.legend()

# Fourier transform

ax_w.plot(w, np.abs(fhat)**2, ".", color="k")
ax_w.set_xlim([-2, 55])
ax_w.set_ylabel("Spectral power [a.u.]")
ax_w.set_xlabel("Angular frequency [rad/s]")

ax_w.axvline(x=10, linestyle='--', linewidth=1, color='k')
ax_w.annotate(f"Cutoff", xy=(1.05*10, 0.95*np.square(np.abs(fhat)).max()), xycoords='data')

ymin, ymax = ax_w.get_ylim()
xmin, xmax = ax_w.get_xlim()
ax_w.fill_betweenx(y=[ymin, ymax], x1=-2, x2=10, alpha=0.1, color='g')
ax_w.text(x=4, y=(ymax - ymin)/2, s="Band pass", color='g', va='center', ha='center')
ax_w.fill_betweenx(y=[ymin, ymax], x1=10, x2=55, alpha=0.1, color='r')
ax_w.text(x=(xmax-10)/2 + 10, y=(ymax - ymin)/2, s="Zeroed region", color='r', va='center', ha='center')
ax_w.set_ylim([ymin, ymax])

# Final output

ax_filtered.plot(t, np.real(f), color='k', label="Pure signal")
ax_filtered.plot(t, filtered, color="g", label="Filtered signal")
ax_filtered.set_xlim([t.min(), t.max()])
ax_filtered.set_xlabel("Time [s]")
ax_filtered.legend()