#include "ninox.h"

static int fetch_stack_u32(int ch, u32 *out);
static int fetch_stack_u16(int ch, unsigned short *out);
static int _write_stack_file(char *fname, int channel);

// How many channels do we support
#define NCHANNELS 32

static double *Data[NCHANNELS];		// Data array
static int    DCount[NCHANNELS];	// Count of active pixels

static int Frames =0;
static int NPixels=0;
static int Width,Height,Depth;
static int Dstwidth, Dstheight;

int stack_max_channels(void) { return NCHANNELS; }

static int
FrameChannel(char *fname)
        {
        char *ptr = fname + strlen(fname)-1;
        int ch = atoi(ptr-4);

        if (ch<0 || ch>6) ch = 0;

        return ch;
        }

int
stack_channel_has_data(int ch)
	{
	int i;

	if (ch<0 || ch >= NCHANNELS) {
	   Print("stack_channel_has_data: channel %d out of range 0 .. %d\n",ch,NCHANNELS);
	   exit(1);
	   }

	if (Data[ch] == NULL) return 0;
	return DCount[ch];
	}

int
stack_reset(void)
	{
	int i;

	for(i=0; i<NCHANNELS; ++i) {
	   if (Data[i]) {free(Data[i]); Data[i] = NULL;}
	   DCount[i] = 0;
	   }

	NPixels=Frames=0;
	Width = Height = Depth = 0;
	}

// must call this before using any of the routines in this file
int
stack_init(void)
	{
	int i;

	for(i=0; i<NCHANNELS; ++i) {
	   Data[i] = NULL;
	   DCount[i] = 0;
	   }

	return stack_reset();
	}

// channel is an integer, used to select the colour
// 2=red 3=green 4=blue
int
stack_frame(struct Image *img)
	{
	int width = img->width;
	int height = img->height;
	int depth = img->depth;
	unsigned char *data = img->data;
	int i,b,g,r,o;
	int threshhold = 1;  // anything not exactly black is ok
	unsigned short *sdata = (unsigned short *)data;
	int channel;
	double *dptr;
	double v;

	if (depth != 16 && depth != 8 && depth != 24) {
	   Print("-stack only supported for 8/16/24 bpp data, not for given value %d\n",depth);
	   return 0;
	   }

	channel = FrameChannel(img->src_fname);
	Print("stack to channel %d from %s\n",channel,img->src_fname);

	if (Width && Width != width) { Print("stack_frame: width has changed from %d to %d\n",Width,width); exit(0); }
	Width = width;

	if (Height && Height != height) { Print("stack_frame: height has changed from %d to %d\n",Height,height); exit(0); }
	Height = height;

	if (Depth && Depth != depth) { Print("stack_frame: depth has changed from %d to %d\n",Depth,depth); exit(0); }
	Depth = depth;

	NPixels = width * height;
	Dstwidth = img->dst_width;
	Dstheight = img->dst_height;

	if (Data[channel] == NULL) {
	   // First frame, allocate buffer
	   Data[channel] = (double *)ZeroMalloc(sizeof(double) * NPixels);
	   DCount[channel] = 0;
	   }

	dptr = Data[channel];

	// Stack the frame
	switch(depth) {
	   case 8:
		for(i=0; i<NPixels; ++i) {
	   	   	v = data[i];
	   	      	dptr[i] += v * 256 ; // upgrade 8 -> 16 bpp
			}
		DCount[channel]++;
		break;
	   case 16:
		threshhold <<= 8;
		for(i=0; i<NPixels; ++i) {
	   	   	v = sdata[i];
	   	      	dptr[i] += v;	// 16bpp data
			}
		DCount[channel]++;
		break;
	   case 24:
		// Use channels 2=red, 3=green, 4=blue
		if (Data[2] == NULL) {
	   	   Data[2] = (double *)ZeroMalloc(sizeof(double) * NPixels);
	   	   DCount[2] = 0;
		   }
		if (Data[3] == NULL) {
	   	   Data[3] = (double *)ZeroMalloc(sizeof(double) * NPixels);
	   	   DCount[3] = 0;
		   }
		if (Data[4] == NULL) {
	   	   Data[4] = (double *)ZeroMalloc(sizeof(double) * NPixels);
	   	   DCount[4] = 0;
		   }

		for(i=o=0; i<NPixels; ++i,o+=3) {
	   	   b = data[o]; g = data[o+1]; r = data[o+2];
	   	   Data[2][i] += r*256;	// 8bpp data, R channel, upgrade to 16bpp
	   	   Data[3][i] += g*256;	// 8bpp data, G channel, upgrade to 16bpp
	   	   Data[4][i] += b*256;	// 8bpp data, B channel, upgrade to 16bpp
	   	   }

		DCount[2]++; DCount[3]++; DCount[4]++;

		break;
	   default:
		Print("Stack: Depth %d unsupported\n",depth);
		break;
	   }
	return ++Frames;
	}

static int
fetch_stack_u32(int ch, u32 *out)
	{
	int i;
	u32 v;
	double *ptr;
	u32 count;

	if (! stack_channel_has_data(ch)) {
	   Print("fetch_stack_u32: Channel %d has no data\n",ch);
	   return 0;
	   }

	if (! Frames) {
	   Print("fetch_stack_u32: no frames stacked\n");
	   return 0;
	   }

	Print("u32 Averaging %d frames, channel=%d\n",Frames,ch);
	ptr = Data[ch];
	count = DCount[ch]; if (count==0) count=1;

	for(i=0; i<NPixels; ++i) out[i] = (ptr[i] * 65536.0) / count;

	return 1;
	}

static int
fetch_stack_u16(int ch, unsigned short *out)
	{
	int i;
	unsigned int v;
	double *ptr;
	u32 count;

	if (! stack_channel_has_data(ch)) {
	   Print("fetch_stack_u16: Channel %d has no data\n",ch);
	   return 0;
	   }

	if (! Frames) {
	   Print("fetch_stack_u16: no frames stacked\n");
	   return 0;
	   }

	Print("u16 Averaging %d frames\n",Frames);
	ptr = Data[ch];
	count = DCount[ch]; if (count==0) count=1;

	for(i=0; i<NPixels; ++i) out[i] = ptr[i] / count;

	return 1;
	}

// Write out one file for each active channel
// with suffix "-2" "-3" "-4" etc

int
write_stack_data(char *fname)
	{
	int i;
	int ch;

	//Print("Write_stack_data '%s'\n",fname);

	for(ch=0; ch<NCHANNELS; ++ch) 
	   if (stack_channel_has_data(ch)) {
	   	_write_stack_file(fname,ch);
		}

	return 1;
	}

static int
_write_stack_file(char *fname, int channel)
	{
	FILE *out;
	double bscale = 1.0;
	int bzero = 0;
	int i,x;
	u32 *data,*ptr;
	struct Image *img;
	char new_fname[1024],*p;

	//Print("write_stack_file '%s' %d\n",fname,channel);
	//fflush(stdout);

	// Allocate the buffer
	data = (u32 *) Malloc(Width * Height * 4);

	// Fetch the image
	fetch_stack_u32(channel,data);
	
	strcpy(new_fname,fname);
	p = new_fname + strlen(new_fname);
	while(p != new_fname && *p != '.') --p;

	// If we can't find a suffix then set up to append
	if (p == new_fname) p = new_fname + strlen(new_fname);

	sprintf(p,"-%d",channel);
	p = new_fname + strlen(new_fname);

	// Monochrome data only. Write out as a 32BPP FITS floating point
	strcpy(p,".fit");

	img = CreateFitsImage(new_fname,Width,Height,32,(unsigned char *)data,bzero,bscale);

	//printf("[%x] %s %d %d %d\n",img,img->dst_fname,Dstwidth,Dstheight,OutputStackFileDepth);
	if (!img) {
	   	Print("write_stack_file: Error creating output img '%s'\n",fname);
	   	exit(1);
	   	}
	img->dst_width = Dstwidth; img->dst_height = Dstheight;

	// Convert from 32bpp unsigned to 32bpp floating point if requested
	if (OutputStackFileDepth == -32) {
		Print("Converting stack to 32 bit floating point\n");
		img = ConvertToType(img, IMG_FIT, -32);
		}

	if (! img) {
		Print("write_stack_file: error converting from U32 to F32\n");
		exit(1);
		}

	Print("Writing stack file '%s' as depth %d\n",fname,img->depth);

       	if (! WriteImage(img)) {
           	Print("Short write on output to %s\n",fname);
           	exit(1);
           	}
	Print("created stack file %s %dx%dx%d\n",new_fname,img->dst_width,img->dst_height,img->depth);
	DestroyImage(img);

	return(1);
	}

//====================================================================================

static struct Image *
load_mergefile(char *fname)
	{
	FILE *in;
	int i,x,count=0;
	double avg;
	int minval8 = ThreshHold;
	int minval16 = minval8<<8 ; // don't process the background
	static struct Image *MImg = NULL;
	unsigned char *ptr;
	unsigned short *uptr;

	// Image already loaded
	if (MImg) {
	   if (!strcmp(MImg->src_fname,fname)) return MImg;
	   Print("Warning: load_mergefile: changing reference image\n");
	   DestroyImage(MImg);
	   }

	MImg = LoadImage(fname,"<internal>");

	if (!MImg) {
	   Print("load_mergefile: cannot open '%s' for reading\n",fname);
	   return NULL;
	   }

	if (MImg->depth != 8 && MImg->depth != 16) {
           Print("load_mergefile: Unsupported depth: %d\n",MImg->depth);
	   DestroyImage(MImg);
           return NULL;
           }

	i = MImg->width * MImg->height;
	ptr = MImg->data;
	uptr = (unsigned short *)MImg->data;

	if (MImg->depth == 8) while(i--) {
             if (*ptr >= minval8) {avg += *ptr; ++count;}
	     ++ptr;
             }
	else if (MImg->depth == 16) while(i--) {
             if (*uptr >= minval16) {avg += *uptr; ++count;}
	     ++uptr;
             }

	MImg->m.avg = (double)avg / count;
	Print("[LoadMerge%d avg=%lf] ",MImg->depth,MImg->m.avg);

	return MImg;
	}

#define MERGE_EQ(a,b) (((a) + ((a)>>1) + ((b)>>1))>>1)

static void
merge_16_16(unsigned short *src, struct Image *r, int npix)
	{
	unsigned short *ref = (unsigned short *)r->data;
	int i,count,high_count=0,low_count=0;
	int minval = ThreshHold << 8;
	double avg,d,scale;

	// calculate average brightness for scaling purposes
	for(i=avg=count=0; i<npix; ++i)
	   if (src[i] >= minval) {avg += src[i]; ++count;}
		
	if (count < MinPixels) {
	   Print("merge_data: found no significant pixels\n");
	   exit(1);
	   }

	avg /= (double) count;
	scale = r->m.avg/avg;

	for(i=0; i<npix; ++i,++src,++ref) {
	   int s = (double)*src * scale + 0.5;
	   if (s >= minval && *ref) {
		double d = s - *ref; d /= *ref;

		if (d >= MergeThreshHold) {
		   *src = MERGE_EQ(s,*ref); ++high_count; }
		else if (d <= -MergeThreshHold) {
		   *src = MERGE_EQ(s,*ref); ++low_count; }
		}
	   }

	Print("%d/%d pixels] ",high_count,low_count);
	return;
	}

static void
merge_8_8(unsigned char *src, struct Image *r, int npix)
	{
	unsigned char *ref = r->data;
	int i,count,high_count=0,low_count=0;
	int minval = ThreshHold;
	double avg,d,scale;

	// calculate average brightness for scaling purposes
	for(i=avg=count=0; i<npix; ++i)
	   if (src[i] >= minval) {avg += src[i]; ++count;}
		
	if (count < MinPixels) {
	   Print("merge_data: found no significant pixels\n");
	   exit(1);
	   }

	avg /= (double) count;
	scale = r->m.avg/avg;

	for(i=0; i<npix; ++i,++src,++ref) {
	   int s = (double)*src * scale + 0.5;
	   if (s >= minval && *ref) {
		double d = s - *ref; d /= *ref;

		if (d >= MergeThreshHold) {
		   *src = MERGE_EQ(s,*ref); ++high_count; }
		else if (d <= -MergeThreshHold) {
		   *src = MERGE_EQ(s,*ref); ++low_count; }
		}
	   }

	Print("%d/%d pixels] ",high_count,low_count);
	return;
	}

static void
merge_16_8(unsigned short *src, struct Image *r, int npix)
	{
	unsigned char *ref = r->data;
	int i,count,high_count=0,low_count=0;
	int minval = ThreshHold << 8;
	double avg,d,scale;

	// calculate average brightness for scaling purposes
	for(i=avg=count=0; i<npix; ++i)
	   if (src[i] >= minval) {avg += src[i]>>8; ++count;}
		
	if (count < MinPixels) {
	   Print("merge_data: found no significant pixels\n");
	   exit(1);
	   }

	avg /= (double) count;
	scale = r->m.avg/avg;

	for(i=0; i<npix; ++i,++src,++ref) {
	   int s = (double)*src * scale + 0.5;
	   int r = *ref << 8;

	   if (s >= minval && r) {
		double d = (s - r) / r;

		if (d >= MergeThreshHold) {
		   *src = MERGE_EQ(s,r); ++high_count; }
		else if (d <= -MergeThreshHold) {
		   *src = MERGE_EQ(s,r); ++low_count; }
		}
	   }

	Print("avg=%lf scale=%lf %d/%d pixels] ",avg,scale,high_count,low_count);
	return;
	}

static void
merge_8_16(unsigned char *src, struct Image *r, int npix)
	{
	unsigned short *ref = (unsigned short *)r->data;
	int i,count,high_count=0,low_count=0;
	int minval = ThreshHold;
	double avg,d,scale;

	// calculate average brightness for scaling purposes
	for(i=avg=count=0; i<npix; ++i)
	   if (src[i] >= minval) {avg += src[i]<<8; ++count;}
		
	if (count < MinPixels) {
	   Print("merge_data: found no significant pixels\n");
	   exit(1);
	   }

	avg /= (double) count;
	scale = r->m.avg/avg;

	for(i=0; i<npix; ++i,++src,++ref) {
	   int s = (double)*src * scale + 0.5;
	   int r = *ref >> 8;

	   if (s >= minval && r) {
		double d = (s - r) / r;

		if (d >= MergeThreshHold) {
		   *src = MERGE_EQ(s,r); ++high_count; }
		else if (d <= -MergeThreshHold) {
		   *src = MERGE_EQ(s,r); ++low_count; }
		}
	   }

	Print("%d/%d pixels] ",high_count,low_count);
	return;
	}

int
merge_data(char *mfile, struct Image *img)
	{
	int width = img->width;
	int height = img->height;
	int depth = img->depth;
	unsigned char *data = img->data;
	struct Image *ref;
	int npix;

	ref = load_mergefile(mfile);

	if (!ref || (width != ref->width) || height != ref->height) {
	   Print("merge_data: Frame size does not match mergefile\n");
	   return 0;
	   }

	npix = width * height;
	Print("[merge %d+%d, ",depth,ref->depth);
	if (depth==16) {
	   if (ref->depth==16)     merge_16_16((unsigned short *)data, ref,npix);
	   else if (ref->depth==8) merge_16_8((unsigned short *)data, ref,npix);
	   }
	else if (depth==8) {
	   if (ref->depth==16) 	   merge_8_16(data, ref,npix);
	   else if (ref->depth==8) merge_8_8(data, ref,npix);
	   }

	return(1);
	}
