#include <opencv2\opencv.hpp>  
#include <opencv2\core\core.hpp>
#include <opencv2\core\mat.hpp>
#include <iostream>  
#include "PoissonDCT.h"
#include "FFT_OT.h"

/******************************************************************************
* 
*   Supplementary Demo Code for the ICCV submission Paper_id 7188
* 
*   FFT-OT: A Fast Algorithm for Optimal Transportation
*   
*   This demo shows the implementation details of the two algorithms in the submission
*   based on the OpenCV built-in DCT, IDCT functions
* 
*   1. DCT Poisson Equation Solver
*   2. FFT OT
*   
*   This is the simplified version the algorithms, the full version is based on
*   libfftw library.
* 
******************************************************************************/


using namespace std;
using namespace cv;

/*
 * Test the FFT_OT algorithm
 * the input is the image name
 * 1. the source density is set by the method set_source_density_Gaussian()
 *    or set_source_density_Bezier()
 * 2. the target gensity is set to be 1
 * 3. the algorithm solves the Monge-Ampere equation
 * 4. the OT mapping is visualized by forward warpping the input image
 * 
 */
void FFT_OT_test(const char * image_name, const int n = 512)
{
    std::cout << "Testing the FFT OT Algorithm" << std::endl;

    Mat src = imread(image_name, 0);
    if (src.empty())
    {
        cout << "the image is not exist" << endl;
        return;
    }
    resize(src, src, Size(n, n));
    src.convertTo(src, CV_64F, 1.0 / 255);
    Mat tgt(n, n, CV_64F, Scalar(0));

    //solving the OT map using FFT
    CFFT_OT FO(n, n);
    FO.solve();
    //forward mapping
    FO.forward_map(FO.gradient(), src, tgt);
    //visualize the source intensity
    FO.Poisson_solver().visualize("Source Density", FO.source_density());
    //visualize the OT map
    imshow("OT Map", tgt);
    std::cout << "Press any key to continue " << std::endl;
    waitKey();
}

/*
 * Test DCT_Poisson algorithm
 * the input is the image name
 *
 *  1. compute the Laplacian of the input image
 *  2. use DCT to solve the Poisson equation
 *  3. display the input image, the solution to the Poisson equation
 *     the two images should look similar
 */

void Poisson_test(const char* image_name)
{
    std::cout << "Testing the DCT Poisson Solver,";
    std::cout << "compare the source imageand the solution image" << std::endl;

    Mat src = imread(image_name, 0);
    if (src.empty())
    {
        cout << "the image is not exist" << endl;
        return;
    }
    resize(src, src, Size(512, 512));
    src.convertTo(src, CV_64F, 1.0 / 255);

    //compute the DCT of the input image
    Mat srcDCT;
    dct(src, srcDCT);
    //display the source image
    imshow("Source image", src);
    //display the DCT of the source
    imshow("DCT of the source", srcDCT);
    std::cout << "Press any key to continue " << std::endl;
    waitKey();

    //solve the Poisson equation using FFT
    CPoissonDCT PD(512, 512);
    PD.test(src);
}

/*
 *  The input argument is the image name
 *  such as "data/sophie.png"
 * 
 */
int main(int argc, char * argv[])
{
    if (argc < 1){
        std::cout << "Usage: " << argv[0] << " input image name" << std::endl;
    }

    //test the DCT Poisson Equation solver algorithm
    Poisson_test(argv[1]);
    //test the FFT-OT algorithm
    FFT_OT_test(argv[1],512);
       
    return 0;
}