// convolve.cpp
//
// Karl's Nifty Convolve by FFT thingy
// Nov 4, 2000
#include <stdio.h> 
#include <stdlib.h>
#include <string>
#include <fstream.h>
#include <math.h>
#include <iostream.h>

#define MAXSIGNALSIZE		520
#define INSIGNALSOURCE		argv[1]	//"2D_White_Box.pgm"
#define KERNELSOURCE		argv[2] //"2D_White_Box.pgm"
#define OUTSIGNALDEST		argv[3] //"2D_Square_Square.pgm"

struct complex {
	float R,I;
};

int MAX(int x, int y);
int MIN(int x, int y);
int getImageFromFile(char *filename, int signalarray[MAXSIGNALSIZE][MAXSIGNALSIZE],
				 int *sizex, int *sizey, int *sdepth);
int writeImageToFile(char *filename, int signalarray[MAXSIGNALSIZE][MAXSIGNALSIZE], 
			int sizex, int sizey, int sdepth);
void TwoDUFFT(complex f[MAXSIGNALSIZE][MAXSIGNALSIZE], int sizex, int sizey, int depth);
void TwoDFFT(complex f[MAXSIGNALSIZE][MAXSIGNALSIZE], int sizex, int sizey, int depth);
void FFT(complex signalarray[], int size);
void uFFT(complex signalarray[], int size);
void FFTs(complex signalarray[], int size);
void uFFTs(complex signalarray[], int size);
void Mult(complex s[MAXSIGNALSIZE][MAXSIGNALSIZE], 
		complex k[MAXSIGNALSIZE][MAXSIGNALSIZE],
		int sizex, int sizey);
int maxpix;


//Main Program **************************************************
int main(int argc, char *argv[]) {

	//cout << "Hello\n";
	maxpix=0;
	if (argc<3) {
		cout << "\nUsage: " << argv[0] << " signal.dat kernel.dat out.dat\n\n";
		exit(0);
	}
	
	int insignal[MAXSIGNALSIZE][MAXSIGNALSIZE];
	int kernel[MAXSIGNALSIZE][MAXSIGNALSIZE];
	complex ker[MAXSIGNALSIZE][MAXSIGNALSIZE];
	complex sig[MAXSIGNALSIZE][MAXSIGNALSIZE];

	//cout << "Allocated\n";	

	int insizex, insizey, depth;
	int ksizex, ksizey, kdepth;
	
	if (!getImageFromFile(INSIGNALSOURCE, insignal, &insizex, &insizey, &depth)) {
		cout << "\nError, could not open file " << INSIGNALSOURCE << "\n\n";
		exit(0);
	}
	
	if (!getImageFromFile(KERNELSOURCE, kernel, &ksizex, &ksizey, &kdepth)) {
		cout << "\nError, could not open file " << KERNELSOURCE << "\n\n";
		exit(0);
	}
	

	int newsizex, newsizey, newksizex, newksizey;

	for (newsizex=1; newsizex<insizex; newsizex*=2) {
	}
	for (newsizey=1; newsizey<insizey; newsizey*=2) {
	}
	for (newksizex=1; newksizex<ksizex; newksizex*=2) {
	}
	for (newksizey=1; newksizey<ksizey; newksizey*=2) {
	}
	newsizex=MAX(newsizex,newksizex);
	newsizey=MAX(newsizey,newksizey);

	for (int x=0; x<newsizex; x++) {
		for (int y=0; y<newsizey; y++) {
			if ((x<((newsizex/2) - (insizex/2))) ||
				(y<((newsizey/2) - (insizey/2))) ||
				(x>=((newsizex/2) + (insizex/2))) ||
				(y>=((newsizey/2) + (insizey/2)))) {
					sig[x][y].R = 0;
			}
			else {
				sig[x][y].R = 
					insignal[x-newsizex/2+insizex/2][y-newsizey/2+insizey/2];
			}
			sig[x][y].I = 0;
			if ((x<((newksizex/2) - (ksizex/2))) ||
				(y<((newksizey/2) - (ksizey/2))) ||
				(x>=((newksizex/2) + (ksizex/2))) ||
				(y>=((newksizey/2) + (ksizey/2)))) {
					ker[x][y].R = 0;
			}
			else {
				ker[x][y].R = 
					kernel[x-newksizex/2+ksizex/2][y-newksizey/2+ksizey/2];
			}
			ker[x][y].I = 0;
		}
	}	
	insizex=newsizex;
	insizey=newsizey;
	ksizex=newsizex;
	ksizey=newsizey;
	
	TwoDFFT(sig,insizex,insizey,depth);
	TwoDFFT(ker,ksizex,ksizey,kdepth);

	Mult(sig,ker,insizex,insizey);
	
	TwoDUFFT(sig,insizex,insizey,depth);

	int midx = insizex/2;
	int midy = insizey/2;
	for (int x=0; x<midx; x++) {
		for (int y=0; y<midy; y++) {
			ker[x][y] = sig[midx + x][midy + y];
			ker[midx+x][midy+y] = sig[x][y];
			ker[midx+x][y] = sig[x][midy+y];
			ker[x][midy+y] = sig[midx+x][y];
		}
	}
	
	for (int x=0; x<insizex; x++) {
		for (int y=0; y<insizey; y++) {
			insignal[x][y] = (int)ker[x][y].R;
		}
	}
	
	if (!writeImageToFile(OUTSIGNALDEST, insignal, insizex, insizey, depth)) {
		cout << "\nError, could not write to file " << OUTSIGNALDEST << "\n";
	}
		
	cout << "Convolved image written:	" << OUTSIGNALDEST << "\n";	
}

//Dump signal file to an array ***********************************
int getImageFromFile(char *filename, int signalarray[MAXSIGNALSIZE][MAXSIGNALSIZE],
			 int *sizex, int *sizey, int *sdepth) {
	ifstream Infile;
	Infile.open(filename, ios::in);
	if (!Infile) return 0;
	else {
		string image_type;
		int height, width, depth, xin, yin;

		Infile >> image_type >> width >> height >> depth;

		//cout << "Reading ... size = " << height << ", " << width << "\n";		

		*sizex = width;
		*sizey = height;
		*sdepth = depth;
		
		if (image_type!="P5") {
			cout << "\nError: " << image_type << " is not a valid image type.\n\n";
			exit(1);
		}
	
		xin = 0; yin = 0;
		unsigned char greyscale;
		while (!Infile.eof()) {
			Infile.read(&greyscale,1);
			if (greyscale>depth) {
				cout << "Greyscale Out of Bounds! = " << greyscale << "\n";
			}
			else {
				signalarray[xin][yin]=(int)greyscale;
			}
			if (xin<(width-1)) {
				xin++;
			}
			else {
				xin=0;
				yin++;
			}
		}
		Infile.close();
		return 1;
	}
}

//Dump array to signal file ***************************************
int writeImageToFile(char *filename, int signalarray[MAXSIGNALSIZE][MAXSIGNALSIZE],
				int sizex, int sizey, int sdepth) {
	maxpix=0;
	for (int x=0; x<sizex; x++) {
		for (int y=0; y<sizey; y++) {
			if (signalarray[x][y]>maxpix) { maxpix=signalarray[x][y]; }
		}
	}
	ofstream outfile;
	outfile.open(filename, ios::out);
	outfile << "P5\n" << sizex << " " << sizey << "\n"
		<< sdepth << "\n";
	if (!outfile) {
		return 0;
	}
	else {
		outfile.setf(ios::fixed);
		for (int y=0; y<sizey; y++) {
			for (int x=0; x<sizex; x++) {
				//cout << signalarray[x][y] << "\t";
				if (signalarray[x][y]<0) signalarray[x][y]=0;
				outfile << (unsigned char)
					(((double)signalarray[x][y]/(double)maxpix)*(sdepth));
			}
		}
		outfile.close();
		return 1;
	}
}
void TwoDUFFT(complex f[MAXSIGNALSIZE][MAXSIGNALSIZE], int sizex, int sizey, int depth) {
	for (int y=0; y<sizey; y++) {
		complex t1[MAXSIGNALSIZE];
		for (int x=0; x<sizex; x++) {
			t1[x].R = f[x][y].R;
			t1[x].I = f[x][y].I;
		}
		uFFTs(t1,sizex);
		for (int x=0; x<sizex; x++) {
			f[x][y].R = t1[x].R;
			f[x][y].I = t1[x].I;
		}
	}
	for (int x=0; x<sizex; x++) {
		complex t1[MAXSIGNALSIZE];
		for (int y=0; y<sizey; y++) {
			t1[y].R = f[x][y].R;
			t1[y].I = f[x][y].I;
		}
		uFFTs(t1,sizey);
		for (int y=0; y<sizey; y++) {
			f[x][y].R = t1[y].R;
			f[x][y].I = t1[y].I;
		}
	}
}


void TwoDFFT(complex f[MAXSIGNALSIZE][MAXSIGNALSIZE], int sizex, int sizey, int depth) {
	for (int y=0; y<sizey; y++) {
		complex t1[MAXSIGNALSIZE];
		for (int x=0; x<sizex; x++) {
			t1[x].R = f[x][y].R;
			t1[x].I = f[x][y].I;
		}
		FFTs(t1,sizex);
		for (int x=0; x<sizex; x++) {
			f[x][y].R = t1[x].R;
			f[x][y].I = t1[x].I;
		}
	}
	for (int x=0; x<sizex; x++) {
		complex t1[MAXSIGNALSIZE];
		for (int y=0; y<sizey; y++) {
			t1[y].R = f[x][y].R;
			t1[y].I = f[x][y].I;
		}
		FFTs(t1,sizey);
		for (int y=0; y<sizey; y++) {
			f[x][y].R = t1[y].R;
			f[x][y].I = t1[y].I;
		}
	}
}

//FFT starting function ********************************************
void FFTs(complex signalarray[], int size) {
	FFT(signalarray,size);
	//size=1;
	for (int loop = 0; loop<size; loop++) {
		signalarray[loop].R = signalarray[loop].R/sqrt(size);
		signalarray[loop].I = signalarray[loop].I/sqrt(size);
	}
}

void FFT(complex signalarray[], int size) {
	if (size==2) {
		complex temp = signalarray[0];
		signalarray[0].R = signalarray[0].R + signalarray[1].R;
		signalarray[0].I = signalarray[0].I + signalarray[1].I;
		signalarray[1].R = temp.R - signalarray[1].R;
		signalarray[1].I = temp.I - signalarray[1].I;
		//cout << "SigArray (Base) -> (" << signalarray[0].R << "," << signalarray[0].I 
		//	<< ") (" << signalarray[1].R << "," << signalarray[1].I << ")\n";
	}
	else {
		complex g[size/2];
		complex h[size/2];
		for (int loop=0; loop<(size/2); loop++) {
			g[loop]=signalarray[2*loop];
			h[loop]=signalarray[2*loop + 1];
		}
		FFT(g, size/2);
		FFT(h, size/2);
		complex gsum,hsum;
		//cout << "Sig array tally -> ";
		for (int loop=0; loop<size; loop++) {
			signalarray[loop].R = g[loop % (size/2)].R 
									+ cos(2*3.1415*loop/size)*(h[loop % (size/2)].R)
									+ sin(2*3.1415*loop/size)*h[loop % (size/2)].I;
			signalarray[loop].I = g[loop % (size/2)].I 
									- sin(2*3.1415*loop/size)*(h[loop % (size/2)].R)
									+ cos(2*3.1415*loop/size)*h[loop % (size/2)].I;
									
			//cout <<"(" << signalarray[loop].R << "," << signalarray[loop].I << ") ";
		}
		//cout << "\n";
	}
}


void uFFTs(complex signalarray[], int size) {
	uFFT(signalarray,size);
	for (int loop = 0; loop<size; loop++) {
		signalarray[loop].R = signalarray[loop].R/sqrt(size);
		signalarray[loop].I = signalarray[loop].I/sqrt(size);
		//cout <<"(" << signalarray[loop].R << "," << signalarray[loop].I << ") ";
	}
	//cout << "\n";
}

void uFFT(complex signalarray[], int size) {
	if (size==2) {
		complex temp = signalarray[0];
		signalarray[0].R = signalarray[0].R + signalarray[1].R;
		signalarray[0].I = signalarray[0].I + signalarray[1].I;
		signalarray[1].R = temp.R - signalarray[1].R;
		signalarray[1].I = temp.I - signalarray[1].I;
		//cout << "SigArray (Base) -> (" << signalarray[0].R << "," << signalarray[0].I 
		//	<< ") (" << signalarray[1].R << "," << signalarray[1].I << ")\n";
	}
	else {
		complex g[size/2];
		complex h[size/2];
		for (int loop=0; loop<(size/2); loop++) {
			g[loop]=signalarray[2*loop];
			h[loop]=signalarray[2*loop + 1];
		}
		uFFT(g, size/2);
		uFFT(h, size/2);
		complex gsum,hsum;
		//cout << "Sig array tally -> ";
		for (int loop=0; loop<size; loop++) {
			signalarray[loop].R = g[loop % (size/2)].R 
									+ cos(2*3.1415*loop/size)*(h[loop % (size/2)].R)
									- sin(2*3.1415*loop/size)*h[loop % (size/2)].I;
			signalarray[loop].I = g[loop % (size/2)].I 
									+ sin(2*3.1415*loop/size)*(h[loop % (size/2)].R)
									+ cos(2*3.1415*loop/size)*h[loop % (size/2)].I;
									
			//cout <<"(" << signalarray[loop].R << "," << signalarray[loop].I << ") ";
		}
		//cout << "\n";
	}
}

void Mult(complex s[MAXSIGNALSIZE][MAXSIGNALSIZE], 
		complex k[MAXSIGNALSIZE][MAXSIGNALSIZE],
		int sizex, int sizey) {

	for (int x=0; x<sizex; x++) {
		for (int y=0; y<sizey; y++) {
			complex t;
			t.R = ((k[x][y].R * s[x][y].R) 
					- (k[x][y].I * s[x][y].I));
			t.I = (k[x][y].R * s[x][y].I + k[x][y].I * s[x][y].R);
			s[x][y].R = t.R;
			s[x][y].I = t.I;
		}
	}
}

int MAX(int x,int y) {
	if (x>y) return x;
	else return y;
}

int MIN(int x,int y) {
	if (x>y) return y;
	else return x;
}