#include <stdio.h>
#include <stdlib.h>

#include "quantize.h"

cQuantizeWu::cQuantizeWu()
{
}

cQuantizeWu::~cQuantizeWu()
{
}

#define WEIG	//>>2
int cQuantizeWu::Quantize(unsigned char * input, int size, int colors)
{
	struct box cube[MAXCOLOR];
    int next;
	int i;
	int j;
	int k;
	int l;
	float vv[MAXCOLOR];
	float temp;
	unsigned char * tag;
	long weight;

	imageInput = input;
	imageSize = size;
	palSize = colors;

	for (j = 0; j < BOX; j++)
		for (k = 0; k < BOX; k++)
			for (l = 0; l < BOX; l++)
			{
				wt[j][k][l] = 0;
				mr[j][k][l] = 0;
				mg[j][k][l] = 0;
				mb[j][k][l] = 0;
				m2[j][k][l] = 0.0;
			}

	Qadd = new unsigned short[imageSize * sizeof(unsigned short)];
    Hist3d((long *)&wt, (long *)&mr, (long *)&mg, (long *)&mb, (float *)&m2);
    Momt3d((long *)&wt, (long *)&mr, (long *)&mg, (long *)&mb, (float *)&m2);

    cube[0].r0 = cube[0].g0 = cube[0].b0 = 0;
    cube[0].r1 = cube[0].g1 = cube[0].b1 = 32;

	next = 0;
	for (i = 1; i < palSize; i++)
	{
		if (Cut(&cube[next], &cube[i]))
		{
			vv[next] = (cube[next].vol > 1) ? Var(&cube[next]) : 0.0;
			vv[i] = (cube[i].vol > 1) ? Var(&cube[i]) : 0.0;
		}
		else
		{
			vv[next] = 0.0;
			i--;
		}
		next = 0;
		temp = vv[0];
		for (k = 1; k <= i; k++)
			if (vv[k] > temp)
			{
				temp = vv[k];
				next = k;
			}
		if (temp <= 0.0)
		{
			palSize = i + 1;
			break;
		}
	}

    tag = new unsigned char[BOX * BOX * BOX];

	unsigned char * palette = (unsigned char *) paletteOutput;
	for (k = 0; k < palSize; ++k)
	{
		Mark(&cube[k], k, tag);
		weight = Vol(&cube[k], wt);
		if (weight)
		{
#ifdef NOINVERT
			palette[k * 4 + 2] = (Vol(&cube[k], mb) / weight) WEIG;
			palette[k * 4 + 0] = (Vol(&cube[k], mr) / weight) WEIG;
#else
			palette[k * 4 + 0] = (Vol(&cube[k], mb) / weight) WEIG;
			palette[k * 4 + 2] = (Vol(&cube[k], mr) / weight) WEIG;
#endif
			palette[k * 4 + 1] = (Vol(&cube[k], mg) / weight) WEIG;
		}
		else
		{
			palette[k * 4 + 0] = 0;
			palette[k * 4 + 1] = 0;
			palette[k * 4 + 2] = 0;
		}
	}

	for (i = 0; i < imageSize; i++)
	{
		imageOutput[i] = tag[Qadd[i]];
	}

	delete[] tag;
	delete[] Qadd;

	return 0;
}

// build 3-D color histogram of counts, r/g/b, c^2
void cQuantizeWu::Hist3d(long *vwt, long *vmr, long *vmg, long *vmb, float *m_2)
{
	int ind;
	int r;
	int g;
	int b;
    int inr;
	int ing;
	int inb;
	int table[256];
	int i;

	for (i = 0; i < 256; i++)
		table[i] = i * i;

	for (i = 0; i < imageSize; i++)
	{
		b = imageInput[i*4];
		g = imageInput[i*4+1];
		r = imageInput[i*4+2];

		inr = (r >> 3) + 1;
		ing = (g >> 3) + 1;
		inb = (b >> 3) + 1;
		Qadd[i] = ind = (inr << 10) + (inr << 6) + inr + (ing << 5) + ing + inb;

		vwt[ind]++;
		vmr[ind] += r;
		vmg[ind] += g;
		vmb[ind] += b;
		m_2[ind] += (float) (table[r] + table[g] + table[b]);
    }
}

// We now convert histogram into moments so that we can rapidly calculate
// the sums of the above quantities over any desired box.
void cQuantizeWu::Momt3d(long *vwt, long *vmr, long *vmg, long *vmb, float *m_2)
{
	unsigned short ind1;
	unsigned short ind2;
 	unsigned char i;
	unsigned char r;
	unsigned char g;
	unsigned char b;
	long line;
	long line_r;
	long line_g;
	long line_b;
	long area[BOX];
	long area_r[BOX];
	long area_g[BOX];
	long area_b[BOX];
    float line2;
	float area2[BOX];

	for (r = 1; r <= 32; ++r)
	{
		for (i = 0; i <= 32; ++i)
			area2[i] = area[i] = area_r[i] = area_g[i] = area_b[i] = 0;
		for (g = 1; g <= 32; ++g)
		{
			line2 = line = line_r = line_g = line_b = 0;
			for (b = 1; b <= 32; ++b)
			{
				ind1 = (r << 10) + (r << 6) + r + (g << 5) + g + b;
				line += vwt[ind1];
				line_r += vmr[ind1];
				line_g += vmg[ind1];
				line_b += vmb[ind1];
				line2 += m_2[ind1];
				area[b] += line;
				area_r[b] += line_r;
				area_g[b] += line_g;
				area_b[b] += line_b;
				area2[b] += line2;
				ind2 = ind1 - 1089; /* [r-1][g][b] */
				vwt[ind1] = vwt[ind2] + area[b];
				vmr[ind1] = vmr[ind2] + area_r[b];
				vmg[ind1] = vmg[ind2] + area_g[b];
				vmb[ind1] = vmb[ind2] + area_b[b];
				m_2[ind1] = m_2[ind2] + area2[b];
			}
		}
	}
}

int cQuantizeWu::Cut(struct box * set1, struct box * set2)
{
	unsigned char dir;
	int cutr;
	int cutg;
	int cutb;
	float maxr;
	float maxg;
	float maxb;
	long whole_r;
	long whole_g;
	long whole_b;
	long whole_w;

	whole_r = Vol(set1, mr);
	whole_g = Vol(set1, mg);
	whole_b = Vol(set1, mb);
	whole_w = Vol(set1, wt);

	maxr = Maximize(set1, RED, set1->r0 + 1, set1->r1, &cutr,
		whole_r, whole_g, whole_b, whole_w);
    maxg = Maximize(set1, GREEN, set1->g0 + 1, set1->g1, &cutg,
		whole_r, whole_g, whole_b, whole_w);
    maxb = Maximize(set1, BLUE, set1->b0 + 1, set1->b1, &cutb,
		whole_r, whole_g, whole_b, whole_w);

	if ((maxr >= maxg) && (maxr >= maxb))
	{
		dir = RED;
		if (cutr < 0)
			return 0;           /* can't split the box */
    }
	else if ((maxg >= maxr) && (maxg >= maxb))
		dir = GREEN;
	else
		dir = BLUE;

	set2->r1 = set1->r1;
	set2->g1 = set1->g1;
	set2->b1 = set1->b1;

	switch (dir)
	{
    case RED:
		set2->r0 = set1->r1 = cutr;
		set2->g0 = set1->g0;
		set2->b0 = set1->b0;
		break;
	case GREEN:
		set2->g0 = set1->g1 = cutg;
		set2->r0 = set1->r0;
		set2->b0 = set1->b0;
		break;
    case BLUE:
		set2->b0 = set1->b1 = cutb;
		set2->r0 = set1->r0;
		set2->g0 = set1->g0;
		break;
    }
	set1->vol = (set1->r1 - set1->r0) * (set1->g1 - set1->g0) * (set1->b1 - set1->b0);
	set2->vol = (set2->r1 - set2->r0) * (set2->g1 - set2->g0) * (set2->b1 - set2->b0);
    return 1;
}

long cQuantizeWu::Vol(struct box * cube, long mmt[BOX][BOX][BOX])
{
	return (mmt[cube->r1][cube->g1][cube->b1]
		- mmt[cube->r1][cube->g1][cube->b0]
		- mmt[cube->r1][cube->g0][cube->b1]
		+ mmt[cube->r1][cube->g0][cube->b0]
		- mmt[cube->r0][cube->g1][cube->b1]
		+ mmt[cube->r0][cube->g1][cube->b0]
		+ mmt[cube->r0][cube->g0][cube->b1]
		- mmt[cube->r0][cube->g0][cube->b0]);
}

float cQuantizeWu::Maximize(struct box * cube, unsigned char dir, int first, int last, int * cut, long whole_r, long whole_g, long whole_b, long whole_w)
{
	long half_r;
	long half_g;
	long half_b;
	long half_w;
	long base_r;
	long base_g;
	long base_b;
	long base_w;
	int i;
	float temp;
	float max;

	base_r = Bottom(cube, dir, mr);
	base_g = Bottom(cube, dir, mg);
	base_b = Bottom(cube, dir, mb);
	base_w = Bottom(cube, dir, wt);
	max = 0.0;
	*cut = -1;
	for (i = first; i < last; i++)
	{
		half_r = base_r + Top(cube, dir, i, mr);
		half_g = base_g + Top(cube, dir, i, mg);
		half_b = base_b + Top(cube, dir, i, mb);
		half_w = base_w + Top(cube, dir, i, wt);
		// now half_x is sum over lower half of box, if split at i
		if (half_w == 0)
		{
			// subbox could be empty of pixels!
			continue; // never split into an empty box
		}
		else
			temp = ((float) half_r * half_r + (float) half_g * half_g + (float) half_b * half_b) / half_w;

		half_r = whole_r - half_r;
		half_g = whole_g - half_g;
		half_b = whole_b - half_b;
		half_w = whole_w - half_w;
		if (half_w == 0)
		{
			// subbox could be empty of pixels!
			continue; // never split into an empty box
		}
		else
			temp += ((float) half_r * half_r + (float) half_g * half_g + (float) half_b * half_b) / half_w;

		if (temp > max)
		{
			max = temp;
			*cut = i;
		}
	}
	return (max);
}

float cQuantizeWu::Var(struct box * cube)
{
	float dr;
	float dg;
	float db;
	float xx;

	dr = Vol(cube, mr);
	dg = Vol(cube, mg);
	db = Vol(cube, mb);
	xx = m2[cube->r1][cube->g1][cube->b1]
		- m2[cube->r1][cube->g1][cube->b0]
		- m2[cube->r1][cube->g0][cube->b1]
		+ m2[cube->r1][cube->g0][cube->b0]
		- m2[cube->r0][cube->g1][cube->b1]
		+ m2[cube->r0][cube->g1][cube->b0]
		+ m2[cube->r0][cube->g0][cube->b1]
		- m2[cube->r0][cube->g0][cube->b0];

	return (xx - (dr * dr + dg * dg + db * db) / (float) Vol(cube, wt));
}

long cQuantizeWu::Bottom(struct box * cube, unsigned char dir, long mmt[BOX][BOX][BOX])
{
	switch (dir)
	{
		case RED:
			return (-mmt[cube->r0][cube->g1][cube->b1]
				+ mmt[cube->r0][cube->g1][cube->b0]
				+ mmt[cube->r0][cube->g0][cube->b1]
				- mmt[cube->r0][cube->g0][cube->b0]);
		case GREEN:
			return (-mmt[cube->r1][cube->g0][cube->b1]
				+ mmt[cube->r1][cube->g0][cube->b0]
				+ mmt[cube->r0][cube->g0][cube->b1]
				- mmt[cube->r0][cube->g0][cube->b0]);
		case BLUE:
			return (-mmt[cube->r1][cube->g1][cube->b0]
				+ mmt[cube->r1][cube->g0][cube->b0]
				+ mmt[cube->r0][cube->g1][cube->b0]
				- mmt[cube->r0][cube->g0][cube->b0]);
	}
	printf("error in Bottom()");
	return 0;
}

long cQuantizeWu::Top(struct box * cube, unsigned char dir, int pos, long mmt[BOX][BOX][BOX])
{
	switch (dir)
	{
		case RED:
			return (mmt[pos][cube->g1][cube->b1]
				- mmt[pos][cube->g1][cube->b0]
				- mmt[pos][cube->g0][cube->b1]
				+ mmt[pos][cube->g0][cube->b0]);
		case GREEN:
			return (mmt[cube->r1][pos][cube->b1]
				- mmt[cube->r1][pos][cube->b0]
				- mmt[cube->r0][pos][cube->b1]
				+ mmt[cube->r0][pos][cube->b0]);
		case BLUE:
			return (mmt[cube->r1][cube->g1][pos]
				- mmt[cube->r1][cube->g0][pos]
				- mmt[cube->r0][cube->g1][pos]
				+ mmt[cube->r0][cube->g0][pos]);
	}
    printf("error in Top()");
	return 0;
}

void cQuantizeWu::Mark(struct box * cube, int label, unsigned char * tag)
{
	int r;
	int g;
	int b;

	for (r = cube->r0 + 1; r <= cube->r1; r++)
		for (g = cube->g0 + 1; g <= cube->g1; g++)
			for (b = cube->b0 + 1; b <= cube->b1; b++)
				tag[(r << 10) + (r << 6) + r + (g << 5) + g + b] = label;

}