// =================
// || Worksheet 10
// || In this exercise, we are going be solving the heat equation. d/dt f(x,t) = - k d^2/dx^2 f(x,t)
// =================

#include "program.h"

using namespace std;
using namespace cv;

int worksheetten(Mat& src)
{
    Mat input = Mat::zeros(200,200, CV_32F );
    
    // a window to show our outputs
    cv::namedWindow( "current frame", WINDOW_FREERATIO );
    
    int MAX_NUMBER_OF_FRAMES = 5000;
    
    // this is the coefficient of dispersion in our heat equation.
    float k = 0.2f;
    
    // draw a little circle in the middle of our image.
    circle(input, Point(input.rows/2, input.cols/2), 1.0f, Scalar( 1.0, 1.0, 1.0 ),CV_FILLED);
    
    // This is our initial data, i.e. f(x,y,t=-1)
    // remember, when solving first order differential equations,
    // we need to specify initial conditions.
    // for now, we'll just make it a copy of our input.
    Mat prev_frame = input.clone() ; // f(x,y,t)
    Mat prev_prev_frame = input.clone(); // f(x,y,t-1)
    
    Mat colorized_frame;
    
    Size imsize = prev_frame.size();
    // repeat this for each of the frames.
    for (int n = 0; n < MAX_NUMBER_OF_FRAMES; n++)
    {
        Mat frame = prev_frame.clone();
        
        for(int i = 1; i < imsize.height -1; i++ )
        {
            for(int j = 1; j < imsize.width -1; j++ )
            {
                // 1d Heat equation: d/dt f(x,t) = - k d^2/dx^2 f(x,t)
                // when discretized, this becomes
                // -f(x,y,t) + f(x,y,t+1) = -k ( f(x-1,y,t)-2*f(x,y,t)+f(x+1,y,t) )
                
                //frame.at<float>(i,j) += k*(prev_frame.at<float>(i+1,j)
                //                           +prev_frame.at<float>(i-1,j)
                //                           -2*prev_frame.at<float>(i,j));
                
                // 2d Heat equation: d/dt f(x,t) = - k (d^2/dx^2  + d^2/dy^2 ) f(x,y,t)
                // let L = (d^2/dx^2  + d^2/dy^2 ), this is called the laplacian operator.
                // when discretized, what does this become?
                // -f(x,y,t) + f(x,y,t+1) = -k L f(x,y,t)
                
                //frame.at<float>(i,j) += k*(prev_frame.at<float>(i+1,j)
                //                           +prev_frame.at<float>(i-1,j)
                //                           -4*prev_frame.at<float>(i,j)
                 //                          +prev_frame.at<float>(i,j+1)
                //                           +prev_frame.at<float>(i,j-1)
                //                           );
                
 // f(x,y,t+1) - 2*f(x,y,t) + f(x,y,t-1) = - k L f(x,y,t)
                // prev_image : f(x,y,t)
                // prev_prev_image : f(x,y,t-1)
                
                frame.at<float>(i,j) = 2 * prev_frame.at<float>(i,j)
                                        - prev_prev_frame.at<float>(i,j)
                                     +    k*(prev_frame.at<float>(i+1,j)
                                           +prev_frame.at<float>(i-1,j)
                                           -4*prev_frame.at<float>(i,j)
                                           +prev_frame.at<float>(i,j+1)
                                           +prev_frame.at<float>(i,j-1)
                                           );
                
            }
        }
        
        // lets colorize our image too, so that it looks better.
        frame.convertTo(colorized_frame,CV_8UC1, 127,127);
        applyColorMap(colorized_frame, colorized_frame,COLORMAP_AUTUMN);
        // lets also make it 4 times the size, so we dont strain our eyes.
        resize(colorized_frame,colorized_frame,Size(colorized_frame.cols*4,colorized_frame.rows*4 ) );
        
        // display the colorized frame.
        cv::imshow( "current frame", colorized_frame);

        // wait for a key press to continue;
        //cv::waitKey();
        
        //wait for 'esc' key press for 15ms. If 'esc' key is pressed, break loop
        if (waitKey(50) == 27)
        {
            cout << "esc key is pressed by user" << endl;
            break;
        }
        
        // now, importantly, we set the previous frame for the next iteration to be the current frame.
        
        prev_prev_frame = prev_frame.clone();
        prev_frame = frame.clone();
    }
    
    return 0;
}