
#include <stdlib.h> 
#include <stdio.h> 
#include "vision_utilities.h"

typedef struct _needle{
	double x;
	double y;
} Needle;

int myround(double x);

int flow(char *fname_im1, char *fname_im2, int step, char *fname_output)
{
	const int w1 = 3;//half windows size of small window (computing)
	const int w2 = 20;//half windows size of large window (searching)


	Image *im1 = new Image;//input image 1
	Image *im2 = new Image;//input image 2
	Image *im_output = new Image;//output image
	readImage(im1, fname_im1);
	readImage(im2, fname_im2);
	int NRows = getNRows(im1);
	int NCols = getNCols(im1);
	setSize(im_output,NRows,NCols);
	setColors(im_output,255);

	Needle **needleMap = new Needle*[NRows];
	for (int i = 0; i < NRows; i++)
		needleMap[i] = new Needle[NCols];

	/* copy im1 to im_output */
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			setPixel(im_output,i,j,getPixel(im1,i,j));		
		}
	}

	/* choose the points */
	double maxLength = 0;//maximum length of optical flow vector
	for (int i1 = 0; i1 < NRows; i1+=step)
	{
		for (int j1 = 0; j1 < NCols; j1+=step)
		{
			if ((i1 < w1)||(i1 > NRows-w1)||(j1 < w1)||(j1 > NCols-w1)) continue;//ignore pixels near the edge

			double maxNCC = 0;
			needleMap[i1][j1].x = 0;
			needleMap[i1][j1].y = 0;

			/* draw a star for visibility */
			setPixel(im_output,i1+1,j1,0);	
			setPixel(im_output,i1-1,j1,0);	
			setPixel(im_output,i1,j1+1,0);
			setPixel(im_output,i1,j1-1,0);

			/* template matching */
			//center of window in im1: (i1,j1)
			//center of window in im2: (i2,j2)
			for (int i2 = MAX(i1-w2,w1); i2 <= MIN(i1+w2,NRows-1-w1); i2++)
			{
				for (int j2 = MAX(j1-w2,w1); j2 <= MIN(j1+w2,NCols-1-w1); j2++)
				{
					double sumProduct = 0;
					double sum1 = 0, sum2 = 0;

					for (int k = -w1; k <= w1; k++)
					{
						for (int l = -w1; l <= w1; l++)
						{
							sum1 += (getPixel(im1,i1+k,j1+l))*(getPixel(im1,i1+k,j1+l));
							sum2 += (getPixel(im2,i2+k,j2+l))*(getPixel(im2,i2+k,j2+l));
							sumProduct += getPixel(im1,i1+k,j1+l) * getPixel(im2,i2+k,j2+l);
						}
					}
					
					double NCC = sumProduct/sqrt(sum1*sum2);
					if (NCC > maxNCC)
					{
						maxNCC = NCC;
						needleMap[i1][j1].x = j2-j1;
						needleMap[i1][j1].y = i2-i1;
					}
				}
			}
			double length = sqrt(needleMap[i1][j1].x*needleMap[i1][j1].x + needleMap[i1][j1].y*needleMap[i1][j1].y);
			if (length > maxLength) {maxLength = length;}			
		}
	}

	/* normalize vectors and draw the needles */
	for (int i1 = 0; i1 < NRows; i1+=step)
	{
		for (int j1 = 0; j1 < NCols; j1+=step)
		{
			if ((i1 < w1)||(i1 > NRows-w1)||(j1 < w1)||(j1 > NCols-w1)) continue;//ignore pixels near the edge
			needleMap[i1][j1].x = needleMap[i1][j1].x/maxLength*step;
			needleMap[i1][j1].y = needleMap[i1][j1].y/maxLength*step;	
			line(im_output, i1, j1, myround(i1+needleMap[i1][j1].y), myround(j1+needleMap[i1][j1].x), 255);
		}
	}

	writeImage(im_output, fname_output);
	printf("Step = %d. ", step);
	return 0;
}

int myround(double x)
{
	if (x<floor(x)+0.50) return floor(x);
	else return (floor(x)+1);
}

int main(int argc, char **argv) 
{ 
	if (argc <= 4){ 
		printf("Usage: %s <input image 1> <input image 2> <step> <output image>\n", argv[0]); 
		return 0; 
	} 
	char *fname_im1 = argv[1]; 
	char *fname_im2 = argv[2]; 
	int step = atoi(argv[3]); 
	char *fname_output = argv[4]; 
	flow(fname_im1,fname_im2,step,fname_output); 
	printf("U got it\n");
	return 0; 
}
