/*****************************************************************************
 * ass3horn.cpp
 ***************************************************************************** 
 *
 * Description: Assignment 3 (COMPSCI775)
 *
 * Topic: Calculation of Optical Flow - Horn-Schunck Algorithm
 *
 * (c) 2002 Christian Graf, Uli Schroeder, YongTao Zou (Group 7)
 *
 *****************************************************************************
 */

#include <math.h>
#include <iostream.h>
#include "HalconCpp.h"
#include "String.h"
#include <stdio.h>

class HORNSCHUNCK {

  private:

	// Simple Forward Difference  
	double computeEx_forwardSimple(HByteImage *firstImage, HByteImage *secondImage, int i, int j) {
		double result = 0;
		result = (*firstImage)(i,j) - (*firstImage)(i+1,j);
		result /= (double)4.0;
		return result;
	}
	double computeEy_forwardSimple(HByteImage *firstImage, HByteImage *secondImage, int i, int j) {
		double result = 0;
		result = (*firstImage)(i,j) - (*firstImage)(i,j+1);
		result /= (double)4.0;
		return result;
	}
	double computeEt_forwardSimple(HByteImage *firstImage, HByteImage *secondImage, int i, int j) {
		double result = 0;
		result = (*firstImage)(i,j) - (*secondImage)(i,j);
		result /= (double)4.0;
		return result;
	}

	// Advanced Forward Differences
	double computeEx_forwardAdv(HByteImage *firstImage, HByteImage *secondImage, int i, int j) {
		double result = 0;
		result = (*firstImage)(i+1,j) + (*secondImage)(i+1,j) + (*firstImage)(i+1,j+1) + (*secondImage)(i+1,j+1) -
		         (*firstImage)(i,j) - (*secondImage)(i,j) - (*firstImage)(i,j+1) - (*secondImage)(i,j+1);
		result /= (double)4.0;
		return result;
	}
	double computeEy_forwardAdv(HByteImage *firstImage, HByteImage *secondImage, int i, int j) {
		double result = 0;
		result = (*firstImage)(i,j+1) + (*secondImage)(i,j+1) + (*firstImage)(i+1,j+1) + (*secondImage)(i+1,j+1) -
		         (*firstImage)(i,j) - (*secondImage)(i,j) - (*firstImage)(i+1,j) - (*secondImage)(i+1,j);
		result /= (double)4.0;
		return result;
	}
	double computeEt_forwardAdv(HByteImage *firstImage, HByteImage *secondImage, int i, int j) {
		double result = 0;
		result = (*secondImage)(i,j) + (*secondImage)(i,j+1) + (*secondImage)(i+1,j) + (*secondImage)(i+1,j+1) -
		         (*firstImage)(i,j) - (*firstImage)(i,j+1) - (*firstImage)(i+1,j) - (*firstImage)(i+1,j+1);
		result /= (double)4.0;
		return result;
	}


  public:

	HORNSCHUNCK(HByteImage *firstImage, HByteImage *secondImage, HByteImage *outputImage, char* resultname, int mode,
		int iterations, float lambda, float initU, float initV, char* temp) {
		// Initialise Variables
		int imageWidth = (*firstImage).Width();
		int imageHeight = (*firstImage).Height();
		int factor, offset, gridwidth, extent;
		int n=0, n0=0;
		double udash=0.0, vdash=0.0, alpha=0.0, maxU=0.0, maxV=0.0;
		double errorU=0.0, errorV = 0.0;
        double trash=0.0;
		FILE * pFile; 

		// Allocate memory for two-dimensional array and store first & second image in it
		int **firstInBuffer = NULL;
		int **secondInBuffer = NULL;
		
		firstInBuffer = (int **) new int*[imageHeight];
		secondInBuffer = (int **) new int*[imageHeight];
		
		for (int i = 0; i < imageHeight; i++) {
			firstInBuffer[i] = (int *) new int[imageWidth];
			secondInBuffer[i] = (int *) new int[imageWidth];
		}
		for (int y = 0; y < imageHeight; y++) {
			for (int x = 0; x < imageWidth; x++) {
				firstInBuffer[x][y] = (*firstImage)(x, y);
				secondInBuffer[x][y] = (*secondImage)(x, y);
			}
		}		
		
		// Allocate memory
		double **myEx = NULL;
		double **myEy = NULL;
		double **myEt = NULL;
		double ***myU = NULL;
		double ***myV = NULL;
	
		myEx = (double **) new double*[imageHeight];
		myEy = (double **) new double*[imageHeight];
		myEt = (double **) new double*[imageHeight];
		myU = (double ***) new double*[imageHeight];
		myV = (double ***) new double*[imageHeight];

		for (int i = 0; i < imageHeight; i++) {
		myEx[i] = (double *) new double[imageWidth];
		myEy[i] = (double *) new double[imageWidth];
		myEt[i] = (double *) new double[imageWidth];
	    myU[i] = (double **) new double[imageWidth];
		myV[i] = (double **) new double[imageWidth][2];
		
			for (int j = 0; j < imageWidth; j++) {
				myU[i][j] = (double *) new double[2];
				myV[i][j] = (double *) new double[2];
			}
		}
	
		cout << "\tMemory for tempory arrays allocated." << endl;
		
	
		// Original Horn-Schunck algorithm
		
		cout << "\tHorn-Schunck is to be run..." << endl;
		if (mode==1) {
 			for (int j = 1; j < imageHeight-1; j++) {
	    			for (int i = 1; i < imageWidth-1; i++) {
				  	myEx[i][j] = (int)floor(computeEx_forwardSimple(firstImage, secondImage, i,j) + 0.5);
					myEy[i][j] = (int)floor(computeEy_forwardSimple(firstImage, secondImage, i,j) + 0.5);
					myEt[i][j] = (int)floor(computeEt_forwardSimple(firstImage, secondImage, i,j) + 0.5);
					myU[i][j][0] = (double)initU;
					myV[i][j][0] = (double)initV;
    				}
			}	
 		}
		if (mode==2) {
 			for (int j = 1; j < imageHeight-1; j++) {
	    			for (int i = 1; i < imageWidth-1; i++) {
					myEx[i][j] = (int)floor(computeEx_forwardAdv(firstImage, secondImage, i,j) + 0.5);
					myEy[i][j] = (int)floor(computeEy_forwardAdv(firstImage, secondImage, i,j) + 0.5);
					myEt[i][j] = (int)floor(computeEt_forwardAdv(firstImage, secondImage, i,j) + 0.5);
					myU[i][j][0] = (double)initU;
					myV[i][j][0] = (double)initV;
				}
			}
 		} 
 		
 		pFile = fopen ("results.txt","at");
  		if (pFile!=NULL)
  		{
  		    fprintf(pFile, "BEGIN file %s.\n", temp);
    		fprintf(pFile, "Modus:%d \titerations: %i \tLambda: %2.2f \tinitU: %2.2f \tinitV: %2.2f\n", mode, iterations, lambda, initU, initV);
    	} else printf("\tFAILURE: output to text file results.txt aborted.");
    	 		
 		n = 1;
		while (n <= iterations) {
	    	for (int j = 1; j < imageHeight-1; j++) {
		    	for (int i = 1; i < imageWidth-1; i++) {
		    		udash = (myU[i-1][j][0]+myU[i+1][j][0]+myU[i][j-1][0]+myU[i][j+1][0])/4.0f;
		    		vdash = (myV[i-1][j][0]+myV[i+1][j][0]+myV[i][j-1][0]+myV[i][j+1][0])/4.0f;
		    		alpha = (((myEx[i][j]*udash) + (myEy[i][j]*vdash) + myEt[i][j])*lambda)/(1+(lambda*((myEx[i][j]*myEx[i][j])+(myEy[i][j]*myEy[i][j]))));
	    			
					trash =	myU[i][j][1];				
					myU[i][j][1] = udash - (alpha * myEx[i][j]);
	    			errorU += (myU[i][j][1]-trash)*(myU[i][j][1]-trash);

					trash = myV[i][j][1];
					myV[i][j][1] = vdash - (alpha * myEy[i][j]);
					errorV += (myV[i][j][1]-trash)*(myV[i][j][1]-trash);
				}
		    }
	    	for (int j = 0; j < imageHeight; j++) {
		    	for (int i = 0; i < imageWidth; i++) {
	    			myU[i][j][0] = myU[i][j][1];
	    			myV[i][j][0] = myV[i][j][1];
		    	}
		    }
		    errorU /= ((imageHeight-2)*(imageWidth-2));
			errorV /= ((imageHeight-2)*(imageWidth-2));
			if (pFile!=NULL)
  			{
  				fprintf(pFile, "\titeration:\t%d \tLSE_u:\t%1.7f \tLSE_v:\t%1.7f\n", n, errorU, errorV);
    		}
			errorU=0.0;
			errorV=0.0;
			n++;
		}
		if (pFile!=NULL)
  		{
			fprintf(pFile, "FINISCH file %s.\n\n", temp);
    		fclose (pFile);
		}
		cout << "\t... gradient and error were computed." << endl;

		// Create needle map
		offset=20; 					// specify some space for border pixel and specify enlargement factor for image
		gridwidth = 5; extent = 3;	// specify grid width and to what extent the gradients should be enlarged
		if (temp=="rubic" || temp=="trees") factor = 1;  else factor = 2; 

		HWindow outputWindow(0,0, imageWidth*factor+offset, imageHeight*factor+offset);
		outputWindow.SetPart(0,0, imageWidth*factor+offset, imageHeight*factor+offset);
		sleep(1);					// give some time for the X-Window to appear
		// as the algorithm runs only from 1 to imagewidth-1, 1 to imageHeight-1 the display is set so as well
		for (int y=1; y<imageHeight-1; y+=gridwidth) {
			for (int x=1; x<imageWidth-1; x+=gridwidth) {
				outputWindow.DispCircle((double)(y*factor+offset/2), (double)(x*factor+offset/2), (double)1.0);
				outputWindow.DispLine((double)(y*factor+offset/2), (double)(x*factor+offset/2), 
				        (double)((y+myV[x][y][0]*extent)*factor+offset/2), (double)((x+myU[x][y][0]*extent)*factor+offset/2));
			}
		}
		cout << "\tSample points were chosen und gradients enlarged and displayed." << endl;

		// Free memory allocated
		for (int i = 0; i < imageHeight; i++) {
			for (int j = 0; j < imageWidth; j++) {
				delete [] myV[i][j];
				delete [] myU[i][j];
			}
			delete [] myV[i];
			delete [] myU[i];
			delete [] myEt[i];
			delete [] myEy[i];
			delete [] myEx[i];
			delete [] secondInBuffer[i];
			delete [] firstInBuffer[i];
		
		}
		delete [] myV;
		delete [] myU;
		delete [] myEt;
		delete [] myEy;
		delete [] myEx;
		delete [] secondInBuffer;
		delete [] firstInBuffer;
		cout << "\tMemory for tempory arrays freed." << endl;

		cout << "\tImage of gradients at sample points is displayed."<<endl;
		outputWindow.DumpWindow("jpeg", resultname);
		//outputWindow.Click();
		   	
	}
	
};

void MyHalconExceptionHandler(const HException& except)
{
	throw except;
}
	

int main(int argc, char* argv[]) {
	// Show informational header
	cout << "----------------------------------------------" << endl;
	cout << " Assignment 3 (COMPSCI775):                   " << endl;
	cout << " Calculation of Optical Flow                  " << endl;
	cout << " (Horn-Schunck Algorithm)                     " << endl;
	cout << "----------------------------------------------" << endl;
	cout << " Group 7: Graf, Schroeder, Zou		       " << endl;
	cout << "----------------------------------------------" << endl;
	// Check for correct number of command line arguments and show help if neccessary
	if (argc != 8) {
		cout << " Wrong number of arguments!" << endl << endl;
		cout << " ./ass3horn [InputFile] [OutputFile] [mode] [iterations] [lambda] [u0] [v0]" << endl << endl;
		cout << "    [inputfile]   path and filename of the image (1.bmp or 2.bmp is added automatically!)" << endl;
		cout << "    [outputfile]  path and filename of the output image (.jpg added automatically)" << endl;
		cout << "    [mode]        1 for simple forward differences, 2 for advanced ones" << endl;
		cout << "    [iterations]  int number of iterations" <<endl;
		cout << "    [lambda]      float value of lambda" << endl;
		cout << "    [u0]		   initial float value for u" << endl;
		cout << "    [v0]		   initial float value for v" << endl;
		cout << " Example:" << endl;
		cout << "   ./ass3horn taxi result 2 8 10 0 0" << endl << endl;
		return(-1);
	}

	// Test if the values of the parameters are valid 
	if (!((atoi(argv[3]))==1)) {
	  if (!((atoi(argv[3]))==2)) {
		cout << " STOP: No valid mode selection: '1' for simple differences" << endl <<
			"\t'2' for advanced ones " << endl;
		return (-1);
	  }
	}
	if ((!((atoi(argv[4]))>0)) || (!((atof(argv[5]))>0))) {
		cout << " STOP: The value for [iterations] and [lambda] " << endl;
		cout << "\tmust not be equal to or less than zero!" << endl;
		return (-1);
	} 
	
	// Initialise variables
	String ext = ".bmp";
	HByteImage firstImage, secondImage, result;
	Herror error_num; char message[100];
	HException::InstallHHandler(&MyHalconExceptionHandler);
	
	// Test if input images are existing and read them
	(void)::set_check("give_error");
	try {
		firstImage = HByteImage(((String)argv[1])+"1"+ext); 
	} catch (HException &except) {
		error_num = except.err;
		if (error_num == H_MSG_FAIL) { cout << " STOP: Image " << ((String)argv[1])+"1"+ext << " not found!" << endl; }
		  else { 
			(void)::get_error_text(error_num, message);
			cout << " HALCON error occured: "<< message << endl;
			cout << " FAILURE! Program terminated." << endl << endl;
		}
		return (-1);	
	}
	
	try {
		secondImage = HByteImage(((String)argv[1])+"2"+ext);
	} catch (HException &except) {
		error_num = except.err;
		if (error_num == H_MSG_FAIL) { cout << " STOP: Image " << ((String)argv[1])+"2"+ext << " not found!" << endl; }
		  else {
		    (void)::get_error_text(error_num, message);
			cout << " HALCON error occured. " << message << endl;
			cout << " FAILURE! Program terminated." << endl << endl;
		}
		return (-1);	
	}

	
	// Test if input images have the same size
	if (!(firstImage.Width()==secondImage.Width() && firstImage.Height()==secondImage.Height())) {
		cout << "\tSTOP: Both input images have to be of same size!" << endl; 
		return (-1);
	}

	// Format result picture
	result = HByteImage(firstImage.Width(), firstImage.Height());

	// Apply Horn-Schunck algorithm
	cout << " Processing images " << (String)(argv[1])+"1"+ext << " & " << (String)(argv[1])+"2"+ext << "..." << endl;
	new HORNSCHUNCK(&firstImage, &secondImage, &result, argv[2], atoi(argv[3]), atoi(argv[4]), atof(argv[5]), atof(argv[6]),
	                atof(argv[7]), argv[1]);

	cout << " Processing finished..." << endl;



	// End program
	cout << " The End!" << endl << endl;
	return(0);
}