#include "stdio.h"
#include "stdlib.h"
#include "math.h"
#include "paulslib.h"
#include "complexlib.h"
#include "sigproclib.h"
#include "bitmaplib.h"

/*
	Workspace for experimenting with 2D image filtering using 2D-FFT
*/

int WritePPM(char *,int,int,COMPLEX **,char);
int WriteOOGL(char *,int,int,COMPLEX **,char);
void WriteOffVector(FILE *,XYZ,XYZ,COLOUR);

int main(int argc,char **argv)
{
	int i,j,c;
	int width,height,filter;
	double mean = 0;
	double tstart,tstop;
	double cutoff1,cutoff2,power=1;
   double r,r1,r2,r3,r4;
	COMPLEX **grid;
	char fname[64];
	FILE *fptr;

	/* Check the arguments */
	if (argc < 5) {
		fprintf(stderr,"Usage: %s filename width height filter\n",argv[0]);
		fprintf(stderr,"Where: filter 1 is low pass\n");
		fprintf(stderr,"              2 is band pass\n");
		fprintf(stderr,"              3 is high pass\n");
		fprintf(stderr,"              4 1 / f^p\n");
		exit(0);
	}
	width  = atoi(argv[2]);
	height = atoi(argv[3]);
	filter = atoi(argv[4]);
	if (filter == 1 || filter == 3) {
		if (argc < 6) {
			fprintf(stderr,"Require a cutoff index for filter types 1 amd 3\n");
			exit(-1);
		}
		cutoff1 = atof(argv[5]);
	}
   if (filter == 2) {
      if (argc < 7) {
         fprintf(stderr,"Require 2 cutoff indices for filter type 2\n");
         exit(-1);
      }
      cutoff1 = atof(argv[5]);
		cutoff2 = atof(argv[6]);
   }
	if (filter == 4) {
      if (argc < 6) {
         fprintf(stderr,"Require a power 1/f filter types 4\n");
         exit(-1);
      }
      power = atof(argv[5]);
   }

	/* Allocate memory for the grid */
	if ((grid = (COMPLEX **)malloc(width * sizeof(COMPLEX))) == NULL) {
      fprintf(stderr,"Unable to allocate memory for the grid\n");
      exit(0);
   }
	for (i=0;i<width;i++) {
		if ((grid[i] = (COMPLEX *)malloc(height * sizeof(COMPLEX))) == NULL) {
			fprintf(stderr,"Unable to allocate memory for the grid\n");
			exit(0);
		}
	}

	/* Read the image */
	if ((fptr = fopen(argv[1],"r")) == NULL) {
		fprintf(stderr,"Unable to open file \"%s\"\n",argv[1]);
		exit(0);
	}
	for (j=0;j<height;j++) {
		for (i=0;i<width;i++) {
			if ((c = fgetc(fptr)) == EOF) {
				fprintf(stderr,"Unexpected end of file\n");
				exit(-1);
			}
			grid[i][j].real = c;
			grid[i][j].imag = 0;
			mean += c;
		}
	}
	fclose(fptr);

   /* Write the mean subtracted image */
   sprintf(fname,"%s_orig_%d.ppm",argv[1],filter);
   WritePPM(fname,width,height,grid,'t');
   sprintf(fname,"%s_orig_%d.oogl",argv[1],filter);
   WriteOOGL(fname,width,height,grid,'t');

	/* Subtract the mean */
	mean /= (width * height);
	fprintf(stderr,"Mean (to be removed) is %g\n",mean);
   for (j=0;j<height;j++) 
      for (i=0;i<width;i++) 
			grid[i][j].real -= mean;

	/* Perform the FFT */
	tstart = GetRunTime();
	FFT2D(grid,width,height,1);
	tstop = GetRunTime();
	fprintf(stderr,"Time for 2D FFT is: %g ms\n",(tstop-tstart)*1000.0);

	/* Write the transformed image */
	sprintf(fname,"%s_fft_%d.ppm",argv[1],filter);
	WritePPM(fname,width,height,grid,'f');
   sprintf(fname,"%s_fft_%d.oogl",argv[1],filter);
   WriteOOGL(fname,width,height,grid,'f');

	/* 
		Perform the filtering 
		0 - no filter
		1 - low pass rectangular
		2 - band pass rectangular
		3 - high pass rectangular
		4 - 1 / f^p
	*/
   for (i=0;i<width;i++) {
      for (j=0;j<height;j++) {
         r1 = i*i + j*j;
         r2 = (i-width)*(i-width) + j*j;
         r3 = i*i + (j-height)*(j-height);
         r4 = (i-width)*(i-width) + (j-height)*(j-height);
         r = MIN(r1,r2);
         r = MIN(r,r3);
         r = MIN(r,r4);

			switch (filter) {
			case 0:		/* No filtering */
				break;
			case 1:		/* Low pass */
            if (r > cutoff1*cutoff1) {
               grid[i][j].real = 0;
               grid[i][j].imag = 0;
            }
				break;
			case 2:		/* Band pass */
            if (r < cutoff1*cutoff1 || r > cutoff2*cutoff2) {
               grid[i][j].real = 0;
               grid[i][j].imag = 0;
            }
				break;
   		case 3:     /* High pass */
            if (r < cutoff1*cutoff1) {
               grid[i][j].real = 0;
               grid[i][j].imag = 0;
            }
				break;
			case 4:
				if (r > 0) {
					grid[i][j].real *= pow(1/r,power);
					grid[i][j].imag *= pow(1/r,power);
				}
				break;
			default:
				fprintf(stderr,"No such filter defined\n");
				exit(-1);
				break;
			}

		} /* j */
	} /* i */
	
	/* Write the filtered transformed image */
	sprintf(fname,"%s_fftfilt_%d.ppm",argv[1],filter);
	WritePPM(fname,width,height,grid,'f');
   sprintf(fname,"%s_fftfilt_%d.oogl",argv[1],filter);
   WriteOOGL(fname,width,height,grid,'f');

	/* Write the filtered image */
	FFT2D(grid,width,height,-1);
	sprintf(fname,"%s_invfft_%d.ppm",argv[1],filter);
	WritePPM(fname,width,height,grid,'t');
   sprintf(fname,"%s_invfft_%d.oogl",argv[1],filter);
   WriteOOGL(fname,width,height,grid,'t');

}

/*
	Write as a PPM image
	type is either 
	t - time signal assumed to be in the real part of "c"
   f - the frequency case is split to put DC in the middle of the image 
		 Plot the magnitude
*/
int WritePPM(char *fname,int nx,int ny,COMPLEX **c,char type) 
{
	int i,j,ii,jj;
	int warning = FALSE;
	double scale=0,mag,rmin=1e32,rmax=-1e32;
	BITMAP *bm,col;
	FILE *fptr;

	/* Determine the maximum magnitude */
	if ((bm = CreateBitmap(nx,ny)) == NULL) {
		fprintf(stderr,"WritePPM - Failed to create bitmap\n");
		return(FALSE);
	}
	for (i=0;i<nx;i++) {
		for (j=0;j<ny;j++) {
			mag = sqrt(c[i][j].real*c[i][j].real+c[i][j].imag*c[i][j].imag);
			scale = MAX(mag,scale);
			rmin = MIN(rmin,c[i][j].real);
			rmax = MAX(rmax,c[i][j].real);
			if (type == 't' && !warning && ABS(c[i][j].imag) > 0.00001) {
				fprintf(stderr,"Warning, imaginary part of time series != 0\n");
				warning = TRUE;
			}
		}
	}
	if (rmin >= rmax) {
		fprintf(stderr,"Warning - The real minimum is >= real maximum\n");
		rmin -= 0.01;
		rmax += 0.01;
	}

	/* Form the greyscale image */
   for (i=0;i<nx;i++) {
      for (j=0;j<ny;j++) {
			if (type == 'f') {
				mag = sqrt(c[i][j].real*c[i][j].real+c[i][j].imag*c[i][j].imag);
				if (i <= nx/2) {
					if (j <= ny/2) {
						ii = i+nx/2-1;
						jj = j+ny/2-1;
					} else {
						ii = i+nx/2-1;
						jj = j-ny/2-1;
					}
				} else {
            	if (j <= ny/2) {
						ii = i-nx/2-1;
						jj = j+ny/2-1;
            	} else {
						ii = i-nx/2-1;
						jj = j-ny/2-1;
            	}
				}
				col.r = 255 * mag / scale;
			} else {
				ii = i;
				jj = j;
				col.r = 255 * (c[ii][jj].real - rmin) / (rmax - rmin);
			}
         col.g = col.r;
         col.b = col.r;
			DrawPixel(bm,nx,ny,ii,jj,col);
      }
   }

	/* Write the file */
	if ((fptr = fopen(fname,"w")) == NULL) {
		fprintf(stderr,"Error opening image file\n");
		DestroyBitmap(bm);
		return(FALSE);
	}
	WriteBitmap(fptr,bm,nx,ny,2);
	DestroyBitmap(bm);
	fclose(fptr);
	
	return(TRUE);
}

/*
	Same as WritePPM but for oogl files
*/
int WriteOOGL(char *fname,int nx,int ny,COMPLEX **c,char type)
{
   int i,j,ii,jj;
   double scale=0,mag,rmin=1e32,rmax=-1e32;
	float **f;
	COLOUR black={0.0,0.0,0.0};
	XYZ p1,p2;
   FILE *fptr;

	if ((f = (float **)malloc(nx*sizeof(float *))) == NULL) {
		fprintf(stderr,"Failed to allocate memory for grid\n");
		return(FALSE);
	}
	for (i=0;i<nx;i++) {
		if ((f[i] = (float *)malloc(ny * sizeof(float))) == NULL) {
      	fprintf(stderr,"Failed to allocate memory for grid\n");
      	return(FALSE);
		}
	}

   /* Determine the maximum magnitude */
   for (i=0;i<nx;i++) {
      for (j=0;j<ny;j++) {
         mag = sqrt(c[i][j].real*c[i][j].real+c[i][j].imag*c[i][j].imag);
         scale = MAX(mag,scale);
         rmin = MIN(rmin,c[i][j].real);
         rmax = MAX(rmax,c[i][j].real);
      }
   }
   if (rmin >= rmax) {
      rmin -= 0.01;
      rmax += 0.01;
   }

   /* Form the greyscale image */
	for (j=0;j<ny;j++) {
   	for (i=0;i<nx;i++) {
         if (type == 'f') {
				mag = sqrt(c[i][j].real*c[i][j].real+c[i][j].imag*c[i][j].imag);
            if (i <= nx/2) {
               if (j <= ny/2) {
                  ii = i+nx/2-1;
                  jj = j+ny/2-1;
               } else {
                  ii = i+nx/2-1;
                  jj = j-ny/2-1;
               }
            } else {
               if (j <= ny/2) {
                  ii = i-nx/2-1;
                  jj = j+ny/2-1;
               } else {
                  ii = i-nx/2-1;
                  jj = j-ny/2-1;
               }
            }
				f[ii][jj] = 0.25 * nx * mag / scale;
         } else {
				ii = i;
				jj = j;
				f[ii][jj] = 0.25 * nx * (c[ii][jj].real - rmin) / (rmax - rmin);
			} 
      }
   }

   /* Open the file */
   if ((fptr = fopen(fname,"w")) == NULL) {
      fprintf(stderr,"Error opening oogl file\n");
      return(FALSE);
   }
	fprintf(fptr,"LIST\n");
   fprintf(fptr,"ZMESH\n%d %d\n",nx,ny);
	for (i=0;i<nx;i++) {
   	for (j=0;j<ny;j++) {
         fprintf(fptr,"%g ",f[i][j]);
      }
      fprintf(fptr,"\n");
   }
	p1.x = 0; p1.y = 0; p1.z = 0;
	p2.x = nx; p2.y = 0; p2.z = 0;
	WriteOffVector(fptr,p1,p2,black);
   p1.x = nx; p1.y = 0; p1.z = 0;
   p2.x = nx; p2.y = ny; p2.z = 0;
   WriteOffVector(fptr,p1,p2,black);
   p1.x = nx; p1.y = ny; p1.z = 0;
   p2.x = 0; p2.y = ny; p2.z = 0;
   WriteOffVector(fptr,p1,p2,black);
   p1.x = 0; p1.y = ny; p1.z = 0;
   p2.x = 0; p2.y = 0; p2.z = 0;
   WriteOffVector(fptr,p1,p2,black);
   p1.x = 0; p1.y = 0; p1.z = 0;
   p2.x = 0; p2.y = 0; p2.z = 0.25 * nx;
   WriteOffVector(fptr,p1,p2,black);

  	fclose(fptr);
	for (i=0;i<nx;i++)
		free(f[i]);
	free(f);
   return(TRUE);
}

/*
   Write a vector in OFF format
   Do not use for points
*/
void WriteOffVector(FILE *fptr,XYZ p1,XYZ p2,COLOUR colour)
{
   fprintf(fptr,"{ = VECT\n");
   fprintf(fptr,"\t1 2 1 2 1\n");
   fprintf(fptr,"\t%g %g %g\n",p1.x,p1.y,p1.z);
   fprintf(fptr,"\t%g %g %g\n",p2.x,p2.y,p2.z);
   fprintf(fptr,"\t%g %g %g 1\n",colour.r,colour.g,colour.b);
   fprintf(fptr,"}\n");
}

#include "sigproclib.c"
#include "bitmaplib.c"
#include "paulslib.c"

