#ifndef _DCT_POISSON_H_
#define _DCT_POISSON_H_

#include <opencv2\opencv.hpp>  
#include <opencv2\core\core.hpp>
#include <opencv2\core\mat.hpp>
#include <iostream>  
#include "FiniteDifferenceOperator.h"
#include <string>
using namespace std;
using namespace cv;

/*
 *  Solve Poisson Equation using DCT
 *  L u = f
 */
class CPoissonDCT : public CFiniteDifferenceOperator {

public:
    /*
    *   Constructor
    *   - width, height : image width and height
    */
	CPoissonDCT(int width, int height):
        m_width(width),
        m_height(height),
		//Laplace( u ) = f
        m_u(width,height,  CV_64F, Scalar(0)), 
		m_f(width, height, CV_64F, Scalar(0)),
        //matrix with ghost cells
        m_U(width+2, height+2, CV_64F, Scalar(0)),

        //intermediate differential values
        m_Dxx(width, height, CV_64F, Scalar(0)),
        m_Dyy(width, height, CV_64F, Scalar(0)),
        m_Dxy(width, height, CV_64F, Scalar(0))
    {
        //the domain is the square [-1,1] X [-1, 1]
        //spacial step length
        m_hx = 2.0 / m_width;
        m_hy = 2.0 / m_height;

        //constant coefficients, used for solving the equation
        m_wx.resize(width);
        m_wy.resize(height);

        for (int i = 0; i < width; i++) {
            m_wx[i] = 2 * (cos(M_PI * i / width) - 1) / (m_hx * m_hx);
        }
        for (int j = 0; j < height; j++) {
            m_wy[j] = 2 * (cos(M_PI * j / width) - 1) / (m_hy * m_hy);
        }

    }

    /*
     *  Destructor
     */
	~CPoissonDCT() {};

    /*
     *  Solve the Poisson equation using DCT
     *  Laplace u = f
     *  hx, hy - spacial step lengths
     * 
     */
    void solve( cv::Mat & u, cv::Mat & f, double hx, double hy)
    {
        int width  = u.cols;
        int height = u.rows;

        cv::Mat fDCT;

        // DCT
        cv::dct(f, fDCT);

        for( int i = 0; i < width; i ++ )
            for (int j = 0; j < height; j++) {
                fDCT.at<double>(i, j) /= (m_wx[i] + m_wy[j]);
            }
        //normalize the solution
        fDCT.at<double>(0, 0) = 0;
        // Inverse DCT
        cv::idct(fDCT, u);
    }

    /*
    *   Compute the Laplacian   f = Laplace(u)
    *   u - input function
    *   U - extension of u with ghost cells
    *   f - output function, the Laplacian of u
    *   Dxx - u_{xx}
    *   Dyy - u_{yy}
    *   hx  - horizontal spacial step length
    *   hy  - vertical   spacial step length
    */
    void Laplace(cv::Mat& u, cv::Mat& U, cv::Mat& f, cv::Mat & Dxx, cv::Mat & Dyy, double hx, double hy) {

        int width  = u.cols;
        int height = u.rows;

        //set the ghost cells
        set_ghost_cells(U, u);
        
        //compute the second derivatives
        _dxx(U, Dxx, hx);
        _dyy(U, Dyy, hy);

        //Laplacian
        f = Dxx + Dyy;

        //visualize("Laplacian", f);
        cv::imshow("Laplacian", f);
        std::cout << "Press any key to continue " << std::endl;
        cv::waitKey();
    }

    //test the Poisson solver
    //input - the input image
    void test(cv::Mat & input)
    {
        m_u = input;
        //compute the Laplacian of the input image
        Laplace(m_u, m_U, m_f, m_Dxx, m_Dyy, m_hx, m_hy);
        //solve the Poisson equation to recover u
        solve(m_u, m_f, m_hx, m_hy);
        //visualize the solution to the Poisson equation
        visualize("Solution to the Poisson Equation", m_u);
    }

protected:

    //the width of the image
    int m_width;
    //the height of the image
    int m_height;

    //the spacial step length
    double m_hx;
    //the spacial step length
    double m_hy;

    //the unknown function 
	cv::Mat m_u;
    //the Laplacian of the unknown function
	cv::Mat m_f;
    //the extension of the unknown function
    //with the ghost cells
	cv::Mat m_U; 

    //the coefficients in the frequency 
    //domain
    std::vector<double> m_wx, m_wy;

    //the derivatives of the unknown function
    cv::Mat m_Dxx;
    cv::Mat m_Dyy;
    cv::Mat m_Dxy;
};

#endif
