import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import ricker, firwin, filtfilt

# 生成Ricker子波
fs = 500  # 采样率
T = 0.5   # 信号时长（秒）
n = int(T * fs)
t = np.linspace(-0.25, 0.25, n, endpoint=False)
f_peak = 10  # 主频10Hz
ricker_wavelet = ricker(n, f_peak * T / 2)
ricker_wavelet = ricker_wavelet / np.max(np.abs(ricker_wavelet))  # 归一化

# 设计高通滤波器
cutoff = 5  # 截止频率
numtaps = 101  # 滤波器阶数
taps = firwin(numtaps, cutoff, pass_zero='highpass', fs=fs)

# 零相位滤波
filtered = filtfilt(taps, 1.0, ricker_wavelet)

# 截断至原支集范围（示例：中心附近）
window = np.zeros_like(filtered)
win_len = len(ricker_wavelet) // 4  # 假设原支集为总长的1/4
center = len(filtered) // 2
window[center - win_len//2 : center + win_len//2] = 1
filtered_windowed = filtered * window

# 绘图比较
plt.figure()
plt.plot(t, ricker_wavelet, label='Original')
plt.plot(t, filtered_windowed, label='Filtered')
plt.legend()
plt.title('Time Domain')
plt.xlabel('Time (s)')

# 频谱比较
fft_orig = np.fft.fft(ricker_wavelet)
fft_filt = np.fft.fft(filtered_windowed)
freq = np.fft.fftfreq(n, d=1/fs)
plt.figure()
plt.semilogy(freq[:n//2], np.abs(fft_orig[:n//2]), label='Original')
plt.semilogy(freq[:n//2], np.abs(fft_filt[:n//2]), label='Filtered')
plt.xlim(0, 20)
plt.axvline(5, color='red', linestyle='--', label='5Hz')
plt.legend()
plt.title('Frequency Domain')
plt.xlabel('Frequency (Hz)')
plt.show()