# -*- coding: utf-8 -*-
import numpy as np  # 导入NumPy库，用于数值计算
import math  # 导入math库，用于数学运算
import cv2  # 导入OpenCV库，用于图像处理
import matplotlib.pyplot as plt  # 导入matplotlib库，用于绘图

# 一维FFT实现（递归 Cooley-Tukey）
def fft_1d(x):
    N = len(x)  # 获取输入序列的长度
    if N <= 1:  # 如果序列长度小于等于1，直接返回序列
        return x
    even = fft_1d(x[::2])  # 递归计算偶数索引位置的FFT
    odd = fft_1d(x[1::2])  # 递归计算奇数索引位置的FFT
    T = [np.exp(-2j * np.pi * k / N) * odd[k] for k in range(N // 2)]  # 计算旋转因子
    return [even[k] + T[k] for k in range(N // 2)] + \
        [even[k] - T[k] for k in range(N // 2)]    # 合并结果

# 一维IFFT实现
def ifft_1d(x):
    N = len(x)  # 获取输入序列的长度
    x_conj = [np.conj(xi) for xi in x]  # 对输入序列取共轭
    y = fft_1d(x_conj)  # 对共轭序列进行FFT
    return [np.conj(yi) / N for yi in y]  # 对结果取共轭并除以N，得到IFFT结果

# 二维FFT
def fft_2d(image):
    image = image.astype(float)  # 将图像转换为浮点型
    rows, cols = image.shape  # 获取图像的行数和列数
    # 对每行进行FFT
    fft_rows = [fft_1d(list(image[i, :])) for i in range(rows)]
    # 对每列进行FFT
    fft_cols = [fft_1d([fft_rows[i][j] for i in range(rows)]) for j in range(cols)]
    return np.array(fft_cols).T  # 返回转置后的结果

# 二维IFFT
def ifft_2d(freq_domain):
    rows, cols = freq_domain.shape  # 获取频域图像的行数和列数
    # 对每列进行IFFT
    ifft_cols = [ifft_1d([freq_domain[i][j] for i in range(rows)]) for j in range(cols)]
    # 对每行进行IFFT
    ifft_rows = [ifft_1d([ifft_cols[j][i] for j in range(cols)]) for i in range(rows)]
    return np.array(ifft_rows)  # 返回IFFT结果

# 填充图像为2的整数次幂大小
def pad_image(image):
    h, w = image.shape  # 获取图像的高度和宽度
    h2 = 1 << (h - 1).bit_length()  # 计算最接近且大于等于h的2的整数次幂
    w2 = 1 << (w - 1).bit_length()  # 计算最接近且大于等于w的2的整数次幂
    padded = np.zeros((h2, w2))  # 创建填充后的零矩阵
    padded[:h, :w] = image  # 将原图像放入填充矩阵的左上角
    return padded  # 返回填充后的图像

# 主程序
def main():
    # 读取图像
    image = cv2.imread("image.jpg", cv2.IMREAD_GRAYSCALE)  # 以灰度模式读取图像
    if image is None:  # 如果图像读取失败，输出错误信息并返回
        print("图像读取失败，请确认 image.jpg 文件存在于当前目录。")
        return

    # 填充图像
    padded_image = pad_image(image)  # 对图像进行填充

    # 执行FFT和IFFT
    freq = fft_2d(padded_image)  # 对填充后的图像进行二维FFT
    restored = ifft_2d(freq).real  # 对频域图像进行二维IFFT，并取实部

    # 可视化结果
    plt.figure(figsize=(12, 4))  # 创建绘图窗口

    plt.subplot(1, 3, 1)  # 创建第一个子图
    plt.title("Original")  # 设置子图标题
    plt.imshow(image, cmap="gray")  # 显示原始图像

    plt.subplot(1, 3, 2)  # 创建第二个子图
    plt.title("FFT Spectrum")  # 设置子图标题
    plt.imshow(np.log(np.abs(freq) + 1), cmap="gray")  # 显示FFT频谱

    plt.subplot(1, 3, 3)  # 创建第三个子图
    plt.title("Restored")  # 设置子图标题
    plt.imshow(restored, cmap="gray")  # 显示恢复后的图像

    plt.tight_layout()  # 调整子图布局
    plt.show()  # 显示绘图窗口

if __name__ == "__main__":
    main()  # 运行主程序