#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>
#include <jsfile.h>
#include <rn.h>
#include <float_image.h>
#include <pst_height_map.h>



float_image_t* height_map_shrink_one_pixel(float_image_t* IZ);

float_image_t* height_map_shrink_one_pixel(float_image_t* IZ){
	int NC = IZ->sz[0];
	int NX = IZ->sz[1];
	int NY = IZ->sz[2];
	float_image_t* IJ = float_image_new(NC,NX-1,NY-1);
	int c;
	for( c = 0; c < NC; c++)
	{
		int x;
		for(x = 0; x < NX -1; x++)
		{
			int y;
      			for(y = 0; y < NY -1; y++)
			{
				double P00 = float_image_get_sample(IZ,c,x,y);
				double P10 = float_image_get_sample(IZ,c,x+1,y);
				double P01 = float_image_get_sample(IZ,c,x,y+1);
				double P11 = float_image_get_sample(IZ,c,x+1,y+1);

				double val = (P00 + P01 + P10 + P11)/4.0;
				float_image_set_sample(IJ,c,x,y,val);
      			}
		}
	}
	return IJ;
}

void add_border_weights(float_image_t* IW);

void add_border_weights(float_image_t* IW){
  int NX = IW->sz[1];
  int NY = IW->sz[2];
  int x,y;
  for(y = 0; y < NY; y++){
    for( x = 0; x < NX; x++){
      if( (x== 0) || (y == 0) ){
	float_image_set_sample(IW,0,x,y,0);
      }
    }
  }
}

float_image_t* readFNI(char* filename);


float_image_t* readFNI(char* filename){
	FILE* arq = open_read(filename,TRUE);
	float_image_t* img = float_image_read(arq);
	fclose(arq);
	return img;
}

void writeFNI(char* filename, float_image_t* img);

void writeFNI(char* filename, float_image_t* img){
	FILE* arq = open_write(filename,TRUE);
	float_image_write(arq,img);
	fclose(arq);
}

int main(int argc, char** argv){
	if(argc < 4){
		fprintf(stderr,"program usage\ncompute_rms <reference map> <compared map> <output_file> [weight_map]\n");
		return 0;
	}
	char* hmREF_filename = argv[1];
	char* hmCPR_filename = argv[2];
	char* output_filename = argv[3];
	char* wm_filename = NULL;
	if(argc == 5){
		wm_filename = argv[4];
	}

        /* First read the reference file */
	float_image_t* imgREF = readFNI(hmREF_filename);
	/* First read the compared file */
	float_image_t* imgCPR = readFNI(hmCPR_filename);
	/* Enforce same number of channels */
	if(imgREF->sz[0] != imgCPR->sz[0])
	{
		fprintf(stderr,"ERROR: Comparing maps with different number of channels\n");
		assert(FALSE);
	}


	float_image_t* imgW = NULL;
	if(wm_filename != NULL){
		imgW = readFNI(wm_filename);
	}

	
	
	int NX_REF = imgREF->sz[1];
	int NY_REF = imgREF->sz[2];

	int NX_CPR = imgCPR->sz[1];
	int NY_CPR = imgCPR->sz[2];

	int NX_WGT;
	int NY_WGT;

   	fprintf(stderr,"Height maps sizes: REF[%d x %d] and CPR[%d x %d]\n",NX_REF,NY_REF,NX_CPR,NY_CPR);
	if(imgW != NULL){ 
		NX_WGT = imgW->sz[1];
		NY_WGT = imgW->sz[2];
		fprintf(stderr,"weight map W[%d x %d]\n",NX_WGT,NY_WGT);
	}

	
	if ( (NX_REF != NX_CPR) || (NY_REF !=NY_CPR) )
	{
		/*The reference file may be one pixel bigger than the compared one so we reduce it, using the weights*/
		if ( (NX_REF == (NX_CPR+1))  && (NY_REF == (NY_CPR+1)) )
		{	
			fprintf(stderr,"Shrunk reference map: 1 pixel larger\n");
			float_image_t* imgREF_red = height_map_shrink_one_pixel(imgREF);
			float_image_free(imgREF);
			imgREF = imgREF_red;
			NX_REF--;
			NY_REF--;
			
		}
		else if ( (NX_REF == (NX_CPR-1))  && (NY_REF == (NY_CPR-1)) )
		{	/*It can happens that we want compare the inverse ! */
			fprintf(stderr,"Shrunk compared map: 1 pixel larger\n");
			float_image_t* imgCPR_red = height_map_shrink_one_pixel(imgCPR);
			float_image_free(imgCPR);
			imgCPR = imgCPR_red;
			NX_CPR--;
			NY_CPR--;
			
		}
		else
		{
			fprintf(stderr,"ERROR: Compared maps with incompatible sizes\n");
			assert(FALSE);
		}
	}
	
	if( imgW != NULL ){
		if( (NX_CPR != NX_WGT)  || (NY_CPR != NY_WGT) ){
			fprintf(stderr,"ERROR: Weight map with incompatible size\n");
			assert(FALSE);
		}
	}
	
	/* We will add a border zero, to avoid confusions.. it will create a weight map if needed */
	if(imgW == NULL){
	  imgW = float_image_new(1,NX_CPR,NY_CPR);
	  float_image_fill_channel(imgW,0,1.0);
	}
	add_border_weights(imgW);


	int NC = imgREF->sz[0];
	float_image_t* img_diff = float_image_new(NC,NX_CPR,NY_CPR);
	int c;
	for( c = 0; c < NC ; c++)
	{
		double sAZP,sBZP; /* The standard deviations of the values of {imgREF,imgCPR}, each from its own mean value. */
		double sEZP; /* The root-mean-square value of the error */
		double sreP; /* The relative error {sEZ/sMZ} where {sMZ = hypot(imgREF,imgCPR)/sqrt(2)}. */
		
		float_image_t* imgREF_channel = float_image_new(1,NX_REF,NY_REF);
		float_image_t* imgCPR_channel = float_image_new(1,NX_CPR,NY_CPR);

		float_image_set_channel(imgREF_channel, 0, imgREF, c);
		float_image_set_channel(imgCPR_channel, 0, imgCPR, c);

		float_image_t* img_diff_channel = pst_height_map_compare(
							imgREF_channel,
							imgCPR_channel,
							imgW,
							TRUE,
							&sAZP,
    							&sBZP,
    							&sEZP,
							&sreP
						);
		fprintf(stdout,"%9.6lf %9.6lf %9.6lf %9.6lf\n",sEZP,sreP,sAZP,sBZP);
		
		float_image_set_channel(img_diff, c, img_diff_channel, 0);
				
		float_image_free(imgREF_channel);
		float_image_free(imgCPR_channel);
		float_image_free(img_diff_channel);
	}
	
	writeFNI(output_filename,img_diff);

	return 0;
}