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

# 生成示例信号：正弦波信号加上高斯白噪声
np.random.seed(0)
t = np.linspace(0, 1, 1000, endpoint=False)  # 时间序列，长度为1000
signal = np.sin(2 * np.pi * 20 * t)  # 20 Hz 的正弦波信号
noise = np.random.normal(0, 0.5, len(t))  # 高斯白噪声
signal_noisy = signal + noise  # 加噪声后的信号

# 进行小波变换
wavelet = 'db4'  # 选择小波基函数
levels = 5  # 分解的层数
coeffs = pywt.wavedec(signal_noisy, wavelet, level=levels)

# 对小波系数进行软阈值处理
threshold = 0.5  # 阈值
coeffs_thresh = [pywt.threshold(c, threshold, mode='soft') for c in coeffs]

# 重构信号
signal_denoised = pywt.waverec(coeffs_thresh, wavelet)

# 绘制结果
plt.figure(figsize=(10, 6))

plt.subplot(3, 1, 1)
plt.plot(t, signal, 'b', label='Original Signal')
plt.title('Original Signal')
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.legend()

plt.subplot(3, 1, 2)
plt.plot(t, signal_noisy, 'g', label='Noisy Signal')
plt.title('Noisy Signal')
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.legend()

plt.subplot(3, 1, 3)
plt.plot(t, signal_denoised, 'r', label='Denoised Signal')
plt.title('Denoised Signal')
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.legend()

plt.tight_layout()
plt.show()