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

def main():
    # ==========================================
    # 1. 系统参数设置与模拟滤波器设计
    # ==========================================
    order = 2           # 为了明显展示频谱混叠，使用较低的阶数 (2阶)
    fc = 200.0          # 模拟滤波器截止频率: 200 Hz
    fs = 1000.0         # 采样频率: 1000 Hz
    T = 1.0 / fs        # 采样周期
    nyq = 0.5 * fs      # 奈奎斯特频率: 500 Hz
    wc = 2 * np.pi * fc # 模拟角频率 (rad/s)

    # 设计一个 2 阶的模拟巴特沃斯低通滤波器 H(s)
    b_a, a_a = signal.butter(order, wc, analog=True)

    # ==========================================
    # 2. 冲击响应不变法 (Impulse Invariance)
    # ==========================================
    # SciPy 的 cont2discrete('impulse') 方法会返回离散系统 H(z)
    # 为了保证数字滤波器的直流增益与模拟滤波器一致，分子系数乘以 T
    sys_imp = signal.cont2discrete((b_a, a_a), dt=T, method='impulse')
    b_imp = sys_imp[0][0] * T
    a_imp = sys_imp[1]

    # ==========================================
    # 3. 双线性变换法 (Bilinear Transform)
    # ==========================================
    # 由于双线性变换存在频率畸变，必须对模拟截止频率进行预畸变 (Pre-warping)
    wd = 2 * np.pi * (fc / fs)                 # 数字截止角频率 (rad)
    wa_pre = (2.0 / T) * np.tan(wd / 2.0)      # 预畸变后的模拟角频率 (rad/s)
    
    # 按照预畸变频率设计新的模拟滤波器，然后再进行双线性变换
    b_a_pre, a_a_pre = signal.butter(order, wa_pre, analog=True)
    sys_bil = signal.cont2discrete((b_a_pre, a_a_pre), dt=T, method='bilinear')
    b_bil = sys_bil[0][0]
    a_bil = sys_bil[1]

    # ==========================================
    # 4. 计算频率响应 (Frequency Response)
    # ==========================================
    # 计算模拟滤波器的频率响应 (范围 0 到 fs，用于观察高频衰减)
    f_a = np.linspace(0, fs, 2000)
    w_a = 2 * np.pi * f_a
    _, h_a = signal.freqs(b_a, a_a, worN=w_a)

    # 计算两种数字滤波器的频率响应 (范围 0 到 Nyquist)
    f_d = np.linspace(0, nyq, 1000)
    w_d = 2 * np.pi * (f_d / fs)
    _, h_imp = signal.freqz(b_imp, a_imp, worN=w_d)
    _, h_bil = signal.freqz(b_bil, a_bil, worN=w_d)

    # ==========================================
    # 5. 计算时域冲击响应 (Impulse Response)
    # ==========================================
    # 模拟滤波器的冲击响应 h_a(t)
    t_a, y_a = signal.impulse((b_a, a_a), T=np.linspace(0, 0.015, 500))

    # 数字滤波器的冲击响应
    # 输入一个离散单位脉冲信号 [1, 0, 0, ...]
    n = np.arange(15)
    t_d_n = n * T
    impulse_seq = np.zeros(len(n))
    impulse_seq[0] = 1.0
    
    # 因为数字脉冲响应理论上满足 h[n] = T * h_a(nT)，所以除以 T 还原回模拟幅值以便对比
    y_imp = signal.lfilter(b_imp, a_imp, impulse_seq) / T
    y_bil = signal.lfilter(b_bil, a_bil, impulse_seq) / T

    # ==========================================
    # 6. 可视化对比
    # ==========================================
    plt.figure(figsize=(12, 10))

    # ---------------- 频域幅度响应对比 ----------------
    plt.subplot(2, 1, 1)
    # 避免 log10(0) 警告
    h_a_db = 20 * np.log10(np.maximum(np.abs(h_a), 1e-10))
    h_imp_db = 20 * np.log10(np.maximum(np.abs(h_imp), 1e-10))
    h_bil_db = 20 * np.log10(np.maximum(np.abs(h_bil), 1e-10))

    plt.plot(f_a, h_a_db, 'k--', linewidth=2, label='Analog Filter $H_a(s)$')
    plt.plot(f_d, h_imp_db, 'r-', linewidth=2, label='Impulse Invariance')
    plt.plot(f_d, h_bil_db, 'b-', linewidth=2, label='Bilinear Transform')
    
    plt.axvline(fc, color='g', linestyle=':', label=f'Cutoff Freq ({fc} Hz)')
    plt.axvline(nyq, color='gray', linestyle='-.', label=f'Nyquist Freq ({nyq} Hz)')
    
    plt.title('Frequency Response Comparison (Magnitude)')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Magnitude (dB)')
    plt.xlim(0, fs)
    plt.ylim(-40, 5)
    plt.legend(loc='upper right')
    plt.grid(True)

    # ---------------- 时域冲击响应对比 ----------------
    plt.subplot(2, 1, 2)
    plt.plot(t_a, y_a, 'k-', linewidth=2, label='Analog Impulse Response $h_a(t)$')
    
    # 冲击响应不变法的离散点（红色）应该与模拟曲线完美重合
    markerline1, stemlines1, baseline1 = plt.stem(t_d_n, y_imp, linefmt='r-', markerfmt='ro', basefmt='r-', label='Impulse Invariance $h_{imp}[n]/T$')
    
    # 双线性变换法的离散点（蓝色，稍作时间偏移以便看清）会偏离模拟曲线
    markerline2, stemlines2, baseline2 = plt.stem(t_d_n + T/10, y_bil, linefmt='b-', markerfmt='b*', basefmt='b-', label='Bilinear Transform $h_{bil}[n]/T$')
    
    plt.title('Impulse Response Comparison in Time Domain')
    plt.xlabel('Time (s)')
    plt.ylabel('Amplitude')
    plt.xlim(0, 0.015)
    plt.legend(loc='upper right')
    plt.grid(True)

    plt.tight_layout()
    
    # 保存图片
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    plt.savefig('iir_transform_methods_sim.png', dpi=300)
    plt.close()

    print("Simulation finished. Image saved as 'iir_transform_methods_sim.png'")

if __name__ == '__main__':
    main()
