#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<math.h>
#include<mvcd.h>
#include<sys/types.h>
#include<unistd.h>
#include <sys/stat.h>


typedef unsigned char datatype;

#define VN 50
#define GM 150
#define WM 250
#define CSF 10
#define LabelNum 256


int cortexLables[] = {
	1,2,4,5,6,7,9,10,15,18,19,21,27,32,37,38,41,50,52,54,56,60,61,63,64,69,70,74,75,80,85,86,88,90,92,94,95,96,97,98,99,100,108,110,112,114,118,119,125,130,132,139,140,145,154,159,165,175,196,251};
int wmLables[] = {13,17,30,43,45,57,59,73,83,105,128,133};/* 43 and 128 are added
						to correct some
						mis-labeling */
int jakobLables[] = {1,2,4,5,6,7,9,10,13,15,18,19,21,27,37,38,41,50,52,60,61,63,64,69,70,74,75,80,85,86,88,90,92,94,95,96,97,98,99,100,108,110,112,114,118,119,125,130,132,139,140,145,154,159,165,196,251,17,30,43,45,57,59,73,83,105,128,133};

int LookupGMLabel(int label){
	int i = 0;
	for( i = 0; i < sizeof(cortexLables)/sizeof(int); i++){
		if(cortexLables[i] ==  label){
			return 1;		
		}
	}
	return 0;
}

int LookupWMLabel(int label){
	int i = 0;
	for( i = 0; i < sizeof(wmLables)/sizeof(int); i++){
		if(wmLables[i] ==  label){
			return 1;		
		}
	}
	return 0;
}

void calculate_hood_byIncreasingRadius(Ivector3d *Hood, int HoodSize, int *pixelNumInHood, float xres,float yres,float zres)
{
	int        i = 0, j = 0, k = 0, l = 0, x = 0, y = 0, z = 0, t = 0, r =0;
	int        half = HoodSize/2;

	t++ ;
	for(r=1; r<=half; r++) {
		for(z=-r; z<=r; z+=2*r){
			for(x=-r; x<=r; x++){
				for(y=-r; y<=r; y++){
					Hood[t].x = x/xres ;
					Hood[t].y = y/yres ;
					Hood[t].z = z/zres ;
					t++;
				}
			}
		}
		for(x=-r; x<=r; x+=2*r){
			for(z=-r+1; z<=r-1; z++){
				for(y=-r; y<=r; y++){
					Hood[t].x = x/xres ;
					Hood[t].y = y/yres ;
					Hood[t].z = z/zres ;
					t++;
				}
			}
		}
		
		for(y=-r; y<=r; y+=2*r){
			for(z=-r+1; z<=r-1; z++){
				for(x=-r+1; x<=r-1; x++){
					Hood[t].x = x/xres ;
					Hood[t].y = y/yres ;
					Hood[t].z = z/zres ;
					t++;
				}
			}
		}
		
	}
	(*pixelNumInHood) = t ;
	printf("search nbr voxels points %d\n",t);
}

int GmNbrhood(int i, int j, int k, int xs, int ys , int zs, int nbrVoxelCnt_InMyHood, datatype ***labeledImg, Ivector3d *nb_pnt,datatype ***gmNbrhoodFlag)
{
	int n, x, y, z;
	for(n=1; n < nbrVoxelCnt_InMyHood; n++){
		x = i + nb_pnt[n].x ;
		y = j + nb_pnt[n].y ;
		z = k + nb_pnt[n].z ;
		if(x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0 && gmNbrhoodFlag[z][x][y] < 255 && (1 == LookupGMLabel(labeledImg[z][x][y])))
			return 1;
	}
	return 0;
}
int VNNbrhood(int i,int j, int k, int xs, int ys, int zs, int nbrhoodSize, datatype ***labeledImg)
{
	int l,m,n;
	for( l = ((k - nbrhoodSize)<0?0:(k - nbrhoodSize)); l < (k + nbrhoodSize > zs?zs:(k + nbrhoodSize)); l++){
		for( m = ((i - nbrhoodSize)<0?0:(i - nbrhoodSize)); m < ((i + nbrhoodSize) > xs?xs:(i + nbrhoodSize)); m++){
			for( n = ((j - nbrhoodSize)<0?0:(j - nbrhoodSize)); n < ((j + nbrhoodSize)>ys?ys:(j + nbrhoodSize)); n++){
				if(labeledImg[l][m][n] == 3 || labeledImg[l][m][n] == 8)
					return 1;
			}
		}
	}
	return 0;
}

void showUsage(char **argv){
    char fullVersion[] = "$Rev: 83 $"; //This line will be automatically 
                                       //populated by svn and should be
                                       //of the form "$Rev: 83 $" else it could
                                       //cause a segmentation fault
    char *shortVersion = strchr(fullVersion,' ')+1; //Eliminate te word "$Rev: " 
    char *secondSpaceInVersion = strchr(shortVersion,' '); //Find the space after the number
    *secondSpaceInVersion ='\0';//Eliminate anything after the number, by zeroing the space
    printf("\t \n\%s version: 1.0.%s \n\ ",argv[0],shortVersion);
	printf("USAGE: \033[40;31m %s orgImg segImg labeledImg outImg [options] \n\
\t  \033[31;32m -x <int>               : height size:256\n\
\t  \033[31;33m -y <int>               : width size:256 \n\
\t  \033[31;34m -z <int>               : slice number, if given, generally there are headers in the volume \n\
\t  \033[31;36m -n <int>               : nbr size for correction:1\n\
\033[0m Conments/Questions to: sbia-software@uphs.upenn.edu\n",argv[0]);
	exit(EXIT_FAILURE);
}
int main(int argc, char **argv)
{
	long  i, j, k, x, y, z, m, n;
	int   xs =  256, ys = 256, zs = 124;
	float x_res = 0.9375, y_res = 0.9375, z_res = 1.5;
	datatype ***orgImg,***segImg,***labeledImg,***labelTempImg,***outImg, ***gmNbrhoodImg;
	datatype ***orgLabelledImg;
	char  filename[256];
	long  fileSize = 0;
	int   nb_size = 3;
	int   increased_nb_size = 5;
	int   increased_nbrVoxelCnt_InMyHood, nbrVoxelCnt_InMyHood;
	int  gmVoxelsUpdated = 1e-5 , gmVoxelsUpdated_last, wmVoxelsUpdated = 1e-5 ,wmVoxelsUpdated_last;
	int VoxelsUpdated = 1e-5 ,VoxelsUpdated_last;
	Ivector3d *nb_pnt;
	Ivector3d *increased_nb_pnt;
	int   num_labels[LabelNum], max ;
	int   hdrFlag = 0, verboseFlag = 1;
	int   c;
	int iterNum = 1;
	extern char *optarg;
	FILE  *fp;

	if(argc<5){
		showUsage(argv);
		exit(EXIT_FAILURE);
	}

	while((c=getopt(argc-4,argv+4,"n:x:y:z:V:v")) != -1){
		switch(c){
			case 'n':
				nb_size=atoi(optarg);
				printf("optarg is %s,nb_size is %d\n",optarg,nb_size);
				break;
			case 'x':
				xs=atoi(optarg);
				break;
			case 'y':
				ys=atoi(optarg);
				break;
			case 'z':
				zs=atoi(optarg);
				hdrFlag = 1;
				break;
			case 'V':
				sscanf(optarg,"%f,%f,%f",&x_res,&y_res,&z_res);
				break;
			case 'v':
				verboseFlag = 1;
				break;
			case '?':
				showUsage(argv);
				exit(EXIT_FAILURE);
		}
	}
	sprintf(filename,argv[1]);
	fp=myopen(filename,"rb");
	fseek(fp,(long)0,SEEK_END);
	fileSize = ftell(fp);
	if(!hdrFlag){
		zs = fileSize/(xs * ys * sizeof(datatype));
		rewind(fp);
	}
	if(verboseFlag)printf("filesize is %d, zs=%d,calced zs = %d",fileSize,zs,fileSize/(xs * ys * sizeof(datatype)));
	if( 0 != fileSize%(xs * ys * zs *sizeof(datatype))){
		printf("Input Image has header in it with size of %d\n",fileSize-(zs * xs * ys * sizeof(datatype))); 
		fseek(fp,fileSize-(zs * xs * ys * sizeof(datatype)),0);
	}
	if(verboseFlag){
		printf("Input image para: xs=%d,ys=%d,zs=%d\n",xs,ys,zs);
	}
	orgImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);
	segImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);
	labelTempImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);
	labeledImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);
	outImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);
	gmNbrhoodImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);
	orgLabelledImg = (datatype ***)MyAlloc3d(sizeof(datatype), xs, ys, zs);

	for(z=0; z<zs;z++) {
		for(x=0; x<xs; x++ ) {
			for(y=0;y<ys;y++) {
				orgImg[z][x][y] = 0;
				segImg[z][x][y] = 0;
				labeledImg[z][x][y] = 0;
				outImg[z][x][y] = 0;
				gmNbrhoodImg[z][x][y] = 0;
				orgLabelledImg[z][x][y] = 0;
			}
		}
	}
	for(k=0; k<zs; k++){
		for(i = 0; i<xs; i++){
			fread(orgImg[k][i],sizeof(datatype),ys,fp);
		}
	}
	fclose(fp);
	/* read in the segmented img */
	sprintf(filename,argv[2]);
	fp = myopen(filename,"rb");
	for(k=0; k<zs; k++){
		for(i = 0; i<xs; i++){
			fread(segImg[k][i],sizeof(datatype),ys,fp);
		}
	}
	fclose(fp);
	/* read in the labelled image */
	sprintf(filename,argv[3]);
	 fp = myopen(filename,"rb");
	 for(k=0; k<zs; k++){
		 for(i = 0; i<xs; i++){
			 fread(labeledImg[k][i],sizeof(datatype),ys,fp);
		 }
	 }
	 fclose(fp);
	 
	 /* Iterative remove non-gray matter */
	nb_pnt= Ivector3dalloc1d(nb_size*nb_size*nb_size) ;
	increased_nb_pnt = Ivector3dalloc1d(increased_nb_size * increased_nb_size * increased_nb_size) ;

	calculate_hood_byIncreasingRadius(          nb_pnt,           nb_size,           &nbrVoxelCnt_InMyHood, x_res, y_res, z_res);
	calculate_hood_byIncreasingRadius(increased_nb_pnt, increased_nb_size, &increased_nbrVoxelCnt_InMyHood, x_res, y_res, z_res);

	 /* remove label on zero segmentation */
	 for(z=0;z<zs;z++){
		 for(x=0;x<xs;x++){
			 for(y=0;y<ys;y++){
				 if(0 == segImg[z][x][y] && labeledImg[z][x][y] != 3 && labeledImg[z][x][y] != 8){
					 labeledImg[z][x][y] = 0;
				 }else if(CSF == segImg[z][x][y] && labeledImg[z][x][y] != 3 && labeledImg[z][x][y] != 8 && labeledImg[z][x][y] != 36 && labeledImg[z][x][y] != 101/*&& !VNNbrhood(x,y,z,xs,ys,zs,3,labeledImg)*/) { 
					 labeledImg[z][x][y] = 255;
				 }else if( 13 == labeledImg[z][x][y]){
					 labeledImg[z][x][y] = 253;
				 }
				 /*else if ( 67 == labeledImg[z][x][y] || 76 == labeledImg[z][x][y]){
					 labeledImg[z][x][y] = 0;
				 }*//* for accord project where cerebelem not clean */
			 }
		 }
	 }

	for(k=0;k < zs; k++){
		for(i=0;i < xs;i++){
			for(j=0; j < ys; j++){
				labelTempImg[k][i][j] = labeledImg[k][i][j] ;
				orgLabelledImg[k][i][j] = labeledImg[k][i][j] ;
			}
		}
	}
/*
	for(k=0;k < zs; k++){printf(".");fflush(stdout);
		for(i=0;i < xs;i++){
			for(j=0; j < ys; j++){
				if(1 == GmNbrhood(i, j, k, xs, ys , zs, increased_nbrVoxelCnt_InMyHood, labeledImg, increased_nb_pnt,gmNbrhoodImg)){
					gmNbrhoodImg[k][i][j] = 255;
				}
			}
		}
	}
*/

	for(k=0;k < zs; k++){printf(".");fflush(stdout);
		for(i=0;i < xs;i++){
			for(j=0; j < ys; j++){
				if(LookupGMLabel(labeledImg[k][i][j])){
					for(n=1; n < increased_nbrVoxelCnt_InMyHood; n++){
						x = i + increased_nb_pnt[n].x ;
						y = j + increased_nb_pnt[n].y ;
						z = k + increased_nb_pnt[n].z ;
						if(x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0 && gmNbrhoodImg[z][x][y] < 255){
							gmNbrhoodImg[z][x][y] = 252;
						}
					}
				}
			}
		}
	}
	 /* for every scattered points assigned a label */
	do{
		VoxelsUpdated_last = VoxelsUpdated;
		VoxelsUpdated = 0 ;
		for(k=0;k < zs;k++){
			for(i=0;i < xs; i++){
				for(j=0;j < ys; j++){
					if(segImg[k][i][j] > 0 && !labeledImg[k][i][j]) {
						for(n=1; n < increased_nbrVoxelCnt_InMyHood; n++){
							x = i + increased_nb_pnt[n].x ;
							y = j + increased_nb_pnt[n].y ;
							z = k + increased_nb_pnt[n].z ;
							if( x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0)
								num_labels[labeledImg[z][x][y]] += 1 ;
						}
						max = 0 ;
						for( m=1; m<LabelNum; m++){
							if( num_labels[m]>max ){
								max = m ;
							}
							num_labels[m] = 0 ;
						}
						if( max > 0 ){/*13/iterNum to get maxim possiblility */
							labelTempImg[k][i][j] = max ;
							VoxelsUpdated++ ;
						}
					}
				}
			}
		}
		for(k=0;k < zs; k++){
			for(i=0;i < xs; i++){
				for(j=0; j < ys; j++){
					labeledImg[k][i][j] = labelTempImg[k][i][j] ;
				}
			}
		}
	}while(abs(VoxelsUpdated - VoxelsUpdated_last) > 0);
	printf("\n Scattered points cleared\n");

	 /* correction for VN for segmentations of value 50 */
	do{
	  VoxelsUpdated_last = VoxelsUpdated;
	  VoxelsUpdated = 0 ;
	  for(k=0;k < zs;k++){
	    for(i=0;i < xs; i++){
	      for(j=0;j < ys; j++){
		if(50 == segImg[k][i][j] && labeledImg[k][i][j] != 3 && labeledImg[k][i][j] != 8 ) {
		  for(n=1; n < increased_nbrVoxelCnt_InMyHood; n++){
		    x = i + increased_nb_pnt[n].x ;
		    y = j + increased_nb_pnt[n].y ;
		    z = k + increased_nb_pnt[n].z ;
		    if( x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0 && (labeledImg[z][x][y] == 3 || labeledImg[z][x][y] == 8))
		      num_labels[labeledImg[z][x][y]] += 1 ;
		  }
		  max = 0 ;
		  for( m=1; m<LabelNum; m++){
		    if( num_labels[m]>max ){
		      max = m ;
		    }
		    num_labels[m] = 0 ;
		  }
		  if( max > 0 ){
		    labelTempImg[k][i][j] = max ;
		    VoxelsUpdated++ ;
		  }
		}
	      }
	    }
	  }
	  for(k=0;k < zs; k++){
	    for(i=0;i < xs; i++){
	      for(j=0; j < ys; j++){
		labeledImg[k][i][j] = labelTempImg[k][i][j] ;
	      }
	    }
	  }
	}while(abs(VoxelsUpdated - VoxelsUpdated_last) > 0);
	printf("\n VN corrections done\n");
	
	/* correction for Frontal Lobe WM */

	do{
		VoxelsUpdated_last = VoxelsUpdated;
		VoxelsUpdated = 0 ;
		for(k=0;k < zs;k++){
			for(i=0;i < xs; i++){
				for(j=0;j < ys; j++){
					if(segImg[k][i][j] == WM && 30 == labeledImg[k][i][j]){
						for(n=1; n < nbrVoxelCnt_InMyHood; n++){
							x = i + nb_pnt[n].x ;
							y = j + nb_pnt[n].y ;
							z = k + nb_pnt[n].z ;
							if( x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0 && labeledImg[z][x][y] == 17)
								num_labels[labeledImg[z][x][y]] += 1 ;
						}
						max = 0 ;
						for( m=1; m<LabelNum; m++){
							if( num_labels[m]>max ){
								max = m ;
							}
							num_labels[m] = 0 ;
						}
						if( max > 0 ){/*13/iterNum to get maxim possiblility */
							labelTempImg[k][i][j] = max ;
							VoxelsUpdated++ ;
						}
					}
				}
			}
		}
		for(k=0;k < zs; k++){
			for(i=0;i < xs; i++){
				for(j=0; j < ys; j++){
					labeledImg[k][i][j] = labelTempImg[k][i][j] ;
				}
			}
		}
	}while(abs(VoxelsUpdated - VoxelsUpdated_last) > 0);

	iterNum = 1;

	do{
		gmVoxelsUpdated_last = gmVoxelsUpdated;
		wmVoxelsUpdated_last = wmVoxelsUpdated;
		gmVoxelsUpdated = 0 ;
		wmVoxelsUpdated = 0 ;

		for(k=0;k < zs;k++){
			for(i=0;i < xs; i++){
				for(j=0;j < ys; j++){
					if( gmNbrhoodImg[k][i][j] > 0 && (( GM == segImg[k][i][j] && !LookupGMLabel(labeledImg[k][i][j])) || (WM == segImg[k][i][j] && !LookupWMLabel(labeledImg[k][i][j]))) ) {
						if(WM == segImg[k][i][j] && !LookupWMLabel(labeledImg[k][i][j])){/* WM but not WM label */
							for(n=1; n < nbrVoxelCnt_InMyHood; n++){
					 x = i + nb_pnt[n].x ;
					 y = j + nb_pnt[n].y ;
					 z = k + nb_pnt[n].z ;
					 if( x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0 && (LookupWMLabel(labeledImg[z][x][y])) && WM == segImg[z][x][y] )
					   num_labels[labeledImg[z][x][y]] += 1 ;
				       }
				       max = 0 ;
				       for( m=1; m<LabelNum; m++){
					 if( num_labels[m]>max ){
					   max = m ;
					 }
					 num_labels[m] = 0 ;
				       }
				       if( max > 0 ){/*13/iterNum to get maxim possiblility */
					 labelTempImg[k][i][j] = max ;
					 wmVoxelsUpdated++ ;
				       }
				     } else if(GM == segImg[k][i][j] && !LookupGMLabel(labeledImg[k][i][j]) && 102 != labeledImg[k][i][j] && 203!= labeledImg[k][i][j] ){ /* GM but not GM Label but do not correct thalamus */
				       for(n=1; n < nbrVoxelCnt_InMyHood; n++){
					 x = i + nb_pnt[n].x ;
					 y = j + nb_pnt[n].y ;
					 z = k + nb_pnt[n].z ;
					 if( x < xs && y < ys && z < zs && x >= 0 && y >= 0 && z >= 0 && (LookupGMLabel(labeledImg[z][x][y])) && GM == segImg[z][x][y] )
					   num_labels[labeledImg[z][x][y]] += 1 ;
				       }
				       max = 0 ;
				       for( m=1; m<LabelNum; m++){
					 if( num_labels[m]>max ){
					   max = m ;
					 }
					 num_labels[m] = 0 ;
				       }
				       if(max > 0){
					 labelTempImg[k][i][j] = max ;
					 gmVoxelsUpdated++ ;
				       }
				     } 
				   }
				 }
			 }
		 }		
		 iterNum++; 
		 /* correction done */
		 for(k=0;k < zs; k++){
			 for(i=0;i < xs; i++){
				 for(j=0; j < ys; j++){
					 labeledImg[k][i][j] = labelTempImg[k][i][j] ;
				 }
			 }
		 }
		 printf("VoxelsUpdated= %d,%d\n", abs(gmVoxelsUpdated-gmVoxelsUpdated_last),abs(wmVoxelsUpdated-wmVoxelsUpdated_last)) ;
		 /*		 RemoveScatterLabelPtsUsing3DRegiongrowing(); */
	 }while( abs(gmVoxelsUpdated-gmVoxelsUpdated_last) > 0 || abs(wmVoxelsUpdated-wmVoxelsUpdated_last) > 0);

	 /* Use the original label to recover the mis-Corrected labels */
	 for(k=0;k < zs; k++){
		 for(i=0;i < xs; i++){
			 for(j=0; j < ys; j++){
				 if(orgLabelledImg[k][i][j] > 0 && gmNbrhoodImg[k][i][j] > 0 && !LookupWMLabel(orgLabelledImg[k][i][j]) && !LookupGMLabel(orgLabelledImg[k][i][j]) && orgLabelledImg[k][i][j] !=0 &&(labeledImg[k][i][j] != 3 || labeledImg[k][i][j] != 8) ){
					 labeledImg[k][i][j] = orgLabelledImg[k][i][j];
				 }
				 /*else if( (segImg[k][i][j] == VN ) && (labeledImg[k][i][j] != 3 || labeledImg[k][i][j] != 8)){
				   printf("E ");
				   }*/
			 }
		 }
	 }

		  
	 /* output the reuslt */
	 sprintf(filename,argv[4]);
	 fp=myopen(filename,"wb");        
	 for(k=0;k<zs;k++){
		 for(i=0;i<xs;i++){
			 fwrite(labeledImg[k][i],sizeof(datatype),ys,fp);
		 }
	 }
	 fclose(fp);
	 return 0;
}
