#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <ctype.h>
#include <dirent.h>

int Usage()
	{

	fprintf(stderr,
"\n"
"FixFits (C)Copyright 2005 Anthony Wesley. All Rights Reserved.\n\n"
"Usage: fixfits [ option ] [ option ] ... input-file(s) or directories\n\n"
"Options are:\n"
"    -add N		Add integer value N to each pixel value (Default 0).\n"
"			N may be negative.\n\n"
"    -scale S		Multiply each pixel value by decimal value S \n"
"			(default 1.0).\n\n"
"    -bzero Z		Set the BZERO header value to the integer value Z\n"
"			Default is to leave it unchanged.\n\n"
"    -force		Force fixfits to process the input files even if it\n"
"			doesn't want to. By default fixfits adds an extra\n"
"			header to indicate that it's processed a file,\n"
"			and if it detects that header in an input file then\n"
"			it will be skipped. This is to let you always say\n"
"			something like 'fixfits *.fit' and have it only\n"
"			process new files.\n\n"
"    -output f		Write the new data to file named f instead of the\n"
"			default behaviour which is to overwrite the original\n"
"			file. Only useful when processing a single file.\n\n"
"    -padbytes B		Add extra B bytes to the end of the file.\n"
"			Can be used to correct for some programs that produce\n"
"			incorrect FITS files. \n\n"
"		NOTE: fixfits automatically detects and fixes files produced\n"
"		by Registax and Astra Image, so no special options are\n"
"		required for it to handle those files.\n\n"
"    -stretch  (experimental)\n"
"			Stretches the range of values in the FIT file.\n"
"			Don't expect this to work yet :-)\n\n"

"Caveats: This program was written exclusively to handle FITS files produced\n"
" by Registax and Astra Image.\n\n"
"It may not work with FITS files from other sources as yet, since it's\n"
"FITS handling is not completely implemented.\n"
"This program is a work in progress. Use with care.\n\n"
	);
	}
 
int BZERO_ORIG = 0;
int BZERO_NEW = 0;
int set_bzero = 0;
int do_stretch = 0;
int force = 0;
int arg_padbytes = 0; 		// user-specified padding
int ADD_VALUE = 0;
double SCALE_VALUE = 1;
char *OUT_FNAME = NULL;
char *BLANK80 = "                                                                                          ";

int fits_header_extract_value(char *str);
void byteswap(unsigned char *buffer, int count, int bpp);
void bzero(char *buffer, int n);
int isFitsFilename(char *str);
int isDirectory(char *str);

// these definitions aren't used yet (work in progress)

#define UNSIGNED_16 	1
#define SIGNED_16 	2
#define UNSIGNED_32 	3
#define SIGNED_32 	4

typedef struct {
	double *data;
	int width, height;
	int type;
	int bpp;
	int bzero;
	char header[36][81];
	} fits;

int main(int argc, char *argv[])
	{
	int i;

	if (argc<2) {
	   Usage();
	   exit(0);
	   }

	for(i=1; i<argc; ++i) {
	   if (!strcmp(argv[i],"-add")) {
		ADD_VALUE = atoi(argv[++i]);
		}
	   else if (!strcmp(argv[i],"-scale")) {
		SCALE_VALUE = atof(argv[++i]);
		}
	   else if (!strcmp(argv[i],"-padbytes")) {
		arg_padbytes = atoi(argv[++i]);
		}
	   else if (!strcmp(argv[i],"-force")) {
		force = 1;
		}
	   else if (!strcmp(argv[i],"-stretch")) {
		do_stretch = 1;
		printf("[do_stretch] ");
		}
	   else if (!strcmp(argv[i],"-bzero")) {
		BZERO_NEW = atoi(argv[++i]);
		set_bzero = 1;
		}
	   else if (!strcmp(argv[i],"-out")) {
		OUT_FNAME = argv[++i];
		}
	   else if (isDirectory(argv[i])) {
		// It's a directory - look for FITS files inside it
		char fname[1024];
		struct dirent *e;
		DIR *d = opendir(argv[i]);
		int len;

		if (d==NULL) {
			fprintf(stderr,"Error opening directory %s\n",argv[i]);
			}
		else {
		   while(e = readdir(d)) {
			if (strlen(e->d_name)>4) {
			   sprintf(fname,"%s/%s",argv[i],e->d_name);
	   		   if (! access(fname,R_OK|F_OK) && isFitsFilename(e->d_name)) {
			      /* We have a xx.fit file to align */
			      Fixup(fname);
			      }
			   }
			}
		      closedir(d);
		      }
		   }
	   else {
		if (! access(argv[i],F_OK)) Fixup(argv[i]);
	   	else {
	   	   fprintf(stderr,"Cannot access file [%s]\n",argv[i]);
	   	   exit(1);
		   }
		}
	   }

	exit(0);
	}

int Fixup(char *fname)
	{
	int i,fd,bytes,bpp;
	int count=0,last_record=0;
	int found_data,pos,min,pad_bytes;
	int add_value = ADD_VALUE;
	char recs[36][81],rec[81];
	unsigned short *buffer_16;
	unsigned int   *buffer_32;
	unsigned char  *buffer,*chr_ptr,tmp;
	struct stat st;
	unsigned int max,hist[10];

	fprintf(stderr,"%s\n",fname);

	if (stat(fname,&st)) {
	   fprintf(stderr,"Error: cannot stat '%s'\n",fname);
	   return(0);
	   }

	// Start with this value, subtract 80 for each header line we see
	bytes = st.st_size;

	if (bytes < 0) {
	   fprintf(stderr,"Error: File '%s' has size too small.\n",fname);
	   return(0);
	   }

	/*
 	 * Open the file, seek past the header of 36 records of 80 chars
	 */
	fd = open(fname,O_RDWR|O_BINARY);

	/*
	 * read the header records
	 */
	for(i=0; i<36; ++i) {
	   if (read(fd,recs[i],80) != 80) {
		fprintf(stderr,"Error reading header line %d in %s\n",i+1,fname);
		close(fd); return(0);
		}

	   bytes -= 80;

	   if (!strncmp(recs[i],"BZERO",5)) {
		int n = fits_header_extract_value(recs[i]);

		fprintf(stderr,"\tBZERO %d\n",n);
		if (!force && (bpp == 16 && n != 32768)) {
		   fprintf(stderr,"\tLooks like this file does not need fixing. Skipping...\n");
		   close(fd); return(1);
		   }
		BZERO_ORIG = n;
		}

	   if (!strncmp(recs[i],"BITPIX",6)) {
		bpp = fits_header_extract_value(recs[i]);
		fprintf(stderr,"\tBPP %d\n",bpp);
		if (bpp != 16 && bpp != 32 && bpp != -32) {
		   fprintf(stderr,"\t**invalid BPP (support only 16 or 32) **\n");
		   close(fd);
		   return(1);
		   }
		}

	   if (! force && !strncmp(recs[i],"COMMENT =Mangled by FIXFITS",27)) {
		   fprintf(stderr,"\tLooks like this file has already been mangled. Skipping...\n");
		   close(fd); return(1);
		   }

	   if (!strncmp(recs[i],"COMMENT= REGISTAX",17)) {
		fprintf(stderr,"\t*Authored by: Registax.\n");
		if (add_value == 0) {
		   add_value = -1;
		   fprintf(stderr,"    ->  Adjusting pixel values by -1\n");
		   }
		}

	   if (!strncmp(recs[i],"ORIGIN  = 'Astra Image",22)) {
		// Astra Image is buggy and truncates its data instead of padding
		// to 2880 bytes. So to be on the safe side we must pad an extra 2880 bytes.
		if (!arg_padbytes) {
		   arg_padbytes = 2880;
		   fprintf(stderr,"\t*Authored by: Astra Image, padding extra 2880 bytes\n");
		   }
		}

	   if (!strncmp(recs[i],"END ",4)) {
		last_record = i;
		}
	   }

	/*
	 * check the file signature to be sure it's a FIT
	 */
	if (strncmp(recs[0],"SIMPLE  =                    T",30)) {
	   fprintf(stderr,"Invalit FIT signature at start of file '%s'\n",fname);
	   close(fd); return(0);
	   }

	/* Allocate buffer */
	buffer = (unsigned char *) malloc(bytes);
	if (buffer == NULL) {
	   fprintf(stderr,"Out of memory allocating %d bytes for file %s\n",bytes,fname);
	   close(fd); return(0);
	   }

	if (bpp == 16) 		buffer_16 = (unsigned short *) buffer;
	else if (bpp == 32)  	buffer_32 = (unsigned int *) buffer;

	/* read the data */
	if (read(fd,(char *)buffer,bytes) != bytes) {
	   fprintf(stderr,"Short read on file '%s', wanted %d bytes\n",fname,bytes);
	   close(fd); free(buffer);
	   return(0);
	   }

	/* Fix the byte order issue */
	byteswap(buffer,bytes,bpp);

	fprintf(stderr,"\t[ADD %d, SCALE %-2.2lf, BZERO %d]\n",
		add_value,SCALE_VALUE,set_bzero?BZERO_NEW:BZERO_ORIG);

	if (bpp == 16) {
	   max = buffer_16[0];
	   if (SCALE_VALUE != 1) for(i=0; i<bytes/2; ++i) {
	      buffer_16[i] += add_value; 
	      buffer_16[i] *= SCALE_VALUE;
	      if (buffer_16[i] > max) max = buffer_16[i];
	      }
	   else for(i=0; i<bytes/2; ++i) {
	      buffer_16[i] += add_value; 
	      if (buffer_16[i] > max) max = buffer_16[i];
	      }
	   }
	else if (bpp == 32) {
	   max = buffer_32[0];
	   if (SCALE_VALUE != 1) for(i=0; i<bytes/4; ++i) {
	      buffer_32[i] += add_value; 
	      buffer_32[i] *= SCALE_VALUE;
	      if (buffer_32[i] > max) max = buffer_32[i];
	      }
	   else for(i=0; i<bytes/4; ++i) {
	      buffer_32[i] += add_value; 
	      if (buffer_32[i] > max) max = buffer_32[i];
	      }
	   }

	if (do_stretch) {
	   double v;
	   if (bpp == 16) {
		v = 65535.0 / (double) max;
		for(i=0; i<bytes/2; ++i) buffer_16[i] *= v;
		}
	   else if (bpp == 32) {
		v = (double)0xffffffff / (double) max;
		for(i=0; i<bytes/4; ++i) buffer_32[i] *= v;
		}
	   }

	// fixup for unsigned 32 -> signed 32
	if (bpp == 32) {
	   for(i=0; i<bytes/4; ++i) buffer_32[i] -= 0x7fffffff;
	   }

	/* revert byte order to file format */
	byteswap(buffer,bytes,bpp);

	/* Maybe write out to a different file */
	if (OUT_FNAME) {
	   close(fd);
	   fd = open(OUT_FNAME,O_CREAT|O_WRONLY|O_BINARY);
	   if (fd<0) {
		fprintf(stderr,"Cannot open output file [%s]\n",OUT_FNAME);
		exit(1);
		}
	   }

	lseek(fd,0,SEEK_SET);

	// Add a comment field indicating that we have mangled the file
	// Make sure we have 36 records (0..35).
	if (last_record < 35) {
	   sprintf(recs[last_record++],"%-8.8s=%-71.71s",
		"COMMENT",
		"Mangled by FIXFITS (C)Copyright 2005 Anthony Wesley");

	   sprintf(recs[last_record++],"%-80.80s","END");
	   while(last_record < 36) sprintf(recs[last_record++],"%-80.80s","");
	   }
	
	/* Write the modified header */
	for(i=0; i<36; ++i) {
	   if (!strncmp(recs[i],"BZERO ",6)) {
		sprintf(recs[i],"%-8.8s=%21d%50s","BZERO",set_bzero?BZERO_NEW:BZERO_ORIG,"");
		if (write(fd,recs[i],80) != 80) {
		   fprintf(stderr,"Short write on file '%s': Probably damaged\n",fname);
		   close(fd); free(buffer);
		   return(0);
		   }
		}
	   else {
		if (write(fd,recs[i],80) != 80) {
		   fprintf(stderr,"Short write on file '%s': Probably damaged\n",fname);
		   close(fd); free(buffer);
		   return(0);
		   }
		}
	   }


	if (write(fd,(char *)buffer, bytes) != bytes) {
	   fprintf(stderr,"Short write on '%s': File probably damaged\n",fname);
	   close(fd); free(buffer); return(0);
	   }

	// Write out the user-specified padding bytes, make them 0
	if (arg_padbytes > 0) {
	   bzero(buffer,arg_padbytes);
	   if (write(fd,buffer,arg_padbytes) != arg_padbytes) {
	   	fprintf(stderr,"Short write on '%s': File probably damaged\n",fname);
	   	close(fd); free(buffer); return(0);
		}
	   bytes += arg_padbytes;
	   }

	/* Must pad data out to multiple of 2880 bytes */
	pad_bytes = 2880 - (bytes%2880);
	if (pad_bytes != 2880) {
	   char pad[2880];
	   if (write(fd,pad,pad_bytes) != pad_bytes) {
	      fprintf(stderr,"Short write on '%s': File probably damaged\n",fname);
	  	}
	   }
	
	close(fd);
	free(buffer);
	fprintf(stderr,"\tOK\n");
	return(1);
	}

int
fits_header_extract_value(char *str)
	{
	int k,n,j=0;
	char tmpbuf[80];

	// skip past keyword
	while(j<79 && str[j]!='=') ++j;

	if (str[j] != '=') {
		fprintf(stderr,"Format error parsing header line '%s'\n",str);
		return(-1);
		}

	++j;
	while(j<79 && str[j]==' ') ++j;

	k=0; while(j<79) tmpbuf[k++] = str[j++];
	tmpbuf[k]=0;
	return (atoi(tmpbuf));
	}

void
byteswap(unsigned char *buffer, int count, int bpp)
	{
	int i;
	unsigned char tmp;

	/* Fix the byte order issue */
	if (bpp == 16) {
	   for(i=0; i<count; i+=2) { tmp=buffer[i]; buffer[i]=buffer[i+1]; buffer[i+1]=tmp; }
	   }
	else if (bpp == 32) {
	   for(i=0; i<count; i+=4) { 
		// 3210 -> 0123
		tmp=buffer[i]; buffer[i]=buffer[i+3]; buffer[i+3]=tmp; 
		tmp=buffer[i+1]; buffer[i+1]=buffer[i+2]; buffer[i+2]=tmp; 

		// 1032 -> 0123 (NO)
		//tmp=buffer[i]; buffer[i]=buffer[i+1]; buffer[i+1]=tmp; 
		//tmp=buffer[i+2]; buffer[i+2]=buffer[i+3]; buffer[i+2]=tmp; 

		// 2301 -> 0123 (NO)
		//tmp=buffer[i]; buffer[i]=buffer[i+2]; buffer[i+2]=tmp; 
		//tmp=buffer[i+1]; buffer[i+1]=buffer[i+3]; buffer[i+3]=tmp; 
		}
	   }
	}

void bzero(char *buffer, int n) { while(n-->0) *(buffer++)=0; }

int
isDirectory(char *str)
	{
	struct stat st;
	int val;

	// Remove trailing "/" or "\"
	val = strlen(str) - 1;
	if (str[val]=='/' || str[val]=='\\') str[val]=0;

	if (stat(str,&st)) {
		fprintf(stderr,"Cannot stat '%s'\n",str);
		return(0);
		}

	val = (st.st_mode & S_IFDIR);
	return (val);
	}

// Return 1 if the argument is a string ending in ".fit" or ".fts"
int
isFitsFilename(char *str)
	{
	char fname[1024];

	// make a copy of the input string cause we're going to lowercase it.
	strncpy(fname,str,1023); fname[1023]=0;
	strlwr(fname);

	// seek to the end, then backup to find the suffix
	str = fname; while(*str && (str-fname < 1023)) str++;

	// filename too short.
	if (str-fname < 5) return 0;

	// rewind to where suffix should be
	str -= 4;
	if (*str != '.') return 0;

	if (!strcmp(str,".fit")) return (1);
	if (!strcmp(str,".fts")) return (1);

	return(0);
	}