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

def dft_matrix(N):
    """生成 N 阶 DFT 矩阵"""
    i, j = np.meshgrid(np.arange(N), np.arange(N))
    omega = np.exp(-2j * np.pi / N)
    W = np.power(omega, i * j)
    return W

def idft_matrix(N):
    """生成 N 阶 IDFT 矩阵"""
    i, j = np.meshgrid(np.arange(N), np.arange(N))
    omega = np.exp(2j * np.pi / N)
    W_inv = (1.0 / N) * np.power(omega, i * j)
    return W_inv

def main():
    # 1. 定义两个有限长序列 x[n] 和 h[n]
    x = np.array([1, 2, 3, 4])  # 长度 N1 = 4
    h = np.array([1, 1, 1])     # 长度 N2 = 3
    
    N1 = len(x)
    N2 = len(h)
    
    # 2. 线性卷积的长度 N = N1 + N2 - 1
    N = N1 + N2 - 1  # 4 + 3 - 1 = 6
    
    # 3. 将 x 和 h 补零到长度 N
    x_padded = np.pad(x, (0, N - N1), 'constant')
    h_padded = np.pad(h, (0, N - N2), 'constant')
    
    # 4. 生成 DFT 矩阵 W_N
    W_N = dft_matrix(N)
    
    # 5. 通过 DFT 矩阵计算频域表示 X(k) 和 H(k)
    X_k = np.dot(W_N, x_padded)
    H_k = np.dot(W_N, h_padded)
    
    # 6. 频域相乘 (对应时域的圆周卷积)
    Y_k = X_k * H_k
    
    # 7. 生成 IDFT 矩阵 W_N_inv 并通过矩阵乘法还原到时域
    W_N_inv = idft_matrix(N)
    y_dft = np.dot(W_N_inv, Y_k)
    
    # 取实部（消除计算中的微小虚部误差）
    y_dft = np.real(y_dft)
    
    # 8. 使用 numpy.convolve 直接计算线性卷积，用作验证
    y_linear = np.convolve(x, h)
    
    # 打印结果对比
    print("x[n]:", x)
    print("h[n]:", h)
    print(f"Zero-padded length N = {N}")
    print("y_dft (DFT Matrix computed):", np.round(y_dft, 4))
    print("y_linear (np.convolve computed):", y_linear)
    
    # 9. 绘图对比
    plt.figure(figsize=(10, 8))
    
    plt.subplot(3, 1, 1)
    plt.stem(np.arange(N1), x, basefmt="b-")
    plt.title('Sequence x[n]')
    plt.grid(True)
    
    plt.subplot(3, 1, 2)
    plt.stem(np.arange(N2), h, basefmt="b-", linefmt='g-', markerfmt='go')
    plt.title('Sequence h[n]')
    plt.grid(True)
    
    plt.subplot(3, 1, 3)
    n_y = np.arange(N)
    # 用离散点画出线性卷积的结果
    markerline1, stemlines1, baseline1 = plt.stem(n_y - 0.1, y_linear, basefmt="b-", linefmt='r-', markerfmt='ro', label='np.convolve')
    markerline2, stemlines2, baseline2 = plt.stem(n_y + 0.1, y_dft, basefmt="b-", linefmt='m-', markerfmt='m*', label='DFT Matrix Method')
    plt.title('Convolution Result y[n]')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    
    # 改变工作目录到当前脚本所在位置并保存图片
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    plt.savefig('convolution_comparison.png')
    plt.close()
    
    print("Image saved as 'convolution_comparison.png'")

if __name__ == '__main__':
    main()
