#include "ninox.h"

static double scan_area(struct Image *,int,int,int,int);
static void scale_area(struct Image *,int,int,int,int,double);

static struct Image * _create_bmp_888(struct Image *in, int pattern);
static unsigned char * _debayer_GBG(struct Image *in, unsigned char *data);
static unsigned char * _debayer_RGR(struct Image *in, unsigned char *data);
static unsigned char * _debayer_BGB(struct Image *in, unsigned char *data);
static unsigned char * _debayer_GRG(struct Image *in, unsigned char *data);

/*
 * Create a rgb image that is the same dimensions as the
 * given monochrome image, Init the data to empty
 */

static char *pstr[] = {
	"",
	"GBG",
	"RGR",
	"BGB",
	"GRG"
	};

static struct Image *
_create_bmp_888(struct Image *in, int pattern)
	{
	char *ptr;
	struct Image *out = CreateImage();
	char *buf = Malloc(strlen(in->dst_fname) + 32);

        struct bmp_info *bi = ZeroMalloc(sizeof(struct bmp_info));

        out->type=IMG_BMP;
        out->width  = in->dst_width  = in->width;
        out->height = in->dst_height = in->height;
        out->depth = 24;

	bi->cmap_entries = 0;
	bi->cmap = NULL;

        out->info = bi;
        out->data = NULL;

        out->src_fname = strdup("<internal>");

	// Change suffix to BMP and denote pattern
	strcpy(buf,in->dst_fname);
	ptr = buf + strlen(buf)-1;
	while(*ptr && *ptr != '.') --ptr;

	sprintf(ptr,"_%s.bmp",pstr[pattern]);
	
        out->dst_fname = buf;
        out->cutout.x1 = out->cutout.y1 = 0;
        out->cutout.x2 = in->width-1;
        out->cutout.y2 = in->height-1;

	out->data = ZeroMalloc(in->width * in->height * 3);
        return out;
	}

#define LEFTRIGHT(ptr) (((int)(*(ptr-1)) + (int)(*(ptr+1)))/2)
#define UPDOWN(ptr) (((int)(*(ptr-w)) + (int)(*(ptr+w)))/2)
#define ORTHO(ptr) (((int)(*(ptr-w)) + (int)(*(ptr+w)) + ((int)(*(ptr-1)) + (int)(*(ptr+1))))/4)
#define DIAG(ptr) (((int)(*(ptr-w-1)) + (int)(*(ptr+w+1)) + ((int)(*(ptr-w+1)) + (int)(*(ptr+w-1))))/4)
#define GREEN(ptr) (((int)(*(ptr-w-1)) + (int)(*(ptr+w+1)) + (int)(*ptr) + ((int)(*(ptr-w+1)) + (int)(*(ptr+w-1))))/5)

// Assume we are starting at (1,1) and omitting conversion of top and bottom row as well
// as left and right edges.
static unsigned char *
_debayer_GBG(struct Image *in, unsigned char *odata)
	{
	int x,y;
	int w = in->width;
	int h = in->height;
	int r,g,b;
	unsigned char *ptr,*p;

	for(y=1; y<h-1; ++y) {
	   ptr = in->data + y * w + 1;
	   p = odata + (y * w + 1) * 3;

	   for(x=1; x<w-1; ++x,++ptr) {

	      if (y&1) { // RGR
		if (x&1) {
		   r = LEFTRIGHT(ptr); g = GREEN(ptr); b = UPDOWN(ptr);
		   }
		else {
		   r = *ptr; g = ORTHO(ptr); b = DIAG(ptr);
		   }
		}
	      else { // GBG
		if (x&1) {
		   r = DIAG(ptr); g = ORTHO(ptr); b = *ptr;
		   }
		else {
		   r = UPDOWN(ptr); g = GREEN(ptr); b = LEFTRIGHT(ptr);
		   }
		}

	      *(p++) = b;
	      *(p++) = g;
	      *(p++) = r;
	      }
	   }

	return odata;
	}

static unsigned char *
_debayer_RGR(struct Image *in, unsigned char *odata)
	{
	int x,y;
	int w = in->width;
	int h = in->height;
	int r,g,b;
	unsigned char *ptr,*p;

	for(y=1; y<h-1; ++y) {
	   ptr = in->data + y * w + 1;
	   p = odata + (y * w + 1) * 3;

	   for(x=1; x<w-1; ++x,++ptr) {

	      if (!(y&1)) { // RGR
		if (x&1) {
		   r = LEFTRIGHT(ptr); g = GREEN(ptr); b = UPDOWN(ptr);
		   }
		else {
		   r = *ptr; g = ORTHO(ptr); b = DIAG(ptr);
		   }
		}
	      else { // GBG
		if (x&1) {
		   r = DIAG(ptr); g = ORTHO(ptr); b = *ptr;
		   }
		else {
		   r = UPDOWN(ptr); g = GREEN(ptr); b = LEFTRIGHT(ptr);
		   }
		}

	      *(p++) = b;
	      *(p++) = g;
	      *(p++) = r;
	      }
	   }

	return odata;
	}

static unsigned char *
_debayer_BGB(struct Image *in, unsigned char *odata)
	{
	int x,y;
	int w = in->width;
	int h = in->height;
	int r,g,b;
	unsigned char *ptr,*p;

	for(y=1; y<h-1; ++y) {
	   ptr = in->data + y * w + 1;
	   p = odata + (y * w + 1) * 3;

	   for(x=1; x<w-1; ++x,++ptr) {

	      if (y&1) {
		if (!(x&1)) {
		   r = LEFTRIGHT(ptr); g = GREEN(ptr); b = UPDOWN(ptr);
		   }
		else {
		   r = *ptr; g = ORTHO(ptr); b = DIAG(ptr);
		   }
		}
	      else {
		if (!(x&1)) {
		   r = DIAG(ptr); g = ORTHO(ptr); b = *ptr;
		   }
		else {
		   r = UPDOWN(ptr); g = GREEN(ptr); b = LEFTRIGHT(ptr);
		   }
		}

	      *(p++) = b;
	      *(p++) = g;
	      *(p++) = r;
	      }
	   }

	return odata;
	}

static unsigned char *
_debayer_GRG(struct Image *in, unsigned char *odata)
	{
	int x,y;
	int w = in->width;
	int h = in->height;
	int r,g,b;
	unsigned char *ptr,*p;

	for(y=1; y<h-1; ++y) {
	   ptr = in->data + y * w + 1;
	   p = odata + (y * w + 1) * 3;

	   for(x=1; x<w-1; ++x,++ptr) {

	      if (!(y&1)) {
		if (!(x&1)) {
		   r = LEFTRIGHT(ptr); g = GREEN(ptr); b = UPDOWN(ptr); 
		   }
		else {
		   r = *ptr; g = ORTHO(ptr); b = DIAG(ptr);
		   }
		}
	      else {
		if (!(x&1)) {
		   r = DIAG(ptr); g = ORTHO(ptr); b = *ptr;
		   }
		else {
		   r = UPDOWN(ptr); g = GREEN(ptr); b = LEFTRIGHT(ptr);
		   }
		}

	      *(p++) = b;
	      *(p++) = g;
	      *(p++) = r;
	      }
	   }

	return odata;
	}

// **************************************************************
// *                                                            *
// * De-Bayer a monochrome (RAW) image into an 888 colour image *
// *                                                            *
// **************************************************************
//
// Image is 8bpp monochrome ("raw") containing a bayer pattern according to the 
// "pattern" variable:
//
//  1 = GBG/RGR
//  2 = RGR/GBG
//  3 = BGB/GRG
//  4 = GRG/BGB
struct Image *
DeBayer(struct Image *in, int pattern)
	{
	struct Image *img = in;
	struct Image *new;
	int need_destroy = 0;
	
	if (img->depth != 8) {
	   // Have to convert the buffer depth
	   img = ConvertToType(img, IMG_FIT, 8);
	   need_destroy = 1;
	   }

	// Create the skeleton of our output image, everything
	// is present except the rgb data
	new = _create_bmp_888(img,pattern);

	// Depending on our bayer pattern, generate the RGB data
	switch(pattern) {
	   case BAYER_GBG: _debayer_GBG(img,new->data); break;
	   case BAYER_RGR: _debayer_RGR(img,new->data); break;
	   case BAYER_BGB: _debayer_BGB(img,new->data); break;
	   case BAYER_GRG: _debayer_GRG(img,new->data); break;
	   default:
		Print("DeBayer: unknown pattern %d\n",pattern);
		exit(0);
		break;
	   }

	if (need_destroy) DestroyImage(img);

	return new;
	}

// Scan the image looking for a Bayer pattern in the grey values. If found then it means that
// our camera is looking through a bayer mask and we must boost the luminosity of the red and 
// blue pixels to match the green.

int
FixBayer8(struct Image *img)
	{
	int sx,sy,ey,ex;
	double s1,s2,s3,s4,m;
	double f;

	/* calculate a box to scan */
	sx = 1; sy = 1;
	ex = img->width-2; ey = img->height-2;

	/* Scan diagonal area and calculate average luminosity */

	m=0;
	s1 = scan_area(img,sx,sy,ex,ey); if (s1>m) m=s1;
	s2 = scan_area(img,sx,sy+1,ex,ey+1); if (s2>m) m=s2;
	s3 = scan_area(img,sx+1,sy+1,ex+1,ey+1); if (s3>m) m=s3;
	s4 = scan_area(img,sx+1,sy,ex+1,ey); if (s4>m) m=s4;

	// printf("lum %lf %lf %lf %lf\n",s1,s2,s3,s4);

	f = BayerBoost;
	if (m-s1 > 0.1) scale_area(img,sx,sy,ex,ey,m/s1 * f);
	if (m-s2 > 0.1) scale_area(img,sx,sy+1,ex,ey+1,m/s2 * f);
	if (m-s3 > 0.1) scale_area(img,sx+1,sy+1,ex+1,ey+1,m/s3 * f);
	if (m-s4 > 0.1) scale_area(img,sx+1,sy,ex+1,ey,m/s4 * f);

	return 1;
	}

static
double scan_area(
	struct Image *img,
	int x1, int y1, // top left corner of sample region
	int x2, int y2	// bottom right corner of sample region
	)
	{
	unsigned char *buffer = img->data;
	double lum=0;
	int count=0;
	int x,y;
	int o;

	if (img->depth==24) {
	   for(y=y1; y<=y2; y+=2) {
	      o = (y * img->width + x1) * 3;
	      for(x=x1; x<=x2; x+=2,o+=6) {
		if (buffer[o] > ThreshHold) {
	   	   lum += buffer[o]; // This is wrong, should compute luminance
		   count++;
		   }
		}
	      }
	   }
	else if (img->depth==8) {
	   for(y=y1; y<=y2; y+=2) {
	      o = y * img->width + x1;
	      for(x=x1; x<=x2; x+=2,o+=2) {
		if (buffer[o] > ThreshHold) {
	   	   lum += buffer[o];
		   count++;
		   }
		}
	      }
	   }
	else {
	   Print("scan_area: depth %d not supported\n");
	   return 0;
	   }

	if (count==0) {
	   Print("oops - divide by zero in scan_area\n");
	   fflush(stderr); return 0;
	   }

	return lum / (double) count;
	}

static void
scale_area(
	struct Image *img,
	int x1, int y1, // top left corner of sample region
	int x2, int y2,	// bottom right corner of sample region
	double s	// scale factor
	)
	{
	unsigned char *buffer = img->data;
	int x,y;
	int o;
	double v;

	if (img->depth == 24) {
	   for(y=y1; y<=y2; y+=2) {
	      o = (y * img->width + x1) * 3;
	      for(x=x1; x<=x2; x+=2,o+=6) {
		v = buffer[o];
		v *= s;  if (v>255) v=255; if (v<0) v=0;
	   	buffer[o] = (int)v;  // Wrong, should scale according to luminance
	   	buffer[o+1] = (int)v;
	   	buffer[o+2] = (int)v;
		}
	      }
	   }
	else if (img->depth==8) {
	   for(y=y1; y<=y2; y+=2) {
	      o = y * img->width + x1;
	      for(x=x1; x<=x2; x+=2,o+=2) {
		v = buffer[o];
		v *= s;  if (v>255) v=255; if (v<0) v=0;
	   	buffer[o] = (int)v;  // Wrong, should scale according to luminance
		}
	      }
	   }

	}