#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/imgcodecs.hpp>
#include <iostream>
using namespace std;

// 四重循环慢得要死，怎么优化呢？
static cv::Mat myDft(const cv::Mat& mat)
{
    CV_Assert(CV_64FC1 == mat.type());
    const int cols = mat.cols, rows = mat.rows;

    cv::Mat res_real(cv::Size(cols, rows), CV_64FC1);  // 保存实部
    cv::Mat res_imag(cv::Size(cols, rows), CV_64FC1);  // 保存虚部

    for (int v = 0; v < rows; ++v)                     // 累加子项
    {
        double* p_real = res_real.ptr<double>(v);
        double* p_imag = res_imag.ptr<double>(v);
        for (int u = 0; u < cols; ++u)
        {
            double sum_real = 0.0;                     // 实部和
            double sum_imag = 0.0;                     // 虚部和
            for (int y = 0; y < rows; ++y)
            {
                const double* ptr = mat.ptr<double>(y);
                for (int x = 0; x < cols; ++x)   // z = r(cosθ + isinθ) => real = rcosθ, imag = rsinθ
                {
                    double theta = -2.0 * CV_PI * (1.0 * u * x / cols + 1.0 * v * y / rows);
                    sum_real += ptr[x] * cos(theta);  
                    sum_imag += ptr[x] * sin(theta);
                }
            }
            p_real[u] = sum_real;
            p_imag[u] = sum_imag;
        }
    }

    // 合并实部与虚部
    cv::Mat result, plane[] = { res_real, res_imag };
    cv::merge(plane, 2, result);
    if (result.empty())
        return cv::Mat();

    return result;
}

int main()
{
    cv::Mat mat(21, 7, CV_8UC1);
    cv::randu(mat, cv::Scalar::all(0), cv::Scalar::all(255));
    cout << "original matrix = \n" << mat << endl;

    cv::Mat paddedMat;
    int optimalRows = cv::getOptimalDFTSize(mat.rows);
    int optimalCols = cv::getOptimalDFTSize(mat.cols);
    cout << "optimal rows = " << optimalRows << endl; 
    cout << "optimal cols = " << optimalCols << endl;

    cv::copyMakeBorder(mat, paddedMat, 0, optimalRows - mat.rows, 0, optimalCols - mat.cols, cv::BORDER_CONSTANT, cv::Scalar::all(0));

    cv::Mat paddedMatDouble;
    paddedMat.convertTo(paddedMatDouble, CV_64FC1);
    cout << "padded matrix in double = \n" << paddedMatDouble << endl;

    cv::Mat dft_res;
    cv::dft(paddedMatDouble, dft_res, cv::DFT_COMPLEX_OUTPUT);
    if (dft_res.empty())
        return -1;

    cv::Mat dft_real_imag[] = { cv::Mat::zeros(dft_res.size(), CV_64FC1), cv::Mat::zeros(dft_res.size(), CV_64FC1) };
    cv::split(dft_res, dft_real_imag);
    cout << "dft real part = \n" << dft_real_imag[0] << endl;
    cout << "dft imag part = \n" << dft_real_imag[1] << endl;

    cv::Mat my_dft_res = myDft(paddedMatDouble);
    if (my_dft_res.empty())
        return -1;

    cv::Mat my_dft_real_imag[] = { cv::Mat::zeros(my_dft_res.size(), CV_64FC1), cv::Mat::zeros(my_dft_res.size(), CV_64FC1) };
    cv::split(my_dft_res, my_dft_real_imag);
    cout << "my dft real part = \n" << my_dft_real_imag[0] << endl;
    cout << "my dft imag part = \n" << my_dft_real_imag[1] << endl;

    cv::Mat diff_real_double = my_dft_real_imag[0] - dft_real_imag[0];
    cv::Mat diff_imag_double = my_dft_real_imag[1] - dft_real_imag[1];

    cv::Mat diff_real_uchar, diff_imag_uchar; // 为了使输出不受精度影响, 转化成uchar
    diff_real_double.convertTo(diff_real_uchar, CV_8UC1);
    diff_imag_double.convertTo(diff_imag_uchar, CV_8UC1);
    cout << "difference real part = \n" << diff_real_uchar << endl;
    cout << "difference imag part = \n" << diff_imag_uchar << endl;

    return 0;
}