import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import os

def main():
    # 1. 系统参数设置
    fs = 1000.0       # 采样频率 1000 Hz
    nyq = 0.5 * fs    # 奈奎斯特频率 500 Hz
    t = np.arange(0, 1.0, 1.0/fs)  # 1秒的时间序列
    
    # 2. 生成模拟信号 (有用信号 + 高频噪声)
    # 有用信号：10 Hz 的正弦波
    # 噪声：100 Hz 和 250 Hz 的正弦波
    f_signal = 10
    f_noise1 = 100
    f_noise2 = 250
    x_clean = np.sin(2 * np.pi * f_signal * t)
    x_noisy = x_clean + 0.6 * np.sin(2 * np.pi * f_noise1 * t) + 0.4 * np.sin(2 * np.pi * f_noise2 * t)
    
    # 3. 滤波器技术指标
    cutoff_freq = 50.0  # 截止频率设为 50 Hz (保留 10 Hz，滤除 100/250 Hz)
    normal_cutoff = cutoff_freq / nyq  # 归一化截止频率 (0 到 1 之间)
    
    # =================================================================
    # 4. 数字滤波器设计 - IIR (无限长脉冲响应)
    # 使用 4 阶巴特沃斯 (Butterworth) 低通滤波器
    # 特点：通带内最平坦，不需要极高的阶数就能达到较好的阻带衰减
    # =================================================================
    order_iir = 4
    b_iir, a_iir = signal.butter(order_iir, normal_cutoff, btype='low', analog=False)
    
    # =================================================================
    # 5. 数字滤波器设计 - FIR (有限长脉冲响应)
    # 使用窗函数法设计 (Hamming 窗)
    # 特点：可以实现严格的线性相位，但为了达到同样的截止效果需要比 IIR 高得多的阶数
    # =================================================================
    numtaps = 101  # 滤波器抽头数 (阶数 + 1)
    b_fir = signal.firwin(numtaps, normal_cutoff, window='hamming')
    a_fir = 1.0    # FIR 滤波器的分母系数为 1
    
    # 6. 对含噪信号进行滤波
    # 使用 filtfilt 进行零相位滤波，避免信号延迟
    y_iir = signal.filtfilt(b_iir, a_iir, x_noisy)
    y_fir = signal.filtfilt(b_fir, a_fir, x_noisy)
    
    # 7. 计算滤波器的频率响应
    w_iir, h_iir = signal.freqz(b_iir, a_iir, worN=8000)
    w_fir, h_fir = signal.freqz(b_fir, a_fir, worN=8000)
    
    # 将角频率转换为物理频率 (Hz)
    freq_hz_iir = (w_iir / np.pi) * nyq
    freq_hz_fir = (w_fir / np.pi) * nyq
    
    # 8. 可视化绘图
    plt.figure(figsize=(14, 10))
    
    # 图 1：时域信号对比
    plt.subplot(2, 1, 1)
    plt.plot(t, x_noisy, label='Noisy Signal', color='lightgray')
    plt.plot(t, x_clean, label='Clean Signal (10 Hz)', color='black', linestyle='--')
    plt.plot(t, y_iir, label=f'IIR Filtered (Butterworth, order={order_iir})', color='red', alpha=0.8)
    plt.plot(t, y_fir, label=f'FIR Filtered (Hamming, taps={numtaps})', color='blue', alpha=0.8)
    
    plt.title('Time Domain: Signal Filtering (IIR vs FIR)')
    plt.xlabel('Time [sec]')
    plt.ylabel('Amplitude')
    plt.xlim(0, 0.4)  # 只显示前 0.4 秒便于观察波形
    plt.legend()
    plt.grid(True)
    
    # 图 2：频域幅度响应对比
    plt.subplot(2, 1, 2)
    # 避免 log10(0) 的警告
    h_iir_db = 20 * np.log10(np.maximum(np.abs(h_iir), 1e-10))
    h_fir_db = 20 * np.log10(np.maximum(np.abs(h_fir), 1e-10))
    
    plt.plot(freq_hz_iir, h_iir_db, label=f'IIR (Butterworth, order={order_iir})', color='red')
    plt.plot(freq_hz_fir, h_fir_db, label=f'FIR (Hamming, taps={numtaps})', color='blue')
    
    plt.title('Frequency Response (Magnitude)')
    plt.xlabel('Frequency [Hz]')
    plt.ylabel('Magnitude [dB]')
    plt.axvline(cutoff_freq, color='green', linestyle='--', label=f'Cutoff Frequency ({cutoff_freq} Hz)')
    plt.axhline(-3, color='gray', linestyle=':', label='-3 dB (Half Power)')
    plt.xlim(0, nyq)
    plt.ylim(-100, 5)
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    
    # 改变工作目录并保存图片
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    plt.savefig('digital_filter_design_sim.png', dpi=300)
    plt.close()
    
    print("Simulation finished. Image saved as 'digital_filter_design_sim.png'")

if __name__ == '__main__':
    main()
