#include "ninox.h"

static double * Comp;
static struct Image *Ref;

// Support for detecting and masking out dead pixels
#define MAX_DEAD_PIXELS 128
static int DeadPixels = 0;
static int DeadPixelList_X[MAX_DEAD_PIXELS];
static int DeadPixelList_Y[MAX_DEAD_PIXELS];

int
LoadGainCompRef(char *fname)
	{
	int npix,i;
	unsigned short *uptr;
	double total=0,avg;
	struct Image *tmp;

	// Check for both Windows and Linux paths
	if (! isFile(fname)) {
	   // Maybe Linux path on Windows?
	   Print("Gaincomp image '%s' not found\n",fname);

	   if (fname[0]=='/' && fname[2]=='/' && isalpha(fname[1])) {
		   fname[0]=fname[1];
		   fname[1]=':';
		   }
	   else if (isalpha(fname[0]) && fname[1]==':') {
		   fname[1]=fname[0];
		   fname[0]='/';
		   }
	   Print("Trying alternate path '%s'\n",fname);
	   if (! isFile(fname)) {
		Print("Gaincomp image '%s' not found\n",fname);
		exit(1);
		}
	   }

	tmp = LoadImage(fname,"reference");
	if (! tmp) {
	   Print("Ref: Error loading reference image '%s'\n",fname);
	   exit(1);
	   }

	// Make sure it's 16bpp monochrome
	Ref = ConvertImage(tmp,IMG_FIT,16);

//ShowImage(Ref);
//Usleep(30 * 1000 * 1000);
//exit(1);

	Print("Ref: dimensions %dx%dx%d ",Ref->width,Ref->height,Ref->depth);
	npix = Ref->width * Ref->height;

	if (! npix) {
	   Print("Ref: Error loading reference image '%s': no pixels?\n",fname);
	   exit(1);
	   }

	uptr = (unsigned short *) Ref->data;

	if (Comp) free(Comp);
	Comp = Malloc(npix * sizeof(double));

	DeadPixels = 0;
	for(i=total=0; i<npix; ++i) {
	   if (uptr[i] >= 65530 && DeadPixels < MAX_DEAD_PIXELS) {
		int y = i / Ref->width;
		int x = i % Ref->width;

		DeadPixelList_X[DeadPixels] = x;
		DeadPixelList_Y[DeadPixels++] = y;
		}
	   else total += uptr[i];
	   }

	Print("Ref: Found %d dead pixels (max value hot pixels)\n",DeadPixels);

	avg = total / (npix - DeadPixels);
	Print("Average: %f\n",avg);

	// need at least 10/16 bits of accuracy
	if (avg < 1024) {
	   Print("Ref: Error, 16 bit average of reference image (%-4.2f) too small\n",avg);
	   exit(1);
	   }

	for(i=0; i<npix; ++i) {
	   double v = (double)uptr[i];

	   if (v>0) Comp[i] = avg/v; else Comp[i]=1;
	   Comp[i] *= GaincompBoost;
	   }

	Print("Ref: Loaded %d entries from '%s'\n",npix,fname);

	return(npix);
	}

int
ApplyGainCompensation(struct Image *img)
	{
	int dc,i;
	unsigned char *ptr = (unsigned char *)img->data;
	unsigned short *uptr = (unsigned short *)img->data;
	unsigned short *ruptr = (unsigned short *)Ref->data;
	double val;
	int w = img->width;
	int h = img->height;
	int rw = Ref->width;
	int rh = Ref->height;
	int npix = w * h;
	int x1 = 0;		// offset to top left corner in gaincomp reference
	int y1 = 0;		// offset to top left corner in gaincomp reference
	int x,y,o,ro;

	if (rw == 0 || rh == 0) {
	   Print("ApplyGainComp: Error: no reference loaded\n");
	   return 0;
	   }

	if (w > rw || h > rh) {
	   Print("ApplyGainCompensation: Image dimensions (%d x %d) are larger than reference (%d x %d)\n", w,h,rw,rh);
	   return(0);
	   }

	// Now we know that our image fits inside the refence frame, so look for xpos/ypos
	// offset in the image in case we are using ROI
	if (img->xpos>0) x1 = img->xpos;
	if (img->ypos>0) y1 = img->ypos;

	// Make sure we stay inside the reference image
	if (x1 + w > rw || y1 + h > rh) {
	   Print("ApplyGainCompensation: computed window (%d,%d) - (%d,%d) outside reference frame dimensions (%d x %d)\n",
		x1,y1,x1+w-1,y1+h-1,rw,rh);
	   return 0;
	   }

	Print("Gaincomp: using window (%d,%d) - (%d,%d) in reference\n",x1,y1,x1+w-1,y1+h-1);

	// Gain Compensation
	switch(img->depth) {
	   case 8:
		for(y=o=0; y<h; ++y) {
		   ro = rw * y1 + x1;
		   for(x=0; x<w; ++x,++ro,++o) {
			val = ptr[o];
			val *= Comp[ro]; val += 0.5;
			if (val<0) val=0; if (val>255) val=255;
			ptr[o] = val;
			}
		   }

		break;

	   case 16:
		for(y=o=0; y<h; ++y) {
		   ro = rw * y1 + x1;
		   for(x=0; x<w; ++x,++ro,++o) if (uptr[o]) {
			val = uptr[o];
			val *= Comp[ro]; val += 0.5;
			if (val<0) val=0; if (val>65535) val=65535;
			uptr[o] = val;
			}
		   }

		break;

	   default:
		Print("ApplyGainCompensation: Unsupported depth %d\n",img->depth);
		return 0;
	   }

	// De-Ring
	if (GaincompDeRingDelay) {
	   int o,x,y;

	   // Starting offset on each row, right hand edge
	   int x1 = w - GaincompDeRingDelay - 2;

	   switch(img->depth) {
	   	case 8:
		   for(y=0; y<h; ++y) {
			o=y * w + x1;  // right -> left
		 	for(x=x1; x>=0; --x,--o) {
				int delta = ptr[o+1]; delta -= ptr[o]; delta *= GaincompDeRingAmount;
				if (delta>0) {
				   double v = ptr[o+GaincompDeRingDelay];
				   v -= delta;
				   if (v<0) v=0; if (v>255) v=255;
				   ptr[o+GaincompDeRingDelay] = (unsigned int) v;
				   }
				}
			}
		   break;

	   	case 16:

		   for(y=0; y<h; ++y) {
			o=y * w + x1;  // right -> left
		 	for(x=x1; x>=0; --x,--o) {
				int delta = uptr[o+1]; delta -= (int)uptr[o];
				int am = delta * GaincompDeRingAmount;
				int am2 = am * GaincompDeRingAmount2;

				if (am>0) {
				   double v = uptr[o+GaincompDeRingDelay]; v -= am;
				   if (v<0) v=0; if (v>65535) v=65535;
				   uptr[o+GaincompDeRingDelay] = (unsigned short) v;
				   }

				if (am2>0) {
				   double v = uptr[o+GaincompDeRingDelay-1]; v -= am2;
				   if (v<0) v=0; if (v>65535) v=65535;
				   uptr[o+GaincompDeRingDelay-1] = (unsigned short) v;

				   //v = uptr[o+GaincompDeRingDelay+1]; v -= am2;
				   //if (v<0) v=0; if (v>65535) v=65535;
				   //uptr[o+GaincompDeRingDelay+1] = (unsigned short) v;
				   }
				}
			}
		break;
	      	}
	      }

	// Dead Column mapping
	switch(img->depth) {
	   case 8:
		for(dc=0; dc<GaincompDeadColCount; ++dc) {
			int o = GaincompDeadColList[dc];
			double v;

			for(i=0; i<h; ++i, o+=w) {
				v = (double)ptr[o-1]; v += (double)ptr[o+1];
				v/=2;
				if (v>255) v=255;
				ptr[o] = (unsigned short) v;
				}
			}
		   break;

	   case 16:
		for(dc=0; dc<GaincompDeadColCount; ++dc) {
			int o = GaincompDeadColList[dc];
			double v;

			for(i=0; i<h; ++i, o+=w) {
				v = (double)uptr[o-1]; v += (double)uptr[o+1];
				v/=2;
				if (v>65535) v=65535;
				uptr[o] = (unsigned short) v;
				}
			}
		break;
	   }

	// Dead pixel mapping from gaincomp reference frame, make sure we account for ROI
	// by subtracting off (x1,y1) which is the ROI top left corner.
	switch(img->depth) {
	   case 8:
		for(i=0; i< DeadPixels; ++i) {
		   int x = DeadPixelList_X[i] - x1;
		   int y = DeadPixelList_Y[i] - y1;
		   double t;
		   if (x>0 && y>0 && x < w-1 && y < h-1) {
			int n = y * w + x;
		   	t = ptr[n-w-1]; t += ptr[n-w]; t += ptr[n-w+1];
		   	t += ptr[n-1]; t += ptr[n+1]; t += ptr[n+w-1];
		   	t += ptr[n+w]; t += ptr[n+w+1];
		   	t /= 8; if (t<0) t=0; if (t>255) t=255;
		   	ptr[n] = (int) t;
		   	}
		   }
		break;

	   case 16:
		for(i=0; i< DeadPixels; ++i) {
		   int x = DeadPixelList_X[i] - x1;
		   int y = DeadPixelList_Y[i] - y1;
		   double t;
		   if (x>0 && y>0 && x < w-1 && y < h-1) {
			int n = y * w + x;
		   	t = uptr[n-w-1]; t += uptr[n-w]; t += uptr[n-w+1];
		   	t += uptr[n-1]; t += uptr[n+1]; t += uptr[n+w-1];
		   	t += uptr[n+w]; t += uptr[n+w+1];
		   	t /= 8; if (t<0) t=0; if (t>65535) t=65535;
		   	uptr[n] = (unsigned short)(t/8);
		   	}
		   }
		break;
		}

	return(1);
	}
