import numpy as np
import matplotlib.pyplot as plt
from dtk.process import butterworth, freq_spectrum

sample_rate = 1000  # Hz
duration = 10.0  # seconds
time = np.linspace(0.0, duration, num=int(sample_rate*duration) + 1)
white_noise = np.random.normal(0.0, 1.0, size=len(time))
cutoff = 200  # Hz
order = 4
filtered = butterworth(white_noise, cutoff, sample_rate, order=order)

freq, amp = freq_spectrum(white_noise, sample_rate)
freq_filt, amp_filt = freq_spectrum(filtered, sample_rate)

fig, ax = plt.subplots(layout='constrained')
ax.plot(freq, amp, label='Unfiltered')
ax.plot(freq_filt, amp_filt, alpha=0.75, label='Filtered')
ax.axvline(cutoff, color='black')
ax.set_ylabel('Amplitude of White Noise with STD=1')
ax.set_xlabel('Frequency [Hz]')
msg = 'Sample rate: {} Hz, Cutoff: {} Hz, Order: {}'
ax.set_title(msg.format(sample_rate, cutoff, order))
ax.legend()