#include <v_lib.h>
#include "framebuf.h"
#include "eval.h"
#include "cspace.h"

typedef struct image_input_s
{
	char		id[MAX_ID];
	framebuf_p	pixels;
	matrix_p	rgba[4];
	matrix_p	hsv[3];
	matrix_p	cie[3];
} image_input_t, *image_input_p;

V_LOCAL vp_stack_t	image_cache = TYPED_STACK_EMPTY_INITIALIZER;

V_LOCAL void
rgbFromCie(vec3_t rgb_ret,vec3_t cie)
{
	static matrix3_t cie2rgb_transform={
		{2.739,-1.145,-0.424},
		{-1.11,2.029,0.033},
		{0.138,-0.333,1.105}};
	mat3VecMul(rgb_ret,cie2rgb_transform,cie);
}

V_LOCAL void
cieFromRgb(vec3_t cie_ret,vec3_t rgb)
{
	static matrix3_t rgb2cie_transform={
        {0.477,0.298,0.174},
        {0.261,0.653,0.081},
        {0.019,0.160,0.907}};
	mat3VecMul(cie_ret,rgb2cie_transform,rgb);
}

V_LOCAL void
rgbFromHsvUnit(vec3_t rgb_ret,vec3_t unit_hsv)
{
	vec3_t hsv;
	hsv[0] = unit_hsv[0] * (vec_t)360.0;
	hsv[1] = unit_hsv[1];
	hsv[2] = unit_hsv[2];
	rgbFromHsv(rgb_ret,hsv);
}

V_LOCAL void
hsvUnitFromRgb(vec3_t hsv_unit_ret,vec3_t rgb)
{
	hsvFromRgb(hsv_unit_ret,rgb);
	hsv_unit_ret[0] *= (vec_t)(1.0 / 360.0);
}

V_LOCAL void
rgbFromRgb(vec3_t rgb_ret,vec3_t rgb)
{
	rgb_ret[0] = rgb[0];
	rgb_ret[1] = rgb[1];
	rgb_ret[2] = rgb[2];
}

void
ctransFill3Matrices(framebuf_p pixels,u32_t target_cspace,
	matrix_p *c1_ret,matrix_p *c2_ret,matrix_p *c3_ret)
{
	vec_p *cells1,*cells2,*cells3;
	vec3_t in, out;
	u32_t x, y;
	void (*transform)(vec3_t dest,vec3_t src);
	if(!*c1_ret || (*c1_ret)->rows != pixels->height ||
		(*c1_ret)->cols != pixels->width)
	{
		if(*c1_ret)
			memFree(*c1_ret);
		*c1_ret = matAlloc(pixels->height,pixels->width);
	}
	if(!*c2_ret || (*c2_ret)->rows != pixels->height ||
		(*c2_ret)->cols != pixels->width)
	{
		if(*c2_ret)
			memFree(*c2_ret);
		*c2_ret = matAlloc(pixels->height,pixels->width);
	}
	if(!*c3_ret || (*c3_ret)->rows != pixels->height ||
		(*c3_ret)->cols != pixels->width)
	{
		if(*c3_ret)
			memFree(*c3_ret);
		*c3_ret = matAlloc(pixels->height,pixels->width);
	}
	cells1 = (*c1_ret)->cell;
	cells2 = (*c2_ret)->cell;
	cells3 = (*c3_ret)->cell;
	switch(target_cspace)
	{
		case CSPACE_HSV:
			transform = hsvUnitFromRgb;
			break;
		case CSPACE_CIE:
			transform = cieFromRgb;
			break;
		default:
			transform = rgbFromRgb;
			break;
	}
	for(y = 0;y < pixels->height;y++)
	{
		for(x = 0;x < pixels->width;x++)
		{
			fbReadPixel3v(pixels,x,y,in);
			(*transform)(out,in);
			cells1[y][x] = out[0];
			cells2[y][x] = out[1];
			cells3[y][x] = out[2];
		}
	}
	return;
}

void
ctransStore3Matrices(framebuf_p pixels,u32_t orig_cspace,
	matrix_p c1,matrix_p c2,matrix_p c3)
{
	vec_p *cells1,*cells2,*cells3;
	vec3_t in, out;
	u32_t x, y;
	void (*transform)(vec3_t dest,vec3_t src);
	V_ASSERT(c1 && c2 && c3);
	V_ASSERT(c1->rows == pixels->height && c1->cols == pixels->width);
	V_ASSERT(c2->rows == pixels->height && c2->cols == pixels->width);
	V_ASSERT(c3->rows == pixels->height && c3->cols == pixels->width);
	cells1 = c1->cell;
	cells2 = c2->cell;
	cells3 = c3->cell;
	switch(orig_cspace)
	{
		case CSPACE_HSV:
			transform = rgbFromHsvUnit;
			break;
		case CSPACE_CIE:
			transform = rgbFromCie;
			break;
		default:
			transform = rgbFromRgb;
			break;
	}
	for(y = 0;y < pixels->height;y++)
	{
		for(x = 0;x < pixels->width;x++)
		{
			in[0] = cells1[y][x];
			in[1] = cells2[y][x];
			in[2] = cells3[y][x];
			(*transform)(out,in);
			fbDrawPixel3v(pixels,x,y,out);
		}
	}
	return;
}

void
ctransFill1Matrix(framebuf_p pixels,matrix_p *m_ret,
	vboolean from_r,vboolean from_g,vboolean from_b,vboolean from_a)
{
	u8_p rgba;
	u32_t x, y, channel_count;
	vec_t scale;
	vec_p row;
	if(!*m_ret || (*m_ret)->rows != pixels->height ||
		(*m_ret)->cols != pixels->width)
	{
		if(*m_ret)
			memFree(*m_ret);
		*m_ret = matAlloc(pixels->height,pixels->width);
	}
	matZero(*m_ret);
	channel_count = 0;
	if(from_r)
	{
    	for(y=0;y < pixels->height;y++)
    	{
    		rgba = pixels->scanlines[y];
    		row = (*m_ret)->cell[y];
    		for(x=0;x < pixels->width;x++,rgba+=4)
    			row[x] += rgba[0];
    	}
    	channel_count++;
    }
	if(from_g)
	{
    	for(y=0;y < pixels->height;y++)
    	{
    		rgba = pixels->scanlines[y];
    		row = (*m_ret)->cell[y];
    		for(x=0;x < pixels->width;x++,rgba+=4)
    			row[x] += rgba[1];
    	}
    	channel_count++;
    }
	if(from_b)
	{
    	for(y=0;y < pixels->height;y++)
    	{
    		rgba = pixels->scanlines[y];
    		row = (*m_ret)->cell[y];
    		for(x=0;x < pixels->width;x++,rgba+=4)
    			row[x] += rgba[2];
    	}
    	channel_count++;
    }
	if(from_a)
	{
    	for(y=0;y < pixels->height;y++)
    	{
    		rgba = pixels->scanlines[y];
    		row = (*m_ret)->cell[y];
    		for(x=0;x < pixels->width;x++,rgba+=4)
    			row[x] += rgba[3];
    	}
    	channel_count++;
    }
    if(channel_count)
    {
	    scale = 1 / (vec_t)(0xFF * channel_count);
    	matScale(*m_ret,*m_ret,scale);
    }
}

void
ctransStore1Matrix(framebuf_p pixels,matrix_p m,
	vboolean to_r,vboolean to_g,vboolean to_b,vboolean to_a)
{
	u8_p rgba;
	u32_t x, y;
	vec_t scaled, scale = 255;
	vec_p row;
	u8_p ptr, one_channel;
	V_ASSERT(m && m->rows == pixels->height && m->cols == pixels->width);
	one_channel = (u8_p)memAlloc(pixels->width * pixels->height);
	V_ASSERT(one_channel);
	/* pre - compute scaled and clamped values */
	ptr = one_channel;
	for(y = 0;y < m->rows;y++)
	{
		row = m->cell[y];
		for(x = 0;x < m->cols;x++,ptr++)
		{
			scaled = row[x] * scale;
			*ptr = CLAMP_FLOAT_TO_BYTE(scaled);
		}
	}
	/* store them to target channels */
	if(to_r)
	{
		ptr = one_channel;
		for(y=0;y < pixels->height;y++)
		{
			rgba = pixels->scanlines[y];
			for(x=0;x < pixels->width;x++,rgba+=4,ptr++)
				*rgba = *ptr;
		}
	}
	if(to_g)
	{
		ptr = one_channel;
		for(y=0;y < pixels->height;y++)
		{
			rgba = pixels->scanlines[y] + 1;
			for(x=0;x < pixels->width;x++,rgba+=4,ptr++)
				*rgba = *ptr;
		}
	}
	if(to_b)
	{
		ptr = one_channel;
		for(y=0;y < pixels->height;y++)
		{
			rgba = pixels->scanlines[y] + 2;
			for(x=0;x < pixels->width;x++,rgba+=4,ptr++)
				*rgba = *ptr;
		}
	}
	if(to_a)
	{
		ptr = one_channel;
		for(y=0;y < pixels->height;y++)
		{
			rgba = pixels->scanlines[y] + 3;
			for(x=0;x < pixels->width;x++,rgba+=4,ptr++)
				*rgba = *ptr;
		}
	}
	memFree(one_channel);
}

vboolean
cspaceHistogram(vector_p dest,matrix_p src,vboolean normalize)
{
	u32_t rn, cn, rx, cx;
	u32_t sum, i, r, c, n_divs = dest->size;
    vec_t lower, upper, min, max, delta, *row;
    matMin(src,&rn,&cn);
    matMax(src,&rx,&cx);
    min = src->cell[rn][cn];
    max = src->cell[rx][cx];
    delta = (max - min) / n_divs;
    for(i = 0;i< n_divs;i++)
    {
    	sum = 0;
        lower = min + i * delta;
        upper = lower + delta;
        for(r = 0;r < src->rows;r++)
        {
        	row = src->cell[r];
        	for(c = 0;c < src->cols;c++)
            {
            	if(row[c] >= lower && row[c] < upper)
                	sum++;
        	}
		}
        dest->elements[i] = sum;
    }
    if(normalize)
    {
    	vec_t vsum = dest->elements[0];
        for(i = 1;i < n_divs;i++)
        	vsum += dest->elements[i];
		if(vsum > 0)
        {
        	vsum = 1 / vsum;
            vecScale(dest,dest,vsum);
        }
        else
        	return false;
    }
    return true;
}

vboolean
cspaceMaximizeContrast(matrix_p dest,matrix_p src)
{
	u32_t r, c;
	vec_t min, max, rscale;
    matMin(src,&r,&c);
    min = src->cell[r][c];
    matMax(src,&r,&c);
    max = src->cell[r][c];
    rscale = max - min;
    if(!IsZero(rscale))
    {
    	rscale = 1 / rscale;
    	matScaleAdd(dest,src,rscale,- min * rscale);
    	return true;
    }
    return false;
}

vboolean
cspaceGammaCorrection(matrix_p dest,matrix_p src,vector_p gamma,
	vboolean unit_input)
{
	u32_t i, r, c, n_divs = gamma->size;
    vec_t min, max, t, delta, *rs, *rd;
    if(unit_input)
    {
    	min = 0;
        max = 1;
	}
    else
    {
		u32_t rn, cn, rx, cx;
    	matMin(src,&rn,&cn);
    	matMax(src,&rx,&cx);
    	min = src->cell[rn][cn];
    	max = src->cell[rx][cx];
	}
    delta = (max - min) / n_divs;
    if(IsZero(delta))
    	return false;
    for(r = 0;r < dest->rows;r++)
    {
    	rd = dest->cell[r];
    	rs = src->cell[r];
        for(c = 0;c < dest->cols;c++)
        {
        	t = (rs[c] - min) / delta; /* 0 <= t < n_divs */
            if(t < 0)
            	t = 0;
            i = (u32_t)t;
            if(i >= n_divs)
            	i = n_divs - 1;
			rd[c] = gamma->elements[i];
		}
	}
    return true;
}

V_LOCAL image_input_p
inputFindCachePic(eval_state_p state,char *image_name)
{
	u32_t i;
	image_input_p pic;
	pic = (image_input_p)vpsGetStorage(&image_cache);
	for(i = 0;i < vpsCount(&image_cache);i++,pic++)
		if(STR_EQ(pic->id,image_name))
			return pic;
	pic = (image_input_p)memAllocZ(sizeof(image_input_t));
	V_ASSERT(pic);
	V_ASSERT(strlen(image_name) < MAX_ID);
	strCopy(pic->id,image_name,MAX_ID);
	pic->pixels = fbLoadEx(image_name,
		state->custom_width,state->custom_height,
		state->input_powers2,state->input_square,true);
	if(!pic->pixels)
	{
		memFree(pic);
		return NULL;
	}
	vpsPush(&image_cache,pic);
	return pic;
}

matrix_p
inputGetMatrix(eval_state_p state,char *image_name,char *channel_id)
{
	image_input_p pic = inputFindCachePic(state,image_name);
	if(!pic)
		return NULL;
	if(STR_EQ(channel_id,"red"))
	{
		if(!pic->rgba[0])
			ctransFill3Matrices(pic->pixels,CSPACE_RGB,
				pic->rgba+0,pic->rgba+1,pic->rgba+2);
		return pic->rgba[0];
	}
	else if(STR_EQ(channel_id,"green"))
	{
		if(!pic->rgba[1])
			ctransFill3Matrices(pic->pixels,CSPACE_RGB,
				pic->rgba+0,pic->rgba+1,pic->rgba+2);
		return pic->rgba[1];
	}
	else if(STR_EQ(channel_id,"blue"))
	{
		if(!pic->rgba[2])
			ctransFill3Matrices(pic->pixels,CSPACE_RGB,
				pic->rgba+0,pic->rgba+1,pic->rgba+2);
		return pic->rgba[2];
	}
	else if(STR_EQ(channel_id,"alpha"))
	{
		if(!pic->rgba[3])
			ctransFill1Matrix(pic->pixels,pic->rgba+3,
				false,false,false,true);
		return pic->rgba[3];
	}
	else if(STR_EQ(channel_id,"hue"))
	{
		if(!pic->hsv[0])
			ctransFill3Matrices(pic->pixels,CSPACE_HSV,
				pic->hsv+0,pic->hsv+1,pic->hsv+2);
		return pic->hsv[0];
	}
	else if(STR_EQ(channel_id,"saturation"))
	{
		if(!pic->hsv[1])
			ctransFill3Matrices(pic->pixels,CSPACE_HSV,
				pic->hsv+0,pic->hsv+1,pic->hsv+2);
		return pic->hsv[1];
	}
	else if(STR_EQ(channel_id,"value"))
	{
		if(!pic->hsv[2])
			ctransFill3Matrices(pic->pixels,CSPACE_HSV,
				pic->hsv+0,pic->hsv+1,pic->hsv+2);
		return pic->hsv[2];
	}
	else if(STR_EQ(channel_id,"cieX"))
	{
		if(!pic->cie[0])
			ctransFill3Matrices(pic->pixels,CSPACE_CIE,
				pic->cie+0,pic->cie+1,pic->cie+2);
		return pic->cie[0];
	}
	else if(STR_EQ(channel_id,"cieY"))
	{
		if(!pic->cie[1])
			ctransFill3Matrices(pic->pixels,CSPACE_CIE,
				pic->cie+0,pic->cie+1,pic->cie+2);
		return pic->cie[1];
	}
	else if(STR_EQ(channel_id,"cieZ"))
	{
		if(!pic->cie[2])
			ctransFill3Matrices(pic->pixels,CSPACE_CIE,
				pic->cie+0,pic->cie+1,pic->cie+2);
		return pic->cie[2];
	}
	else
	{
		conlPrintf(LOG_ERROR,"inputGetMatrix(%s,%s): unknown color space.\n",
			image_name,channel_id);
		return NULL;
	}
}

void
inputCollectGarbage(void)
{
	u32_t i;
	image_input_p ptr;
	while(vpsPop(&image_cache,(void_p*)&ptr))
	{
		if(ptr->pixels)
			memFree(ptr->pixels);
		for(i=0;i<3;i++)
		{
			if(ptr->rgba[i])
				memFree(ptr->rgba[i]);
			if(ptr->hsv[i])
				memFree(ptr->rgba[i]);
			if(ptr->cie[i])
				memFree(ptr->cie[i]);
		}
		if(ptr->rgba[3])
			memFree(ptr->rgba[3]);
		memFree(ptr);
	}
	vpsClear(&image_cache);
}
