enum Filter{lowpass, highpass, notch, bandpass, lpn, hpbp};

void utilities::swapQuads(cv::Mat &img) {
    // Rearranges amplitude image
    int cx = img.cols / 2;
    int cy = img.rows / 2;
    cv::Mat q0(img, cv::Rect(0, 0, cx, cy));   // Top-Left - Create a ROI per quadrant
    cv::Mat q1(img, cv::Rect(cx, 0, cx, cy));  // Top-Right
    cv::Mat q2(img, cv::Rect(0, cy, cx, cy));  // Bottom-Left
    cv::Mat q3(img, cv::Rect(cx, cy, cx, cy)); // Bottom-Right

    // swap quadrants (Top-Left with Bottom-Right)
    cv::Mat tmp;
    q0.copyTo(tmp);
    q3.copyTo(q0);
    tmp.copyTo(q3);

    // swap quadrant (Top-Right with Bottom-Left)
    q1.copyTo(tmp);
    q2.copyTo(q1);
    tmp.copyTo(q2);
}

cv::Mat utilities::create_mask(Filter filter, int rows, int cols, int d0, int d1, int d2){
    cv::Mat mask(rows, cols, CV_32FC2);
    int cx = cols / 2, cy = rows / 2, di;
    for(int i = 0; i < rows; i++){
        for(int j = 0; j < cols; j++){
            di = sqrt(pow(i-cy, 2) + pow(j-cx, 2));
            switch(filter){
                case lowpass :
                    if(di <= d0){
                        mask.at<cv::Vec2f>(i, j)[0] = 1;
                        mask.at<cv::Vec2f>(i, j)[1] = 1;
                    }
                    else{
                        mask.at<cv::Vec2f>(i, j)[0] = 0;
                        mask.at<cv::Vec2f>(i, j)[1] = 0;
                    }
                    break;

                case highpass:
                    if(di >= d0){
                        mask.at<cv::Vec2f>(i, j)[0] = 1;
                        mask.at<cv::Vec2f>(i, j)[1] = 1;
                    }
                    else{
                        mask.at<cv::Vec2f>(i, j)[0] = 0;
                        mask.at<cv::Vec2f>(i, j)[1] = 0;
                    }
                    break;

                case notch:
                    if(di <= d1 || di >= d2){
                        mask.at<cv::Vec2f>(i, j)[0] = 1;
                        mask.at<cv::Vec2f>(i, j)[1] = 1;
                    }
                    else{
                        mask.at<cv::Vec2f>(i, j)[0] = 0;
                        mask.at<cv::Vec2f>(i, j)[1] = 0;
                    }
                    break;

                case bandpass:
                    if(di >= d1 && di <= d2){
                        mask.at<cv::Vec2f>(i, j)[0] = 1;
                        mask.at<cv::Vec2f>(i, j)[1] = 1;
                    }
                    else{
                        mask.at<cv::Vec2f>(i, j)[0] = 0;
                        mask.at<cv::Vec2f>(i, j)[1] = 0;
                    }
                    break;

                case lpn:
                    if(di <= d0 || (di >= d1 && di <= d2)){
                        mask.at<cv::Vec2f>(i, j)[0] = 1;
                        mask.at<cv::Vec2f>(i, j)[1] = 1;
                    }
                    else{
                        mask.at<cv::Vec2f>(i, j)[0] = 0;
                        mask.at<cv::Vec2f>(i, j)[1] = 0;
                    }
                    break;

                case hpbp:
                    if(di >= d0 && (di <= d1 || di >= d2)){
                        mask.at<cv::Vec2f>(i, j)[0] = 1;
                        mask.at<cv::Vec2f>(i, j)[1] = 1;
                    }
                    else{
                        mask.at<cv::Vec2f>(i, j)[0] = 0;
                        mask.at<cv::Vec2f>(i, j)[1] = 0;
                    }
                    break;

                default : break;
            }

        }
    }

    return mask;
}

void utilities::cvidft(cv::Mat &tgt, cv::Mat &magI, cv::Mat &complexI){
    // Creates fourier domain amplitude image
    double minv, maxv;
    cv::minMaxLoc(tgt, &minv, &maxv);
    cv::Mat planes[2];
    cv::split(complexI, planes);
    cv::magnitude(planes[0], planes[1], magI);
    magI += cv::Scalar::all(1);
    cv::log(magI, magI);
    cv::normalize(magI, magI, minv, maxv, CV_MINMAX);

    // Converts back to image, crops padding, and saves roi to target
    utilities::swapQuads(complexI);
    cv::idft(complexI, complexI, cv::DFT_SCALE | cv::DFT_REAL_OUTPUT);
    cv::normalize(complexI, complexI, 0, 255, CV_MINMAX);
    complexI.convertTo(complexI, CV_8U);
    complexI = complexI(cv::Rect(0, 0, tgt.cols, tgt.rows));
    complexI.copyTo(tgt);
}

void utilities::cvdft(cv::Mat &src, cv::Mat &magI, cv::Mat &complexI){
    // Gets min/max
    double minv, maxv;
    cv::minMaxLoc(src, &minv, &maxv);

    // Pads image if needed
    cv::Mat padded;
    int m = cv::getOptimalDFTSize(src.rows);
    int n = cv::getOptimalDFTSize(src.cols);
    cv::copyMakeBorder(src, padded, 0, m - src.rows, 0, n - src.cols, cv::BORDER_CONSTANT, cv::Scalar::all(0));

    // Create planes and run dft
    cv::Mat planes[] = {cv::Mat_<float>(padded), cv::Mat::zeros(padded.size(), CV_32F)};
    cv::merge(planes, 2, complexI);
    cv::dft(complexI, complexI);

    // Swaps quadrants
    utilities::swapQuads(complexI);

    // Creates fourier domain amplitude image
    cv::split(complexI, planes);
    cv::magnitude(planes[0], planes[1], magI);
    magI += cv::Scalar::all(1);
    cv::log(magI, magI);
    cv::normalize(magI, magI, minv, maxv, CV_MINMAX);
}

void utilities::dft_filter(cv::Mat &tgt, cv::Mat &magbefore, cv::Mat &magafter, int d0, int d1, int d2, Filter filter){
    // Runs dft on src
    cv::Mat complexI;
    utilities::cvdft(tgt, magbefore, complexI);

    // Runs low pass filter mask
    cv::Mat mask = utilities::create_mask(filter, complexI.rows, complexI.cols, d0, d1, d2);

    // Applies mask to complexI
    cv::multiply(complexI, mask, complexI);

    // Inverse DFT
    utilities::cvidft(tgt, magafter, complexI);
}