import numpy as np
import matplotlib.pyplot as plt
import math

plt.style.use('default')
plt.rcParams['font.family'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 1. 定义混合波形
srate = 3000
t = np.linspace(0, 1, srate)
s1 = 1.5 * np.sin(-2 * np.pi * 30 * t)
s2 = 3 * np.cos(-2 * np.pi * 5 * t)
s3 = 10 * np.sin(-2 * np.pi * 100 * t)
s4 = 20 * np.cos(-2 * np.pi * 120 * t)
m = s1 + s2 + s3 + s4
n = m + 2.5 * np.random.randn(len(t))

fcoefs = np.fft.fft(m)
amp_spec = 2 * np.abs(fcoefs / srate)

freqs = np.fft.fftfreq(len(m), 1 / srate) 
amp_spec_smoothed = np.copy(amp_spec) 
amp_spec_smoothed[amp_spec < 0.05 * np.max(amp_spec)] = 0 
amp_spec_smoothed[(np.abs(freqs) > 110)] = 0 

freq_shifted = np.fft.fftshift(freqs)
amp_shifted = np.fft.fftshift(amp_spec) 
amp_shifted_smoothed = np.fft.fftshift(amp_spec_smoothed) 

plt.figure(figsize=(36, 10))
plt.subplot(2, 2, 1)
plt.plot(m[:500])
plt.xlim(0,500)
plt.xlabel('时间')
plt.ylabel('振幅')
plt.grid()
plt.title('原始未过滤数据的复合波形时域图')



m_filtered = np.fft.ifft(fcoefs * (amp_spec_smoothed>1))

plt.subplot(2, 2, 2)
plt.plot(np.real(m_filtered)[:500])
plt.xlim(0,500)
plt.ylim(-15,15)
plt.xlabel('时间')
plt.ylabel('振幅')
plt.grid()
plt.title('双重过滤后，去除了最高频率之后的混合波形时域图')


plt.subplot(2, 2, 3)
plt.stem(freq_shifted, amp_shifted)
plt.xlim(-150, 150)
plt.xlabel('频率')
plt.ylabel('幅值')
plt.grid()
plt.title('原始波形频域图')

plt.subplot(2, 2, 4)
plt.stem(freq_shifted, amp_shifted_smoothed)
plt.xlim(-150, 150)
plt.ylim(-1,21)
plt.xlabel('频率')
plt.ylabel('幅值')
plt.grid()
plt.title('去除高频之后的波形频域图')

plt.tight_layout

plt.show()