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) 

# # 定义去噪函数，将振幅小于阈值的噪声去除
# def denoise(arr, thresh):
#     mask = arr > thresh
#     return mask * arr

# 1. 获得混合波形
srate=3000
t=np.linspace(0,1,srate)
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
#s_noise = s_clean + 2.5 * np.random.randn(len(t))

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

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

# 把频率轴从0~1000 转变成 0~499 然后 -500~-1
freqs = np.fft.fftfreq(len(amp_list), 1/srate)
# 然后把 频率轴 和 数据 都变成 0hz 在中间，向左是负频率，向右是正频率的形式
amp_shifted=np.fft.fftshift(amp_list)
freq_shift=np.fft.fftshift(freqs)

#双重去噪，过滤数据
amp_list_copy = np.copy(amp_list)
mask = amp_list_copy > 1
fCoefs_cl = mask * fCoefs
amp_list_copy[(freqs > 110) | (freqs < -110)] = 0
#将fCoefs_cl数组中与amp_list_copy数组中值为0的位置相对应的元素的值设置为0。
x = np.where(amp_list_copy == 0)[0]
fCoefs_cl[x] = 0

amp_list_cl = 2*np.abs(fCoefs_cl / srate)
amp_cl = np.fft.ifft(fCoefs_cl)
amp_shifted_cl = np.fft.fftshift(amp_list_cl)

# 把振幅列表画出来
fg,ax = plt.subplots(2,2,figsize = (20,5))

#原始未过滤数据的复合波形时域图
ax[0][0].plot(m[0:500])
ax[0][0].set_xlim([0,500])
ax[0][0].grid()
ax[0][0].set_title('原始未过滤数据的复合波形时域图')

#双重过滤后的复合波形时域图
ax[0][1].plot(np.real(amp_cl[0:500]))
ax[0][1].set_xlim([0,500])
ax[0][1].set_ylim([-15,15])
ax[0][1].grid()
ax[0][1].set_title('双重过滤后的复合波形时域图')

#原始未过滤数据的复合波形频域图
ax[1][0].stem(freq_shift,amp_shifted)
ax[1][0].set_xlim([-150,150])
ax[1][0].grid()
ax[1][0].set_title('原始未过滤数据的复合波形频域图')

#双重过滤后的复合波形频域图
ax[1][1].stem(freq_shift,amp_shifted_cl)
ax[1][1].set_xlim([-150,150])
ax[1][1].set_ylim([-1,21])
ax[1][1].grid()
ax[1][1].set_title('双重过滤后的复合波形频域图')

#显示图像
plt.tight_layout()
plt.show()