/*
 *	greenscreen1.3.c
 *	a yuv based plugin for greenscreen removal
 *	dv4.1.1 filter
 *	edge blur and green spill control
 *	
 *
 *
 * 2/28/07         paprmh ->gmail.com	License: GPL of course
 */


#include <math.h>

#include "plugin.h"
#include "stdio.h"
#include "iff.h"
#include "util.h"
#include "string.h"

#include "stdlib.h"
#include "floatpatch.h"


/*for the TOG3 buttons*/
#define TOG3 (7<<9)
#define SHO  64
#define BIT 256

/* ******************** GLOBAL VARIABLES ***************** */

char name[24]= "greenscreen1.3";

/* structure for buttons,
 *  butcode      name           default  min  max  0
 */

VarStruct varstr[]= {

	{TOG3|SHO,	"greenscreen1.3",		1.0,	0.0,	0.1,	"toggle greenscreen plugin off || on || invert selection"},
	{TOG|INT,	"dv4.1.1",		0.0,	0.0,	0.0,	"dv 4:1:1 compression filter"},
	{NUMSLI|FLO,	"Ymin",			0.1,	0.0,	1.0,	"lower luminance clip"},
	{NUMSLI|FLO,	"Ymax",			0.9,	0.0,	1.0,	"upper luminance clip"},
	{NUMSLI|FLO,	"Uclip",			0.5,	0.0,	1.0,	"blue factor clipping value"},
	{NUMSLI|FLO,	"Vclip",			0.5,	0.0,	1.0,	"red factor clipping value"},

	{NUM|INT,	"display:",		0.0,	0.0,	5.0,	"display: 0)composited output, 1)selection, 2)alpha mask, 3) Y, 4) U, 5) V"},
	{TOG3|SHO,"check / ibuf2",	0.0,	0.0,	1.0,	"composite background: color(off) || checker board || use ibuf2"},
	{COL|CHA,	"back Color",		0.0,	0.0, 255.0, 	"background color"},
	{TOG|INT,	"edge blur",		0.0,	0.0,	1.0,	"alpha blur of mask edge"},
	{NUM|FLO,	"blur",			0.6,	0.0,	10.0,	"Amount of edge blur"},
	{NUM|INT,	"radius",			1.0,	1.0,	5.0,	"radius of edge blur [kernal=(2*radius)+1]"},

	{TOG3|SHO,"info / TMI",		1.0,	0.0,	1.0,	"output plugin info to the console or (probably)Too Much Info"},
	{TOG3|SHO,"spill / spill+",		0.0,	0.0,	1.0,	"green spill reduction w/ respill: manual adjustment || auto to bk color"},
	{NUM|FLO,	"factor",			0.0,	0.0,	1.0,	"green spill reduction factor"},
	{NUM|FLO,	"Yadj",			0.0,	-1.0,	1.0,	"luminance adjustment of spill area -> not active with spill+"},
	{NUM|FLO,	"Uadj",			0.0,	-1.0,	1.0,	"blue adjustment of spill area -> not active with spill+"},
	{NUM|FLO,	"Vadj",			0.0,	-1.0,	1.0,	"red adjustment of spill area -> not active with spill+"},

	
	
	
	
};
/* The cast struct is for input in the main doit function
   Varstr and Cast must have the same variables in the same order */

typedef struct Cast {

	short doit[2];
	unsigned int dv411;
	float ymin;
	float ymax;
	float uclip;
	float vclip;

	int showchannel;
	short check[2];
	unsigned char bcol[4];
	int edgeblur;
	float eblur;
	int erad;
	
	short talk[2];
	short bleed[2];
	float tgfact;
	float yadj;
	float uadj;
	float vadj;

	
	
} Cast;


/* cfra: the current frame */

float cfra;

void plugin_seq_doit(Cast *, float, float, int, int, ImBuf *, ImBuf *, ImBuf *, ImBuf *);

/* ******************** Fixed functions ***************** */

int plugin_seq_getversion(void)
{
	return B_PLUGIN_VERSION;
}

void plugin_but_changed(int but)
{
}

void plugin_init(void)
{
}

void plugin_getinfo(PluginInfo *info)
{
	info->name= name;
	info->nvars= sizeof(varstr)/sizeof(VarStruct);
	info->cfra= &cfra;

	info->varstr= varstr;

	info->init= plugin_init;
	info->seq_doit= (SeqDoit) plugin_seq_doit;
	info->callback= plugin_but_changed;
}


/***my functions*************************************************************/


/*turn 32bit rgba into float(0->1)*/

void ch2fl(unsigned char *rgba,float *frgba, int x){
	int i;
	for(i=0;i<x; i++) frgba[i]=(float)rgba[i]/255.0;
}

/*****turn float array into unsigned char- clipped ********/

void fl2ch(float *frgba,unsigned char *rgba, int x){
	int i, temp;

	for(i=0; i<x; i++){
		temp=(int)(frgba[i]*255)+0.5;
		if(temp<0)temp=0;
		else if(temp>255)temp=255;
		rgba[i]=(unsigned char)temp;
	}
}
/*****these are set up so that you can call with same pointer***/
/**************RGB to YUV*********************/

void frgb2fyuv(float *frgb, float *fyuv) {
	float y,u,v;
	y= 0.299*frgb[0] + 0.587*frgb[1] + 0.114*frgb[2];
	u= -(0.147*frgb[0])-(0.289*frgb[1])+(0.436*frgb[2])+0.5;
	v= (0.615*frgb[0])-(0.515*frgb[1])-(0.100*frgb[2])+0.5;

	fyuv[0]= y;
	fyuv[1]= u;
	fyuv[2]= v;
	fyuv[3]= frgb[3];
}

/***********YUV to RGB***********************/

void fyuv2frgb(float *fyuv, float *frgb) {
	float r,g,b;
	r= fyuv[0] + (1.14*(fyuv[2]-0.5));
	g= fyuv[0] - (0.395*(fyuv[1]-0.5)) - (0.581*(fyuv[2]-0.5));
	b= fyuv[0] + (2.032*(fyuv[1]-0.5));

	frgb[0]= r;
	frgb[1]= g;
	frgb[2]= b;
	frgb[3]= fyuv[3];
}

/*************float gaussian blur w/ channel select, h-v separate options*********************/

void gaussianBlur(float *fsrc, float *fdest, int xo, int yo, float bfactor, int rad, int talk,int bmode) {
	float k, fpixel[4];
	float total=0,t=0;
	int h = rad,totop=xo*yo*4;
	int i,x,y,p,p2,bchmode;

	float *temp, *filter, *fptr;

	if ((bfactor <= 0)||(!rad)) {
		if(talk) printf("no blur\n");
		return;
	}
	if(talk) printf("\nstart blur...");

	temp = mallocN(sizeof(float)*totop,"tempbuf");
	filter = mallocN(sizeof(float)*(rad*2+1),"filterbuf");
	if(!filter || !temp){
		printf("memory allocation falure...blur aborted\n");
		return;
	}
	if(talk>1) printf("malloc ok...");

	k = -1.0/(2.0*3.14159*bfactor*bfactor);
	for (i = -rad; i <= rad; i++) {
		filter[i+rad] =(float)expf(k*(i*i));
		t += filter[i+rad];
	}

	for (i = -rad; i <= rad; i++) {
		filter[i+rad] /= t;
		total+=filter[i+rad];
	if(talk>1) printf("filter %d: %f...subtotal=%f\n",i,filter[i+rad],total);
	}

	bchmode=(bmode<0)? (-bmode):bmode;
	bchmode--;
	if(talk&&bchmode) printf("blur channel %d only...",bchmode);

/************** horizontal pass *************************/
	for (x = 0; x < xo; x++) {
		for (y = 0; y < yo; y++) {

			p = (x+(y*xo))*4;
			if(bchmode){
				fdest[p]=fsrc[p];
				fdest[p+1]=fsrc[p+1];
				fdest[p+2]=fsrc[p+2];
				fdest[p+3]=fsrc[p+3];
			}
			for (i = -rad; i <= rad; i++) {

				if (x+i >= xo)	p2 = (xo-1+(y*xo))*4;
				else if (x+i <= 0) p2 = (1+(y*xo))*4;
				else p2 = (x+i+(y*xo))*4;
				fptr=(float*)fsrc+p2;
				if(bchmode){
					temp[p+bchmode]+=fptr[bchmode]*filter[i+h];
				}
				else{
					temp[p]+=fptr[0]*filter[i+h];
					temp[p+1]+=fptr[1]*filter[i+h];
					temp[p+2]+=fptr[2]*filter[i+h];
					temp[p+3]+=fptr[3]*filter[i+h];
				}
			}

		}
	}

	if(talk>1) printf("horiz pass %dx ok...",i+h);
	if(bmode>0){					/*****no vert blur******/
		for(x=0;x<totop;x+=4){
			if(bchmode){			/***if just one channel**/
				fdest[x+bchmode]=temp[x+bchmode];
			}
		    else{
				fdest[x]=temp[x];
				fdest[x+1]=temp[x+1];
				fdest[x+2]=temp[x+2];
				fdest[x+3]=temp[x+3];
			}
		}
		if(talk) printf("no vert blur...\n");
	}

	/* vertical pass */
	else {
		for (x = 0; x < xo; x++) {
			for (y = 0; y < yo; y++) {
				p = (x+(y*xo))*4;
				fpixel[0]=fpixel[1]=fpixel[2]=fpixel[3]=0;

				for (i = -rad; i <= rad; i++) {
					if(y+i >= yo) p2 = (x+((yo-1)*xo))*4;
					else if(y+i <= 0) p2 = x*4;
					else p2 = (x+((y+i)*xo))*4;

					if(bchmode){
						fpixel[bchmode]+=temp[p2+bchmode]*filter[i+h];
					}
					else{
						fpixel[0]+= temp[p2]*filter[i+h];
						fpixel[1]+= temp[p2+1]*filter[i+h];
						fpixel[2]+= temp[p2+2]*filter[i+h];
						fpixel[3]+= temp[p2+3]*filter[i+h];
					}
				}
				if(bchmode){
					fdest[p+bchmode]=fpixel[bchmode];
				}
				else{
					fdest[p]=fpixel[0];
					fdest[p+1]=fpixel[1];
					fdest[p+2]=fpixel[2];
					fdest[p+3]=fpixel[3];
				}
			}
		}
		if(talk>1) printf("vert pass %dx ok\n",i+h);
	}


	/* Take out the trash */
	freeN(temp);
	freeN(filter);

}

/***************makecheck...make grey 8x8 pixel checker board**********/
void makecheck(float *board, int sx, int sy){

	float gx,gy=0.333;
	int x,y;

	for(y=0;y<sy;y++){
		if(!(y%8))gy= 1-gy;
		gx=gy;
		for(x=0;x<sx;x++,board+=4){
			if(!(x%8))gx=1-gx;
			board[0]=gx;
			board[1]=gx;
			board[2]=gx;
			board[3]=0;
		}
	}
}


/*************************************************************************/






/****main function*********************************************************/

void plugin_seq_doit(Cast *cast, float facf0, float facf1,
	int sx, int sy, ImBuf *ibuf1, ImBuf *ibuf2, ImBuf *out, ImBuf *use)
{

	unsigned char *csrc, *cdest;
	float *fsrc, *fdest, *fwork, *fback; 
	int ix, pc, ri, gi, bi, bmode, bleed, check, talk;
	int totpix, totop, totopf;
	int keytog=0;
	float ftemp, fcol[4],fyuv[4], min2max[3];
	char showstr[6][25]= {
		"normal output","selection","key mask...","luminence(y)","chrominance(u)","chrominance(v)"
		};
	talk= cast->talk[0]+cast->talk[1];
	check= cast->check[0]+cast->check[1];
	bleed= cast->bleed[0]+cast->bleed[1];
	if(talk)printf("\ngreenscreen...frame %d...setting up\n",(int)cfra);
	
/***********setup**********************/
	totpix=sx*sy;
	totop=totpix*4;
	totopf=totop*sizeof(float);
	bmode=(-1);
	
	csrc=(unsigned char *) ibuf1->rect;
	cdest=(unsigned char *) out->rect;
	fsrc=(float *) ibuf1->rect_float;
	fdest=(float *) out->rect_float;
	
		if(!cast->doit[0]){
		if(talk)printf("just lurking...\n");
		if(!fdest) memcpy(cdest, csrc, totop);
		else memcpy(fdest, fsrc, totopf);
		return;
	}

/*******set up work buffer***************/

	fwork = mallocN(totopf,"workbuf");
	if(!fwork){
		printf("buffer allocation failure\n");
		if(fdest) memcpy(fdest,fsrc,totopf);
		else if(cdest) memcpy(cdest, csrc, totop);
		return;
	}
		
	if(!fdest){
		if(talk)printf("\nno floatbuffer...this plugin uses float values->converting...\n");
		if(!csrc){
			if(talk)printf("heh...no char buffer either...aborting\n");
			return;
		}
		for(ix=0;ix<totop;ix+=4){
			ch2fl(csrc+ix, fcol, 4);
			frgb2fyuv(fcol, fwork+ix);
		}

		fback = mallocN(totopf,"malfback");
		if(!fback){
			printf("buffer allocation failure\n");
			if(cdest){
				memcpy(cdest, csrc, totop);
				return;
			}
		}
	}
	else{
		for(ix=0;ix<totop;ix+=4)frgb2fyuv(fsrc+ix, fwork+ix);
		fback= fdest;
	}
	
	
	
	/******set up background ********/
	ch2fl(cast->bcol, fcol, 4);
	if(check){
		if(check>1){
			if(ibuf2 && (ibuf2!= ibuf1)){
				if(talk && cast->showchannel!=1) printf("-> background: ibuf2...\n");
				if(ibuf2->rect_float) memcpy(fback, ibuf2->rect_float, totopf);
				else ch2fl((unsigned char *)ibuf2->rect, fback, totop);
				for(ix=3;ix<totop+3;ix+=4) fback[ix]=0.0;
			}
			else{
				if(talk && cast->showchannel!= 1)printf("...no ibuf2...defaulting...");
				cast->check[1]=0;
				cast->check[0]=0;
				check=0;
			}
		}
		else{
			makecheck( fback, sx, sy);
			if(talk&&cast->showchannel!=1)printf("-> background: checkerboard...\n");
		}
	}

	if(!check){
		for(ix=0; ix<totop;ix+=4){
			fback[ix]= fcol[0];
			fback[ix+1]= fcol[1];
			fback[ix+2]= fcol[2];
			fback[ix+3]= 0.0;
		}
		if(talk&&cast->showchannel!=1)printf("-> background color  R:%4.3f G:%4.3f B:%4.3f A:%4.3f ...\n",fcol[0], fcol[1], fcol[2], fcol[3]);
	}



/*******dv 4:1:1 filter**********/

	if(cast->dv411){
		if(talk) printf("dv 4:1:1 filter U & V blur...");
		gaussianBlur( fwork, fwork, sx, sy, 1.0, 2, talk, bmode*2);
		gaussianBlur( fwork, fwork, sx, sy, 1.0, 2, talk, bmode*3);
		if(talk) printf("dv filter done\n");
	}

/************display options****************************/

	if(talk) printf("show: %d) %s ",cast->showchannel,showstr[cast->showchannel]);
	
	if(cast->showchannel==1){
		if(talk) printf("channels...");
		for(ix=0;ix<totop;ix+=4){
			if(fwork[ix]< cast->ymin) ftemp= 0.5;
			else	ftemp= (fwork[ix]> cast->ymax)? 1:0;
			
			fback[ix]= (fwork[ix+2]> cast->vclip)? 1:0;
			fback[ix+1]= ftemp;
			fback[ix+2]= (fwork[ix+1]> cast->uclip)? 1:0;
		}
	}
	
	else if(cast->showchannel>2){
		pc= (cast->showchannel-3);
		if(talk) printf("(yuv[%d])...", pc);
		for(ix=0;ix<totop;ix+=4){
			ftemp= fwork[ix+pc];
			fback[ix]= ftemp;
			fback[ix+1]= ftemp;
			fback[ix+2]= ftemp;
			fback[ix+3]= ftemp;
		}
	}

/*********************key process loop****************************************/
	else{
		if(talk) printf("\nprocessing image...");
		frgb2fyuv(fcol,fcol);

		for(ix=0;ix<totop;ix+=4){
			if(fwork[ix+3]){

				if((fwork[ix]<= cast->ymin) || (fwork[ix]>= cast->ymax) || 
				(fwork[ix+1]>= cast->uclip) || (fwork[ix+2]>= cast->vclip)) keytog=1;
				else keytog=0;
				
				if(cast->doit[1]) keytog= (keytog)? 0:1;
				if(keytog){
					fyuv2frgb((float*)fwork+ix,(float*)fback+ix);
					
					if(bleed){
						ri=gi=bi=0;

						/*min to max array for average of lower 2 values*/
						(fback[ix]<fback[ix+1])? gi++:ri++;
						(fback[ix]<fback[ix+2])? bi++:ri++;
						(fback[ix+1]<fback[ix+2])? bi++:gi++;
						min2max[ri]=fback[ix];
						min2max[gi]=fback[ix+1];
						min2max[bi]=fback[ix+2];
						ftemp= (min2max[1]+min2max[0])/2;
						
						ftemp= cast->tgfact*(fback[ix+1]-ftemp);

						if(fback[ix+1]>min2max[0]) fback[ix+1]-=ftemp; /* reduce green */
						frgb2fyuv(fback+ix,fyuv);
						
						if(bleed>1){
							fyuv[0]= fwork[ix]+ ((fcol[0]-fwork[ix])*ftemp*2);	/*use origional Y*/
							fyuv[1]+= ((fcol[1]-fyuv[1])*ftemp);
							fyuv[2]+= ((fcol[2]-fyuv[2])*ftemp);
						}
						else{
							fyuv[0]= fwork[ix]+cast->yadj*ftemp;	/*use origional Y*/
							fyuv[1]+= cast->uadj* ftemp;
							fyuv[2]+= cast->vadj* ftemp;
						}
						
						fyuv2frgb(fyuv,fback+ix);
					}
				}
			}
		}

		if(talk){
			printf("key mask set...");
			if(cast->doit[1]) printf("inverting mask...");
			if(bleed> 1) printf("\nspill+ control at %4.3f%%...auto respill to background color\nR:%4.3f G:%4.3f B:%4.3f A:%4.3f ...\n", cast->tgfact, fcol[0], fcol[1], fcol[2], fcol[3]);
			else if(bleed) printf("\nspill control at %4.3f%%...manual respill\n",cast->tgfact);
			 
		}

/************** blur mask for edge blur********************/
		if(cast->edgeblur){
			memcpy(fwork, fback, totopf);
			if(talk) printf("edge blur...");
			gaussianBlur(fwork,fwork,sx,sy,cast->eblur,cast->erad,talk,bmode);
		}


/**********show key mask*************************/
		if(cast->showchannel==2){
			fsrc= (cast->edgeblur)? fwork:fback;
			for(ix=0;ix<totop;ix+=4){
				ftemp= fsrc[ix+3];
				fback[ix]= ftemp;
				fback[ix+1]= ftemp;
				fback[ix+2]= ftemp;
				fback[ix+3]= ftemp;
			}
		}
/**********composit *****************/
		else{
			if(cast->edgeblur){
				for(ix=0; ix<totop;ix+=4){
					if((fwork[ix+3]<.95) && (fwork[ix+3]>.05)){
						fback[ix]= fwork[ix];
						fback[ix+1]= fwork[ix+1];
						fback[ix+2]= fwork[ix+2];
						fback[ix+3]= fwork[ix+3];
					}
				}
			}
		}

	}

/**********output**************/
	if(!fdest){
		if(talk) printf("\nconverting back to char values...");
		fl2ch(fback, cdest, totop);
		freeN(fback);
	}
	freeN(fwork);
	if(talk) printf("done\n");

}