#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)
	{
	FILE *z;
	int npix,i;
	unsigned char *ptr;
	unsigned short *uptr;
	double total=0,avg;

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

	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);
	   }

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

	if (Comp) free(Comp);
	Comp = Malloc(npix * sizeof(double));
	switch(Ref->depth) {
	   case 8:
		for(i=total=0; i<npix; ++i) {
		   total += ptr[i];
		   if (ptr[i] >= 255 && DeadPixels < MAX_DEAD_PIXELS) DeadPixelList[DeadPixels++] = i;
		   }  
		break;
	   case 16:
		for(i=total=0; i<npix; ++i) {
		   total += uptr[i];
		   if (uptr[i] >= 65535 && DeadPixels < MAX_DEAD_PIXELS) DeadPixelList[DeadPixels++] = i;
		   }
		break;
	   default:
		Print("LoadGainCompRef: Unsupported depth %d for image '%s'\n",Ref->depth,fname);
		exit(1);
		break;
	   }

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

	switch(Ref->depth) {
	   case 8:
		for(i=0; i<npix; ++i) {
		   double v = (double)ptr[i];
		   if (v>0 && avg>0) Comp[i] = avg/v; else Comp[i]=1;
		   }
		break;
	   case 16:
		for(i=0; i<npix; ++i) {
		   double v = (double)uptr[i];
		   if (v>0 && avg>0) Comp[i] = avg/v; else Comp[i]=1;
		   }
		break;
	   }

	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);
	}
