#include "ninox.h"
#include <sys/stat.h>

static void Canonise_src_fname(char *src_fname);
static void Set_CurrentFile(char *src_fname);
static char * Generate_dst_fname(char *src_fname);
static int ProcessImage(struct Image *img);

static int AlignSingleFile(char *name);
static int AlignArchive(char *name);

static int AlignBMP(char *fname, char *out_fname);
static int AlignFIT(char *fname, char *out_fname);
int isBrokenFrame(struct Image *img);

int
Align(char *fname)
	{
	char *ptr = fname + strlen(fname)-1;
	int count=0;

	if (ImageCount == 0) return 0;

	if (StackCount>0 && StackCount>=StackMax) {
	   printf("stack count reached, processing stopped\n");
	   return 0;
	   }

	while(ptr>fname && *ptr != '.') --ptr;

	// Is it an archive FITS file ?
	if (!strcasecmp(ptr,".fta") || !strcasecmp(ptr,".ftz")) {
again:
	   count += AlignArchive(fname);

	   if (ChainArchives && ImageCount>0) {
		char nfname[256];
		int found_next=ArchiveNextFilename(fname,nfname);

	  	if (found_next) {
		   strcpy(fname,nfname);
		   goto again;
		   }
		}

	   return count;
	   }

	// Otherwise it must be a simple file
	--ImageCount;
	return AlignSingleFile(fname);
	}

static void
Canonise_src_fname(char *src_fname)
	{
	int i;

	// Canonise filename
	for(i=0; src_fname[i]; ++i)
	   if (src_fname[i]=='\\') src_fname[i]='/';

	// Get rid of leading "./"
	if (src_fname[0]=='.' && src_fname[1]=='/') {
	   for(i=2; src_fname[i]; ++i) src_fname[i-2]=src_fname[i];
	   src_fname[i-2]=0;
	   }

	}
	
static void
Set_CurrentFile(char *src_fname)
	{
	char *p;

	p = src_fname;
	while(*p) ++p; --p;
	while(p != src_fname && *(p-1) != '/') --p;
	CurrentFile = strdup(p);
	}

static char *
Generate_dst_fname(char *src_fname)
	{
	char *dst_fname = strdup(src_fname);
	
	// If -outdir option is given write the files to the named directory
	// instead of overwriting the original files. OutDir is relative to
	// the source file directory
	if (OutDir) {
	   char *strtmp = strdup(src_fname);
	   char *ptr = strtmp + strlen(strtmp) -1;
	   int num;
	   char sfx[10],type[10],*pfx_dir;

	   while(ptr != strtmp && *ptr != '/') --ptr;
	   if (ptr != strtmp) {
		*(ptr++)=0;
		pfx_dir = strtmp;
		}
	   else {
		pfx_dir=".";
		}

	   // Look for filename "dddddd-X.sss" d=digits, sss=suffix, X is
	   // identifier. If found, create subdirectory "X"
	   type[0] = 0;
	   if (SubDirs && sscanf(ptr,"%d-%1s.%s",&num,type,sfx) != 3) type[0]=0;

	   free(dst_fname);
	   dst_fname = malloc(strlen(pfx_dir) + strlen(OutDir) + 
			strlen(type) + strlen(ptr) + 4);
	   if (! dst_fname) {
		Print("Align: Out of Memory\n");
		exit(1);
		}

	    // If OutDir is absolute, then just use it, else if it is relative
	    // then prepend our directory prefix
	    if (OutDir[0]=='/' || OutDir[1]==':' || !strncmp(OutDir,"./",2))
		strcpy(dst_fname,OutDir);
	    else
	    	sprintf(dst_fname,"%s/%s",pfx_dir,OutDir);

	    if (! NoSave) Mkdir(dst_fname);

	     if (type[0]) {
		strcat(dst_fname,"/");
		strcat(dst_fname,type);
		if (! NoSave) Mkdir(dst_fname);
		if (forceWriteEmptyFiles) 
		     WriteEmptyFiles = 1;
		else WriteEmptyFiles=0;
		}

	   strcat(dst_fname,"/");
	   strcat(dst_fname,ptr);
	   free(strtmp);
	   }

	return dst_fname;
	}

// Return 1 if we should keep going, 0 if we should stop
static int
CheckRunfile()
	{
	static unsigned long last_runfile_check = 0;

	if (RunFile && time(0) - last_runfile_check > 0) {
	   struct stat st;

	   if (access(RunFile,0)) {
		if (!Quiet) Print("RunFile has been removed, exiting\n");
		return(0);
		}

	   while(1) {
	      stat(RunFile,&st);
	      int age = time(0) - st.st_mtime;
	      if (age <= 1) break;
	      if (age>=5) {
		 Print("Runfile is stale, quitting.\n");
	         unlink(RunFile);
	         return(0);
		 }
	      Usleep(100 * 1000);
	      }
	   }

	return (1);
	}

static int 
AlignSingleFile(char *src_fname)
	{
	char *p,*dst_fname;
	int i,rval=1;
	struct Image *img = NULL;
	static int _process_skip=9999;
	static int last_runfile_check = 0;

	if (! CheckRunfile()) { exit(1); }

	if (ProcessOffset>0) {--ProcessOffset; goto end;}
	if (_process_skip < ProcessSkip) {_process_skip++; goto end;}
	_process_skip=0;
   
	Canonise_src_fname(src_fname);
	dst_fname = Generate_dst_fname(src_fname);

	Set_CurrentFile(src_fname);

	img = LoadImage(src_fname,dst_fname);
	if (! img) {
	   Print("Error Loading '%s'\n",src_fname);
	   rval=0; goto end;
	   }


	rval = ProcessImage(img);  // img is destroyed in here

	end:

	free(dst_fname);
	return rval;
	}

// Open the named archive and load/process each of the files it contains.
// Return the number of files processed.
static int
AlignArchive(char *archive_name)
	{
	char *src_dir;
	char strtmp[128];
	char *p,*src_fname;
	int count,i,flags,err=1;
	struct Image *img;
	static int archive_num = 0;   // When overriding filenames in the archive
	static int last_runfile_check = 0;
        FILE *in = OpenArchive(archive_name);

        if (in == NULL) {
           Print("Cannot open archive '%s'\n",archive_name);
           perror(NULL);
           exit(1);
           }

	flags=0;
	if (! strcmp(archive_name+strlen(archive_name)-4,".ftz")) flags = ARCHIVE_COMPRESSED;

	// Re-Init this so we don't keep data from an unrelated source
	if (StreamFilter) init_streamfilter();

	Canonise_src_fname(archive_name);

	// Get the directory where the archive is located
	// If there's no directory then src_dir will be an empty string
	src_dir = strdup(archive_name);
	p = src_dir + strlen(src_dir)-1;
	while(p>src_dir && *p != '/') --p; *p=0;

	count=0;
	while(ImageCount>0) {
	   if (StackCount>0 && StackCount>=StackMax) break;

	   if (! CheckRunfile()) {exit(1);}

	   img = ArchiveLoadNextImage(in,&err,flags|ARCHIVE_LOAD);
	   if (err==0) break; // End of archive
	   if (err<0) {
		Print("fatal error while reading archive '%s'\n",archive_name);
		exit(1);
		}

	   ++count;

	   if (ArchiveOverrideFilenames) {
		char *n = Malloc(32);
		char *p1,*p2;

		// find the filename portions before and after the numeric section
		p1 = img->src_fname; while(*p1 && !isdigit(*p1)) ++p1;
		p2 = p1; while(isdigit(*p2)) ++p2;

		if (*p1==0 || *p2==0 || p1==p2) {
		   Print("ArchiveOverrideFilenames: Could not find numeric portion in [%s]\n",img->src_fname);
		   exit(1);
		   }

		Print("ArchiveOverrideFilename: Renaming %s => ",img->src_fname);
		*p1=0;
		sprintf(n,"%s%06d%s",img->src_fname,archive_num,p2);
		Print("%s\n",n);
		++archive_num;
		free(img->src_fname); img->src_fname = n;
		}

	   // Generate a pseudo-filename that reflects what this file would
	   // be called if it were not in the archive, by prepending the
	   // archive directory to the given filename
	   if (src_dir[0] && isDirectory(src_dir)) {
	      src_fname = Malloc(strlen(src_dir) + strlen(img->src_fname) + 2);
	      sprintf(src_fname,"%s/%s",src_dir,img->src_fname);
	      }
	   else src_fname = strdup(img->src_fname);

	   free(img->src_fname);
	   img->src_fname = src_fname;

	   // now that we have a pseudo filename we can generate
	   // a destination
	   img->dst_fname = Generate_dst_fname(src_fname);

	   p = src_fname + strlen(src_fname)-1;
	   while(p>src_fname && *(p-1) != '/') --p;
	   sprintf(strtmp,"%s:%s",archive_name,p);
	   Set_CurrentFile(strtmp);

	   if (ProcessOffset>0) {
		--ProcessOffset;
		 DestroyImage(img);
		 free(CurrentFile); CurrentFile=NULL;
		goto end;
		}

	   // process the image and destroy it
	   ProcessImage(img);
	   if (ImageCount < IMAGECOUNT_MAX) --ImageCount;

	   end:
   
	   for(i=0; ProcessOffset==0 && i<ProcessSkip; ++i)
		ArchiveLoadNextImage(in,&err,flags|ARCHIVE_SKIP);
	   }

	fclose(in);
	return count;
	}

static int
Hist(struct Image *img)
	{
	int width = img->width;
	int height = img->height;
	int h[65536];
	int npix = width * height;
	unsigned short *data = (unsigned short *)img->data;
	int o;

	
	for(o=0; o<65536; ++o) h[o]=0;
	for(o=0; o<npix; ++o) h[data[o]>>4]++;

	for(o=0; o<65536; ++o)
	   printf("%d: %d\n",o,h[o]);

	exit(0);
	}

static double
do_sum_flux(struct Image *img)
	{
	int width = img->width;
	int height = img->height;
	int depth = img->depth;
	int i,npix = img->width * img->height;
	unsigned char *ptr = (unsigned char *) img->data;
	unsigned short *uptr = (unsigned short *) img->data;

	double total =0;

	switch(depth) {
	   case 8:
		for(i=2; i<npix; ++i) total += ptr[i];
		break;
	   case 16:
		for(i=2; i<npix; ++i) {
		   unsigned int v = uptr[i]>>8;
		   if (v >= ThreshHold) total += v;
		   }
		break;
	   case 24:
		for(i=9; i<npix; i+=3) {
		   double v = uptr[i]*0.29 + uptr[i+1]*0.58 + uptr[i+2]*0.29;
		   total += v;
		   }
		break;
	   default:
		Print("do_sum_flux: depth %d not supported\n",depth);
		exit(1);
		break;
	   }

	return total;
	}

static int
ProcessImage(struct Image *img)
	{
	static int first=1;
	static int last_runfile_check = 0;
	int rval=1;

	if (first) {
	   if (UpScale != DownScale) {
	      double scalef = (double)UpScale / (double) DownScale;
              Print("Rescaling by %-2.2f, smoothing is %s\n",scalef,UpScale_Smoothing?"on":"off");
	      }

	   }

	/*
	 * Check if we're counting the flux in this image
	 */
	if (SumFlux) {
	   double flux = do_sum_flux(img);
	   static int FCount=0;
	   FILE *z;

	   if (first) {
		Print("Appending image flux to '%s'\n",FluxFile);
		unlink(FluxFile);
		}

	   z = fopen(FluxFile,"a");
	   if (! z) {
		Print("Cannot open Flux file '%s'\n",FluxFile);
		exit(1);
		}
	   else {
		double fl = img->tm_sec;
		int hr = fl/3600;
		int min = (fl - hr*3600) / 60;
		double secs = fl - hr*3600 - min*60;

		fprintf(z,"%-5d,\"%s\",\"%02d:%02d:%2.6lf\",\"%lf\"\n",
			++FCount,img->src_fname,hr,min,secs,flux);
		fclose(z);
		}
	   }

	/*
	 * De-bayer if needed
	 */
	if (DoDeBayer) {
	   struct Image *new = DeBayer(img,DoDeBayer);
	   DestroyImage(img);
	   img = new;
	   Print("Done Debayer\n");
	   }

/*
	Hist(img);
*/

	/*
	 * de-protect (remove magic pixels 0 and 1 
	 */
	if (img->depth==8)
	   img->data[0] = img->data[1]=0;
	else if (img->depth == 16) {
	   unsigned short *udata = (unsigned short *)img->data;
	   udata[0]=udata[1]=0;
	   }
	
	if (DarkFrame != NULL) {
	   SubtractDarkFrame(img);
	   }

	/*
	 * Test mode: Subsample input image and write out
	 */
	if (SubsampleMode) {
	   Test_Subsample(img);
	   exit(0);
	   }

	/* levels adjust on input data */
	if (doLevelsAdjust && !LevelsAdjust(img,LevelsMin,LevelsMax)) {
	   printf("Levels adjust failed\n");
	   exit(1);
	   }

	/* If requested, debayer the image */
	if (Bayer8) FixBayer8(img);

	// Detect Broken Frames
	if (DetectBrokenFrames && isBrokenFrame(img)) {
	   printf("Frame is broken, skipping\n");
	   rval=1; goto end;
	   }
	   
	// Detect hot pixels
	if (DetectHotPixels) {
	   do_DetectHotPixels(img);
	   }

       // Gain compensation
       if (ApplyGainComp)
          ApplyGainCompensation(img);

	// POP filter
	if (PopFilter && PopFilter_When == POPFILTER_BEFORE)
	   pop_filter(img);

	if (DoFFT) {
	   static struct Image *ref = NULL;
	   int dx,dy;

	   if (! ref) ref = LoadImage(FFT_Ref,"FFT REF");

	   if (! ref) {
		Print("DoFFT: Can't load reference image '%s'\n",FFT_Ref);
		exit(0);
		}

	   if (! fft_register_pad(ref,img,FFT_LPF_Radius1,FFT_LPF_Radius2,FFT_LPF_Pow,FFT_Region,&dx,&dy)) {
		Print("fft_register: failed\n");
		exit(0);
		}
	   Print("Done FFT, offset (%d,%d)\n",dx,dy);
	   }

	if (! process(img)) {
	   Print("Image processing failed\n");

	   // Make sure the output file exists even if it is size 0
	   if (access(img->dst_fname,0) && WriteEmptyFiles && NoSave==0) {
		FILE *out = fopen(img->dst_fname,"wb");
		fclose(out);
		rval=1; goto end;
		}
	   rval=1; goto end;
	   }

	if (PopFilter && PopFilter_When == POPFILTER_AFTER)
	   pop_filter(img);

        if (UpScale_Smoothing && UpScale>1 && UpScale_Smoothing_When == SMOOTHING_AFTER)
             smooth_image(img,UpScale);

	if (FFT_Filter) fft_filter(img,FFT_LPF_Radius1,FFT_LPF_Radius2,FFT_LPF_Pow);

	if (OutputFileType>0 || OutputFileDepth != -1) {
	   int newtype = OutputFileType;
	   int newdepth = OutputFileDepth;

	   if (newtype<0) newtype = img->type;
	   if (newdepth == -1) newdepth = img->depth;

	   img = ConvertToType(img,newtype,newdepth);
	   if (img==NULL) {
		Print("Cannot convert image '%s'(%s/%dbpp) to (%s/%dbpp)\n",
				img->src_fname,TypeName(img->type),img->depth,TypeName(newtype),newdepth);
		exit(1);
		}
	   }

	if (QEstimator) {
	   if (!Quiet) Print("Quality: ");
	   double q = QualityEstimate(img);
	   if (!Quiet) Print(" = %-2.2lf\n",q);
	   QAssociate(img->dst_fname,q);
	   }

	if (NoSave==0) {

	   if (!Quiet) Print("-> %s (%dx%dx%d)\n",img->dst_fname,img->dst_width,img->dst_height,img->depth);

	   if (! AllowWrite(img->dst_fname)) {
	         if (!Quiet) Print("Not overwriting output file %s\n",img->dst_fname);
	         }
	   else if (!WriteImage(img)) {
	      Print("Write on image '%s' failed\n",img->dst_fname);
	      exit(1);
	      }
	   }
	else {
	   Print("Not saved\n");
	   }

         // Possibly stack the input frame
       	if (StackFile && StackCount < StackMax) {
              stack_frame(img);
              StackCount++;
              if (!Quiet) Print("[Stack %d] ",StackCount);
              }

	// This has to be the last thing we do with the image since we're going to
	// draw all over it with text etc.

	if (DisplayFrames) {
	   ShowImage(img);
	   }

end:

	DestroyImage(img);
	free(CurrentFile); CurrentFile=NULL;
	Print("\n");

	first=0;
	return rval;
	}

// Detect "broken frames", i.e. frames where one or more packets of data were lost
// during capture, so part of the image is horizontally shifted.
int
isBrokenFrame(struct Image *img)
	{
	int x,y,i,total,avg,count;
	int w = img->width;
	int h = img->height;
	int depth = img->depth;
	unsigned char *data = (unsigned char *)img->data;
	unsigned short *udata = (unsigned short *)img->data;
	int left_edge,right_edge,broken;
	int *rowdiffs;

	rowdiffs = ZeroMalloc(h * sizeof(int));

	// For planets, if they are off the screen then this is broken
	left_edge=right_edge=broken=0;

	for(y=1; y<h-1; ++y) {
	   int diff,diff_above=0;
	   int diff_below=0;
	   int o = y * w;

	   // difference between this line and line below and above
	   if (depth==8) {

		if (DBF_TYPE == DBF_PLANET) {
		   if (data[o] > ThreshHold) ++left_edge;
		   if (data[o+w-1] > ThreshHold) ++right_edge;
		   }

		for(x=0; x<w; ++x,++o) {
		   diff = data[o]; diff -= data[o+w]; if (diff<0) diff=-diff;
		   diff_below += diff;
		   diff = data[o]; diff -= data[o-w]; if (diff<0) diff=-diff;
		   diff_above += diff;
		   }
		}
	   else if (depth==16) {

		if (DBF_TYPE == DBF_PLANET) {
		   if (udata[o] > (ThreshHold<<8)) ++left_edge;
		   if (udata[o+w-1] > (ThreshHold<<8)) ++right_edge;
		   }

		for(x=0; x<w; ++x,++o) {
		   diff = udata[o]; diff -= udata[o+w]; if (diff<0) diff=-diff;
		   diff_below += diff>>8;
		   diff = udata[o]; diff -= udata[o-w]; if (diff<0) diff=-diff;
		   diff_above += diff>>8;
		   }
		}
	   else if (depth==24) {
		for(x=0,o*=3; x<w; ++x,o+=3) { // BGR
		   int v_o =  (double)data[o]     * 0.114 + (double)data[o+1]    *0.587 + (double)data[o+2]    *0.299;
		   int v_oa = (double)data[o-w*3] * 0.114 + (double)data[o-w*3+1]*0.587 + (double)data[o-w*3+2]*0.299;
		   int v_ob = (double)data[o+w*3] * 0.114 + (double)data[o+w*3+1]*0.587 + (double)data[o+w*3+2]*0.299;

		   diff = v_o - v_ob; if (diff<0) diff=-diff;
		   diff_below += diff;
		   diff = v_o - v_oa; if (diff<0) diff=-diff;
		   diff_above += diff;
		   }
		}

	   diff = diff_above - diff_below;
	   if (diff<0) diff=-diff;

	   rowdiffs[y] = diff;
	   }

	if (DBF_TYPE == DBF_PLANET && (left_edge >= DBF_EDGE_COUNT || right_edge >= DBF_EDGE_COUNT)) {
	   if (!Quiet) printf("Image laterally displaced, probably broken.\n");
	   broken=1;
	   }

	// Scan through the row differences and look for 2 sequential rows that are more than
	// 10x different from the average
	total = rowdiffs[1];
	count=0;
	for(y=2; y<h-2; ++y) {
	   avg = total/y;
	   if (rowdiffs[y] > avg*DBF_DIFF && y > 4) { total += avg; ++count; }
	   else total += rowdiffs[y];

	   //printf("%d:%d avg=%d count=%d\n",y,rowdiffs[y],avg,count);
	   }

//printf("\n");

	if (broken || count >= 2) {
		FILE *z = OpenLogfile("brokenframes.txt","a");
		if (!z) printf("Cannot open file\n");
		else {
		   fprintf(z,"%s\n",img->src_fname);
		   fclose(z);
		   }
		if (!Quiet) printf("%s: broken frame detected!\n",img->src_fname);
		free(rowdiffs);
		return 1;
		}

	free(rowdiffs);
	return 0;
	}

int
SubtractDarkFrame(struct Image *img)
	{
	static struct Image *df = NULL;
	int w = img->width;
	int h = img->height;
	int d = img->depth;
	int npix = w * h;
	int o;
	unsigned char *data = (unsigned char *)img->data;
	unsigned short *udata = (unsigned short *)img->data;
	unsigned char *df_data;
	unsigned short *df_udata;

	if (df == NULL) {
	   df = LoadImage(DarkFrame,NULL);
	   if (! df) {
	  	printf("Error loading darkframe %s\n",DarkFrame);
	 	exit(1);
		}
	   }

	if (w != df->width || h != df->height || d != df->depth) {
	   printf("DarkFrame %s has different WxHxD to image\n");
	   exit(1);
	   }

	df_data = (unsigned char *)df->data;
	df_udata = (unsigned short *)df->data;

	switch(d) {
	   case 8:
		for(o=0; o<npix; ++o) {
		   int val = data[o]; val -= df_data[o] * DarkFrame_Scale;
		   if (val<0) val=0;
		   data[o] = val;
		   }
		break;
	   case 16:
		for(o=0; o<npix; ++o) {
		   int val = udata[o]; val -= df_udata[o] * DarkFrame_Scale;
		   if (val<0) val=0;
		   udata[o] = val;
		   }
		break;
	   default:
		printf("DarkFrame subtraction not supported for depth %d\n",d);
		exit(1);
		break;
	   }

	return 0;
	}