/**
 * 手动实现离散傅里叶变换(DFT)和图像滤波
 * 使用OpenCV进行图像读取和显示，手动实现DFT/IDFT核心算法
 */

#include <opencv2/opencv.hpp>  // OpenCV核心库，用于图像处理
#include <iostream>            // 标准输入输出
#include <complex>             // 复数支持
#include <vector>              // 向量容器
#include <cmath>               // 数学函数

using namespace cv;   // OpenCV命名空间
using namespace std;  // 标准命名空间

#ifndef M_PI
#define M_PI 3.14159265358979323846  // 定义π值
#endif

/**
 * 一维离散傅里叶变换(DFT)
 * 公式: X[k] = Σ_{n=0}^{N-1} x[n] * e^{-j*2πkn/N}
 * @param input 输入复数序列
 * @return 输出复数序列(频域表示)
 */
vector<std::complex<float>> dft1D(const vector<std::complex<float>>& input) {
    int N = input.size();  // 输入序列长度
    vector<std::complex<float>> output(N);  // 输出频域序列
    
    // 对每个频率分量k进行计算
    for (int k = 0; k < N; k++) {
        std::complex<float> sum(0.0f, 0.0f);  // 累加器初始化为0
        
        // 对每个时间点n进行计算
        for (int n = 0; n < N; n++) {
            // 计算旋转因子相位角: -2πkn/N
            float angle = -2 * M_PI * k * n / N;
            
            // 计算当前点的贡献并累加: x[n] * e^{-j*2πkn/N}
            sum += input[n] * std::exp(std::complex<float>(0, angle));
        }
        output[k] = sum;  // 存储当前频率分量的结果
    }
    return output;
}

/**
 * 一维离散傅里叶逆变换(IDFT)
 * 公式: x[n] = (1/N) * Σ_{k=0}^{N-1} X[k] * e^{j*2πkn/N}
 * @param input 频域复数序列
 * @return 输出时域复数序列
 */
vector<std::complex<float>> idft1D(const vector<std::complex<float>>& input) {
    int N = input.size();  // 输入序列长度
    vector<std::complex<float>> output(N);  // 输出时域序列
    
    // 对每个时间点k进行计算
    for (int k = 0; k < N; k++) {
        std::complex<float> sum(0.0f, 0.0f);  // 累加器初始化为0
        
        // 对每个频率分量n进行计算
        for (int n = 0; n < N; n++) {
            // 计算旋转因子相位角: 2πkn/N
            float angle = 2 * M_PI * k * n / N;
            
            // 计算当前频率分量的贡献并累加: X[n] * e^{j*2πkn/N}
            sum += input[n] * std::exp(std::complex<float>(0, angle));
        }
        // 除以N完成归一化
        output[k] = sum / std::complex<float>(N, 0);
    }
    return output;
}

/**
 * 二维离散傅里叶变换(2D DFT)
 * 通过行列分离法实现: 先对每行做1D DFT，再对每列做1D DFT
 * @param input 输入图像矩阵(单通道浮点型)
 * @return 复数矩阵(CV_32FC2类型)表示频域结果
 */
Mat dft2D(const Mat& input) {
    int rows = input.rows;  // 图像行数(高度)
    int cols = input.cols;  // 图像列数(宽度)
    Mat output(rows, cols, CV_32FC2);  // 创建输出复数矩阵
    
    // 第一步: 对每行做1D DFT
    for (int i = 0; i < rows; i++) {
        // 提取当前行的像素值并转换为复数形式(虚部为0)
        vector<std::complex<float>> row(cols);
        for (int j = 0; j < cols; j++) {
            row[j] = std::complex<float>(input.at<float>(i, j), 0);
        }
        
        // 对当前行执行1D DFT
        auto rowDFT = dft1D(row);
        
        // 将结果存储回矩阵
        for (int j = 0; j < cols; j++) {
            output.at<Vec2f>(i, j) = Vec2f(rowDFT[j].real(), rowDFT[j].imag());
        }
    }
    
    // 第二步: 对每列做1D DFT
    for (int j = 0; j < cols; j++) {
        // 提取当前列的复数数据
        vector<std::complex<float>> col(rows);
        for (int i = 0; i < rows; i++) {
            col[i] = std::complex<float>(output.at<Vec2f>(i, j)[0], output.at<Vec2f>(i, j)[1]);
        }
        
        // 对当前列执行1D DFT
        auto colDFT = dft1D(col);
        
        // 将结果存储回矩阵
        for (int i = 0; i < rows; i++) {
            output.at<Vec2f>(i, j) = Vec2f(colDFT[i].real(), colDFT[i].imag());
        }
    }
    
    return output;
}

/**
 * 二维离散傅里叶逆变换(2D IDFT)
 * 通过行列分离法实现: 先对每列做1D IDFT，再对每行做1D IDFT
 * @param input 复数矩阵(CV_32FC2类型)表示频域数据
 * @return 输出图像矩阵(单通道浮点型)
 */
Mat idft2D(const Mat& input) {
    int rows = input.rows;  // 图像行数(高度)
    int cols = input.cols;  // 图像列数(宽度)
    Mat output(rows, cols, CV_32FC2);  // 创建输出复数矩阵
    
    // 第一步: 对每列做1D IDFT
    for (int j = 0; j < cols; j++) {
        // 提取当前列的复数数据
        vector<std::complex<float>> col(rows);
        for (int i = 0; i < rows; i++) {
            col[i] = std::complex<float>(input.at<Vec2f>(i, j)[0], input.at<Vec2f>(i, j)[1]);
        }
        
        // 对当前列执行1D IDFT
        auto colIDFT = idft1D(col);
        
        // 将结果存储回矩阵
        for (int i = 0; i < rows; i++) {
            output.at<Vec2f>(i, j) = Vec2f(colIDFT[i].real(), colIDFT[i].imag());
        }
    }
    
    // 第二步: 对每行做1D IDFT
    for (int i = 0; i < rows; i++) {
        // 提取当前行的复数数据
        vector<std::complex<float>> row(cols);
        for (int j = 0; j < cols; j++) {
            row[j] = std::complex<float>(output.at<Vec2f>(i, j)[0], output.at<Vec2f>(i, j)[1]);
        }
        
        // 对当前行执行1D IDFT
        auto rowIDFT = idft1D(row);
        
        // 将结果存储回矩阵
        for (int j = 0; j < cols; j++) {
            output.at<Vec2f>(i, j) = Vec2f(rowIDFT[j].real(), rowIDFT[j].imag());
        }
    }
    
    return output;
}

/**
 * 频谱中心化(将低频分量移到图像中心)
 * 通过交换四个象限实现:
 * 1. 左上↔右下
 * 2. 右上↔左下
 * @param mag 输入/输出的幅度谱矩阵(会被原地修改)
 */
static void fftShift(Mat &mag) {
    // 计算中心点坐标
    int cx = mag.cols / 2;
    int cy = mag.rows / 2;

    // 定义四个象限的ROI
    Mat q1(mag, Rect(0, 0, cx, cy));   // 第一象限(左上)
    Mat q2(mag, Rect(cx, 0, cx, cy));  // 第二象限(右上)
    Mat q3(mag, Rect(0, cy, cx, cy));  // 第三象限(左下)
    Mat q4(mag, Rect(cx, cy, cx, cy)); // 第四象限(右下)

    Mat tmp;  // 临时存储矩阵
    
    // 交换第一象限和第四象限(左上↔右下)
    q1.copyTo(tmp);
    q4.copyTo(q1);
    tmp.copyTo(q4);

    // 交换第二象限和第三象限(右上↔左下)
    q2.copyTo(tmp);
    q3.copyTo(q2);
    tmp.copyTo(q3);
}

/**
 * 创建理想低通滤波器(圆形通带)
 * 在频域中心创建一个圆形区域(值为1)，其余区域为0
 * @param size 滤波器尺寸(应与图像尺寸一致)
 * @param radius 通带半径(像素)
 * @return 单通道浮点型滤波器矩阵
 */
static Mat createLowPassFilter(Size size, int radius) {
    Mat filter = Mat::zeros(size, CV_32F);  // 创建全零矩阵
    Point center = Point(filter.cols / 2, filter.rows / 2);  // 中心点坐标
    
    // 在中心绘制实心圆(半径radius，值为1)
    circle(filter, center, radius, Scalar(1), -1);
    
    return filter;
}

/**
 * 创建理想高通滤波器(圆形阻带)
 * 在频域中心创建一个圆形区域(值为0)，其余区域为1
 * @param size 滤波器尺寸(应与图像尺寸一致)
 * @param radius 阻带半径(像素)
 * @return 单通道浮点型滤波器矩阵
 */
static Mat createHighPassFilter(Size size, int radius) {
    Mat filter = Mat::ones(size, CV_32F);  // 创建全1矩阵
    Point center = Point(filter.cols / 2, filter.rows / 2);  // 中心点坐标
    
    // 在中心绘制实心圆(半径radius，值为0)
    circle(filter, center, radius, Scalar(0), -1);
    
    return filter;
}

/**
 * 主函数：实现图像频域处理流程
 * 1. 读取并预处理图像
 * 2. 执行手动2D DFT变换
 * 3. 创建并应用低通/高通滤波器
 * 4. 执行逆变换并显示结果
 * @return 程序执行状态(0表示成功)
 */
int main() {
    // ===== 1. 图像读取与预处理 =====
    // 读取输入图像并转换为灰度图
    Mat image = imread("../input.png", IMREAD_GRAYSCALE);
    if (image.empty()) {
        cout << "Could not open or find the image" << endl;
        return -1;
    }
    resize(image, image, Size(image.cols / 3, image.rows / 3)); // 统一大小方便处理

    // ===== 2. 图像扩展 =====
    // 扩展图像到最优DFT尺寸(提高计算效率)
    // 在图像右侧和底部补零
    Mat padded;
    int m = getOptimalDFTSize(image.rows);
    int n = getOptimalDFTSize(image.cols);
    copyMakeBorder(image, padded, 0, m - image.rows, 0, n - image.cols,
                   BORDER_CONSTANT, Scalar::all(0));

    // ===== 3. 数据类型转换 =====
    // 将图像转换为32位浮点型并归一化到[0,1]范围
    // 这是DFT处理的标准输入格式
    Mat floatImg;
    padded.convertTo(floatImg, CV_32F);
    floatImg /= 255.0;

    // ===== 4. 频域变换 =====
    // 执行手动实现的2D离散傅里叶变换
    // 结果是一个复数矩阵(CV_32FC2类型)
    Mat complexImg = dft2D(floatImg);

    // ===== 5. 频谱可视化处理 =====
    // 计算幅度谱并进行对数变换增强显示效果
    // 1. 分离复数矩阵的实部和虚部
    // 2. 计算幅度谱
    // 3. 对数变换压缩动态范围
    // 4. 归一化到[0,1]范围
    // 5. 频谱中心化(低频移到中心)
    Mat planes[2];
    split(complexImg, planes);
    Mat mag;
    magnitude(planes[0], planes[1], mag);
    mag += Scalar::all(1); // 避免零值
    log(mag, mag);
    normalize(mag, mag, 0, 1, NORM_MINMAX); // 归一化便于显示
    fftShift(mag);                          // 中心化

    // ===== 6. 滤波器创建 =====
    // 创建理想低通和高通滤波器
    // 滤波器半径设为图像尺寸的1/10
    // 注意: 滤波器需要中心化以匹配频谱中心
    int radius = min(padded.rows, padded.cols) / 10;
    Mat lowPass = createLowPassFilter(padded.size(), radius);
    Mat highPass = createHighPassFilter(padded.size(), radius);
    fftShift(lowPass); // 中心化滤波器
    fftShift(highPass);

    // ===== 7. 频域滤波处理 =====
    // 7.1 低通滤波处理
    // 将单通道滤波器扩展为双通道复数滤波器
    // 在频域执行点乘运算(相当于时域卷积)
    Mat lowPassPlanes[] = {lowPass, lowPass};
    Mat lowPassFilter;
    merge(lowPassPlanes, 2, lowPassFilter);
    Mat lowPassComplex;
    multiply(complexImg, lowPassFilter, lowPassComplex); // 频域相乘

    // 7.2 高通滤波处理
    // 同样将单通道滤波器扩展为双通道复数滤波器
    // 在频域执行点乘运算
    Mat highPassPlanes[] = {highPass, highPass};
    Mat highPassFilter;
    merge(highPassPlanes, 2, highPassFilter);
    Mat highPassComplex;
    multiply(complexImg, highPassFilter, highPassComplex);

    // ===== 8. 频域逆变换 =====
    // 8.1 低通部分逆变换
    // 执行手动2D IDFT将频域结果转换回时域
    // 提取实部作为处理结果
    // 归一化到[0,255]范围并转换为8位无符号整型
    // 低通部分
    Mat invLowPass = idft2D(lowPassComplex);
    Mat invLowPassPlanes[2];
    split(invLowPass, invLowPassPlanes);
    Mat lowPassImg = invLowPassPlanes[0];
    normalize(lowPassImg, lowPassImg, 0, 255, NORM_MINMAX, CV_8U);

    // 高通部分
    Mat invHighPass = idft2D(highPassComplex);
    Mat invHighPassPlanes[2];
    split(invHighPass, invHighPassPlanes);
    Mat highPassImg = invHighPassPlanes[0];
    normalize(highPassImg, highPassImg, 0, 255, NORM_MINMAX, CV_8U);

    // ===== 9. 结果后处理 =====
    // 裁剪处理结果到原始图像尺寸
    // 去除DFT扩展的填充区域
    Rect roi(0, 0, image.cols, image.rows);
    Mat lowPassResult = lowPassImg(roi);
    Mat highPassResult = highPassImg(roi);

    // ===== 10. 结果显示 =====
    // 显示原始图像、频谱图、低通和高通滤波结果
    // 按任意键退出
    imshow("Original Image", image);
    imshow("Frequency Spectrum", mag);
    imshow("Low Frequency Component", lowPassResult);
    imshow("High Frequency Component", highPassResult);

    waitKey(0);
    return 0;
}