#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[MAX_DEAD_PIXELS];

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

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


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

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

	if (! npix) {
	   Print("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) DeadPixelList[DeadPixels++] = i;
	   else total += uptr[i];
	   }

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

	// need at least 10/16 bits of accuracy
	if (avg < 1024) {
	   Print("LoadGainCompRef: Error, 16 bit average of reference image (%-4.2lf) 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;
	   }

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

	return(npix);
	}

int
ApplyGainCompensation(struct Image *img)
	{
	int i,npix;
	unsigned char *ptr = (unsigned char *)img->data;
	unsigned short *uptr = (unsigned short *)img->data;
	double val;

	if (Ref->width == 0 || Ref->height == 0) {
	   Print("ApplyGainComp: Error: no reference loaded\n");
	   return 0;
	   }

	if (img->width != Ref->width || img->height != Ref->height) {
	   Print("ApplyGainCompensation: Image has different dimensions to reference\n");
	   return(0);
	   }

	npix = Ref->width * Ref->height;

	for(i=0; i< DeadPixels; ++i) {
		int n = DeadPixelList[i];
		int w = img->width;
		int y = n / w;
		int x = n - y * w;
		int t=0;
		if (x>0 && y>0 && x < Ref->width && y < Ref->height) {
		   if (img->depth==8) t += ptr[n-w-1]; else t += uptr[n-w-1];
		   if (img->depth==8) t += ptr[n-w];   else t += uptr[n-w];
		   if (img->depth==8) t += ptr[n-w+1]; else t += uptr[n-w+1];
		   if (img->depth==8) t += ptr[n-1];   else t += uptr[n-1];
		   if (img->depth==8) t += ptr[n+1];   else t += uptr[n+1];
		   if (img->depth==8) t += ptr[n+w-1]; else t += uptr[n+w-1];
		   if (img->depth==8) t += ptr[n+w];   else t += uptr[n+w];
		   if (img->depth==8) t += ptr[n+w+1]; else t += uptr[n+w+1];

		   if (img->depth==8) ptr[n] = t/8; else uptr[n] = t/8;
		   }
		}

	switch(img->depth) {
	   case 8:
		for(i=0; i<npix; ++i) if (ptr[i]) {
			val = ptr[i];
			val *= Comp[i]; val += 0.5;
			if (val<0) val=0; if (val>255) val=255;
			ptr[i] = val;
			}
		break;
	   case 16:
		for(i=0; i<npix; ++i) if (uptr[i]) {
			val = uptr[i];
			val *= Comp[i]; val += 0.5;
			if (val<0) val=0; if (val>65535) val=65535;
			uptr[i] = val;
			}
		break;
	   default:
		Print("ApplyGainCompensation: Unsupported depth %d\n",img->depth);
		return 0;
	   }

	return(1);
	}