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

# 定义一个函数，用于生成正弦波
def getSin(amp, freq, phase, sampleList):
    return amp * np.sin(2 * math.pi * freq * sampleList + phase)

# 定义一个函数，用于生成余弦波
def getCos(amp, freq, phase, sampleList):
    return amp * np.cos(2 * math.pi * freq * sampleList + phase)

# 1. 获得混合波形
srate = 3000           # 设置采样率为 3000
t = np.linspace(0, 1, srate)     # 在 0 到 1 的范围内生成等间距的采样点

# 分别生成正、余弦波
s1 = getSin(amp=1.5, freq=-30, phase=0, sampleList=t)
s2 = getCos(amp=3, freq=-5, phase=0, sampleList=t)
s3 = getSin(amp=10, freq=-100, phase=0, sampleList=t)
s4 = getCos(amp=20, freq=-120, phase=0, sampleList=t)
m= s1+s2+s3+s4             # 将正、余弦波相加，得到混合波形

# 2. 获得傅里叶系数
fCoefs = np.fft.fft(m,srate)

# 3. 获得振幅列表：每一个绕线的重心到原点的距离
amp_list=2*np.abs(fCoefs/srate)

# 把频率轴从0~3000 转变成 0~1500 然后 -1500~-1
freqs = np.fft.fftfreq(len(amp_list), 1/srate)

fg,ax=plt.subplots(1,2,figsize=(15,6),sharex=True,sharey=True)

# 然后把 频率轴 和 数据 都变成 0hz 在中间，向左是负频率，向右是正频率的形式
amp_shifted=np.fft.fftshift(amp_list)
freq_shift=np.fft.fftshift(freqs)

amp_list1=np.copy(amp_list)
mask = (amp_list1>0.5)
list = mask*amp_list1        #大于1的归0
list[:-110][110:]=0         #选取大于110hz为0

amp_shifted1=np.fft.fftshift(list)
freq_shift1=np.fft.fftshift(freqs)

ax[0].stem(freq_shift,amp_shifted)
ax[0].set_xlim([-150,150])
ax[0].set_title("原始振幅谱")

ax[1].stem(freq_shift1,amp_shifted1)
ax[1].tick_params(labelleft=True)
ax[1].set_title("去噪后的振幅谱")
for ax in ax:
    ax.grid(True)
plt.tight_layout()
plt.show()