/*! 
*   \file  Fourier.cpp
*   \brief Defines basic functions in Fourier domain.
*    
*   \details  
*   TODO.
*
*   \date      March 22, 2012
*   \copyright eBay Research Labs.
*/

// Include standard C++ headers.  
//#include <algorithm>

// Include openCV-specific headers.
#include <highgui.h>
#include <opencv2/contrib/contrib.hpp> // For Timer.

#include "Fourier.h"

// Namespaces.
using namespace std;
using namespace cv;

namespace Fourier 
{

// **** Helper functions ****

/*!
*   \brief   Evaluates Gaussian function in the desired rectangle.
*   \details Gaussian is isotropic and is centered at the origin. 
*            Rectangle is defined by [x1, x2) x [y1, y2). Y increases downwards.
*            Values will be in [0, 1].
*   \param [in]  x1   Left-most x value.
*   \param [in]  x2   Right-most x value.
*   \param [in]  y1   Top-most y value.
*   \param [in]  y1   Bottom-most y value.
*   \param [in]  den  Denominator of Gaussian exponent.
*   \param [out] img  Gaussian subset image of type CV_32F.
*/
void GaussianInRectangle(int x1, int x2, int y1, int y2, float den, Mat &img)
{
    MatIterator_<float> it = img.begin<float>();
    for (int i=y1; i<y2; i++)
    {
        for (int j=x1; j<x2; j++)
        {
            *it++ = exp(-((float) i * (float) i + (float) j * (float) j) / den);
        }
    }
}


// *** Functions declared in Fourier.h "***

/*!
*   \brief   Swaps diagonal quadrants as fftshift in Matlab does.
*   \details Quadrants 1 and 3 are swapped. Quadrants 2 and 4 are swapped.
*            The input image must have even number of rows and columns.
*   \param [in, out] img  Image that needs its quadrants swapped.
*/
void SwapDiagonalQuadrants(Mat &img)
{
    // Get half dimensions.
    int nx = img.cols/2;
    int ny = img.rows/2;

    CV_Assert( 2*nx==img.cols && 2*ny==img.rows); // Expect #rows and #columns to be even.

    // Define quadrants.
    Rect qTopLeft     (0,   0, nx, ny);
    Rect qTopRight    (nx,  0, nx, ny);
    Rect qBottomLeft  (0,  ny, nx, ny);
    Rect qBottomRight (nx, ny, nx, ny);

    // Copy image.
    Mat tmp = img.clone();      

    // Swap Top-Left and Bottom-Right quadrants.
    tmp(qTopLeft).copyTo(img(qBottomRight));
    tmp(qBottomRight).copyTo(img(qTopLeft));
 
     // Swap Top-Right and Bottom-Left quadrants.
    tmp(qTopRight).copyTo(img(qBottomLeft));
    tmp(qBottomLeft).copyTo(img(qTopRight));
}


/*!
*   \brief  Gaussian function in Fourier domain, packed a 2-channel image.
*   \details This Gaussian function is ready to be used with dft. 
*            (i.e.) even and odd quadrants are swapped already. 
*   \param [in]  n             #samples along each positive frequency axis. [-n/2:1:n/2-1].
*   \param [in]  denOfExponent Denominator of Gaussian exponent.
*   \return  Gaussian. Type = CV_32F. channel[0] = real part, channel[1] = imaginary part.
*/
Mat FourierOfGaussian(int n, float denOfExponent)
{
    int nc = n/2; // Get half-dimension.
    CV_Assert( 2*nc==n ); // Expect dimensions to be even.

    Mat gaussFourier = Mat::zeros(n, n, CV_32F);

    // Define quadrants.
    Rect qTopLeft     (0,   0, nc, nc);
    Rect qTopRight    (nc,  0, nc, nc);
    Rect qBottomLeft  (0,  nc, nc, nc);
    Rect qBottomRight (nc, nc, nc, nc);
   
   // Evaluate Gaussian in each quadrant.
    Mat topLeft     = gaussFourier(qTopLeft);
    Mat topRight    = gaussFourier(qTopRight);
    Mat bottomLeft  = gaussFourier(qBottomLeft);
    Mat bottomRight = gaussFourier(qBottomRight);

    GaussianInRectangle(0,  nc,  0, nc, denOfExponent, topLeft);
    GaussianInRectangle(-nc, 0,  0, nc, denOfExponent, topRight);
    GaussianInRectangle(0,  nc, -nc, 0, denOfExponent, bottomLeft);
    GaussianInRectangle(-nc, 0, -nc, 0, denOfExponent, bottomRight);

    Mat gfPlanes[] = {Mat_<float>(gaussFourier), Mat::zeros(gaussFourier.size(), CV_32F)};
    Mat gfComplex;
    merge(gfPlanes, 2, gfComplex); 
        
    return gfComplex;
}


/*!
*   \brief  Evaluates of magnitude of Fourier transform of given single channel image.
*   \details Applies fourier transform and then evaluates magnitude. 
*   \param [in]  img    Single-channel image.
*   \return  Magnitude of Fourier transform of img. Type is CV_32F.
*/
Mat FourierMagnitude(const Mat &img)
{
    // Copy image.
    Mat imgCopy  = img.clone();

    // Create 2-channel image with 0 values for imaginary part.
    Mat planes[] = {Mat_<float>(imgCopy), Mat::zeros(imgCopy.size(), CV_32F)};
    Mat imgFourier;
    merge(planes, 2, imgFourier);    

    // Apply Fourier transform.
    dft(imgFourier, imgFourier);

    // Separate real and imaginary parts.
    split(imgFourier, planes);

    // Evaluate magnitude of complex entries.
    magnitude(planes[0], planes[1], planes[0]);

    // For Debugging: Display log of magnitude of Fourier Transform, with dc at center of image.
    //Mat mag = planes[0].clone();
    //SwapDiagonalQuadrants(mag);
    //log(1+mag, mag);
    //normalize(mag, mag, 0, 1, CV_MINMAX); // Normalize image to have range in [0, 1]
    //imshow("Mag", mag);
    //waitKey(0);

    return planes[0];
}


/*!
*   \brief  Display magnitude of complex matrix comprising of 2 channels.
*   \details channel[0] = real part, channel[1] = imaginary part. 
*   \param [in]  complex  Complex matrix comprising of 2 channels.
*   \param [in]  figName  Title of figure.
*/
void ViewMagnitudeOfComplex(const Mat &complex, const string &figName)
{
    // Separate real and imaginary parts.
    vector <Mat> planes;
    split(complex, planes);                    // planes[0] = Re(DFT(I), planes[1] = Im(DFT(I))

    // Extract magnitude.
    magnitude(planes[0], planes[1], planes[0]); // planes[0] = magnitude
    Mat magI = planes[0];

    // Switch to logarithmic scale.
    magI += Scalar::all(1);                    
    log(magI, magI);

    // Normalize values to be in [0, 1] for better visualization.
    normalize(magI, magI, 0, 1, CV_MINMAX); 

    // Swap quadrants so that origin is at center of image.
    SwapDiagonalQuadrants(magI);

    // Dispay image.
    imshow(figName, magI);
    waitKey(0);
}


/*!
*   \brief  Returns magnitude of product of spectrums.
*   \details Fourier Transform of given image is modulated by the given Fourier
*            transform of filter. The magnitude of the result is returned. 
*   \param [in]  img            Input image. Type is CV_32F.
*   \param [in]  filterFourier  Fourier transform of filter. Type is CV_32FC2.
*   \param [in]  flagReturnMagnitude  When true, return value will be magnitude. Else,
*                real part will be returned.
*   \return  Either magnitude or real part of product of spectrums. Type is CV_32F.
*/
Mat ApplyFilter(const Mat &img, Mat &filterFourier, bool flagReturnMagnitude)
{
    CV_Assert( !img.data && img.type() == CV_32F &&  !filterFourier.data && filterFourier.type() == CV_32FC2);

    Mat imgCopy = img.clone();

    Mat planes[] = {Mat_<float>(imgCopy), Mat::zeros(imgCopy.size(), CV_32F)};
    Mat imgFourier;
    merge(planes, 2, imgFourier);    

    dft(imgFourier, imgFourier);
    mulSpectrums(imgFourier, filterFourier, imgFourier, 0);

    Mat lowPassComplex;
    idft(imgFourier, lowPassComplex, DFT_SCALE);

    split(lowPassComplex, planes); // planes[0] = Re(DFT(low pass)), planes[1] = Im(DFT(low pass))

    if (flagReturnMagnitude)
    {
        magnitude(planes[0], planes[1], planes[0]); // planes[0] = magnitude
    }

    return planes[0];
}


} // namespace Fourier