#include "stdafx.h"
#include "Quantizer.h"
#include <MATH.H>

/////////////////////////////////////////////////////////////////////////////
CQuantizer::CQuantizer (UINT nMaxColors, UINT nColorBits)
{
	m_nColorBits = nColorBits < 8 ? nColorBits : 8;

	m_pTree	= NULL;
	m_nLeafCount = 0;
	for	(int i=0; i<=(int) m_nColorBits; i++)
		m_pReducibleNodes[i] = NULL;
	m_nMaxColors = nMaxColors;
}

/////////////////////////////////////////////////////////////////////////////
CQuantizer::~CQuantizer	()
{
	if (m_pTree	!= NULL)
		DeleteTree (&m_pTree);
}

/////////////////////////////////////////////////////////////////////////////
BOOL CQuantizer::ProcessImage (HANDLE hImage)
{
	BYTE r,	g, b;
	int	i, j;

	BITMAPINFOHEADER ds;
	memcpy(&ds,hImage, sizeof(ds));
	int effwdt = ((((ds.biBitCount * ds.biWidth ) + 31) / 32) * 4);

	int	nPad = effwdt - (((ds.biWidth *	ds.biBitCount) + 7) / 8);

	BYTE* pbBits = (BYTE*)hImage + *(LPDWORD)hImage;

	switch (ds.biBitCount) {

	case 1:	// 1-bit DIB
	case 4:	// 4-bit DIB
	case 8:	// 8-bit DIB
		for	(i=0; i<ds.biHeight;	i++) {
			for	(j=0; j<ds.biWidth; j++)	{
				BYTE idx=GetPixelIndex(j,i,ds.biBitCount,effwdt,pbBits);
				BYTE* pal = (BYTE*)(hImage) + sizeof(BITMAPINFOHEADER);
				long ldx = idx*sizeof(RGBQUAD);
				b = pal[ldx++];
				g = pal[ldx++];
				r = pal[ldx];
				AddColor (&m_pTree,	r, g, b, m_nColorBits, 0, &m_nLeafCount,
					m_pReducibleNodes);
				while (m_nLeafCount	> m_nMaxColors)
					ReduceTree (m_nColorBits, &m_nLeafCount,
						m_pReducibleNodes);
			}
		}
		break;
	case 15:
	case 16://16bit
		for	(i=0; i<ds.biHeight;	i++) 
		{
			for	(j=0; j<ds.biWidth; j++)
			{
				b = (*pbBits)&0x1F;
				g = (*pbBits++)>>5;
				g |= ((*pbBits)&0x03)<<3;
				r = ((*pbBits++)>>2)&0x1F;

				r *= 8;
				b *= 8;
				g *= 8;
			
				AddColor (&m_pTree,	r, g, b, m_nColorBits, 0, &m_nLeafCount,
					m_pReducibleNodes);
				while (m_nLeafCount	> m_nMaxColors)
					ReduceTree (m_nColorBits, &m_nLeafCount, m_pReducibleNodes);
			}

			pbBits += nPad;
		}
		break;
	case 24: //	24-bit DIB
		for	(i=0; i<ds.biHeight;	i++) {
			for	(j=0; j<ds.biWidth; j++)	{
				b =	*pbBits++;
				g =	*pbBits++;
				r =	*pbBits++;
				AddColor (&m_pTree,	r, g, b, m_nColorBits, 0, &m_nLeafCount,
					m_pReducibleNodes);
				while (m_nLeafCount	> m_nMaxColors)
					ReduceTree (m_nColorBits, &m_nLeafCount, m_pReducibleNodes);
			}
			pbBits += nPad;
		}
		break;
	case 32: //	32-bit DIB
		for	(i=0; i<ds.biHeight;	i++) {
			for	(j=0; j<ds.biWidth; j++)	{
				b =	*pbBits++;
				g =	*pbBits++;
				r =	*pbBits++;
				pbBits++;

				AddColor (&m_pTree,	r, g, b, m_nColorBits, 0, &m_nLeafCount,
					m_pReducibleNodes);
				while (m_nLeafCount	> m_nMaxColors)
					ReduceTree (m_nColorBits, &m_nLeafCount, m_pReducibleNodes);
			}
			pbBits += nPad;
		}
		break;
	default: //	Unrecognized color format
		return FALSE;
	}
	return TRUE;
}

/////////////////////////////////////////////////////////////////////////////
void CQuantizer::AddColor (NODE** ppNode, BYTE r, BYTE g, BYTE b,
	UINT nColorBits, UINT nLevel, UINT*	pLeafCount,	NODE** pReducibleNodes)
{
	static BYTE	mask[8]	= {	0x80, 0x40,	0x20, 0x10,	0x08, 0x04,	0x02, 0x01 };

	// If the node doesn't exist, create it.
	if (*ppNode	== NULL)
		*ppNode	= (NODE*)CreateNode (nLevel, nColorBits, pLeafCount, pReducibleNodes);

	// Update color	information	if it's	a leaf node.
	if ((*ppNode)->bIsLeaf)	
	{
		(*ppNode)->nPixelCount++;
		(*ppNode)->nRedSum += r;
		(*ppNode)->nGreenSum +=	g;
		(*ppNode)->nBlueSum	+= b;
	} else 
	{	// Recurse a level deeper if the node is not a leaf.
		int	shift =	7 -	nLevel;
		int	nIndex =(((r & mask[nLevel]) >> shift) << 2) |
					(((g & mask[nLevel]) >>	shift) << 1) |
					(( b & mask[nLevel]) >> shift);
		AddColor(&((*ppNode)->pChild[nIndex]),	r, g, b, nColorBits,
					nLevel + 1,	pLeafCount,	pReducibleNodes);
	}
}

/////////////////////////////////////////////////////////////////////////////
void* CQuantizer::CreateNode (UINT nLevel, UINT	nColorBits,	UINT* pLeafCount,
	NODE** pReducibleNodes)
{
	NODE* pNode = (NODE*)calloc(1,sizeof(NODE));

	if (pNode== NULL) return NULL;

	pNode->bIsLeaf = (nLevel ==	nColorBits)	? TRUE : FALSE;
	if (pNode->bIsLeaf) (*pLeafCount)++;
	else {
		pNode->pNext = pReducibleNodes[nLevel];
		pReducibleNodes[nLevel]	= pNode;
	}
	return pNode;
}

/////////////////////////////////////////////////////////////////////////////
void CQuantizer::ReduceTree	(UINT nColorBits, UINT*	pLeafCount,
	NODE** pReducibleNodes)
{
	int i;
	// Find	the	deepest	level containing at	least one reducible	node.
	for	(i=nColorBits -	1; (i>0) &&	(pReducibleNodes[i]	== NULL); i--);

	// Reduce the node most	recently added to the list at level	i.
	NODE* pNode	= pReducibleNodes[i];
	pReducibleNodes[i] = pNode->pNext;

	NODE* pNodeTmp = pNode;
	NODE* pNeedReduceNode = NULL,*pPreNode = NULL,*pPreNodeTmp = NULL;
	long nMinPixelCount = 0;

	while(pNodeTmp)
		{
		long nPixelCount = 0;
		for	(int m=0; m<8; m++)	
			{
			if (pNodeTmp->pChild[m] !=	NULL) 
				{
				nPixelCount += pNodeTmp->pChild[m]->nPixelCount;
				}
			}

		if(nMinPixelCount == 0)
			{
			nMinPixelCount = nPixelCount;
			pNeedReduceNode = pNodeTmp;
			}
		else if(nPixelCount < nMinPixelCount)
			{
			nMinPixelCount = nPixelCount;
			pNeedReduceNode = pNodeTmp;
			pPreNode = pPreNodeTmp;
			}

		pPreNodeTmp = pNodeTmp;
		pNodeTmp = pNodeTmp->pNext;
		}

	if(pPreNode)
		{
		pPreNode->pNext = pNeedReduceNode->pNext;
		pReducibleNodes[i] = pNode;
		}

	pNode = pNeedReduceNode;

	UINT nRedSum = 0;
	UINT nGreenSum = 0;
	UINT nBlueSum =	0;
	UINT nChildren = 0;

	for	(i=0; i<8; i++)	
		{
		if (pNode->pChild[i] !=	NULL) 
			{
			nRedSum	+= pNode->pChild[i]->nRedSum;
			nGreenSum += pNode->pChild[i]->nGreenSum;
			nBlueSum +=	pNode->pChild[i]->nBlueSum;
			pNode->nPixelCount += pNode->pChild[i]->nPixelCount;
			free(pNode->pChild[i]);
			pNode->pChild[i] = NULL;
			nChildren++;
			}
		}

	pNode->bIsLeaf = TRUE;
	pNode->nRedSum = nRedSum;
	pNode->nGreenSum = nGreenSum;
	pNode->nBlueSum	= nBlueSum;
	*pLeafCount	-= (nChildren -	1);
}

/////////////////////////////////////////////////////////////////////////////
void CQuantizer::DeleteTree	(NODE**	ppNode)
{
	for	(int i=0; i<8; i++)	{
		if ((*ppNode)->pChild[i] !=	NULL) DeleteTree (&((*ppNode)->pChild[i]));
	}
	free(*ppNode);
	*ppNode	= NULL;
}


BOOL CQuantizer::FindColorIndex(NODE *pNode,BYTE r,BYTE g,BYTE b,int nLevel,UINT *pColorIndex)
{
	static BYTE	mask[8]	= {	0x80, 0x40,	0x20, 0x10,	0x08, 0x04,	0x02, 0x01 };

	// If the node doesn't exist, create it.
	if (pNode == NULL)
	{	
		return FALSE;
	}

	// Update color	information	if it's	a leaf node.
	if (pNode->bIsLeaf)	
	{
		*pColorIndex = pNode->nColorIndex;
		return TRUE;
	}
	
	// Recurse a level deeper if the node is not a leaf.
	int	shift =	7 -	nLevel;
	int	nIndex =(((r & mask[nLevel]) >> shift) << 2) |
				(((g & mask[nLevel]) >>	shift) << 1) |
				(( b & mask[nLevel]) >> shift);

	return FindColorIndex(pNode->pChild[nIndex], r, g, b, nLevel+1, pColorIndex);
}

BOOL CQuantizer::GetColorIndex(BYTE r,BYTE g,BYTE b,UINT *pColorIndex)
{
	return FindColorIndex(m_pTree,r,g,b,0,pColorIndex);
}

/////////////////////////////////////////////////////////////////////////////
void CQuantizer::GetPaletteColors (NODE* pTree,	RGBQUAD* prgb, UINT* pIndex)
{
	if (pTree){
		if (pTree->bIsLeaf)	{
			prgb[*pIndex].rgbRed = (BYTE)((pTree->nRedSum)/(pTree->nPixelCount));
			prgb[*pIndex].rgbGreen = (BYTE)((pTree->nGreenSum)/(pTree->nPixelCount));
			prgb[*pIndex].rgbBlue = (BYTE)((pTree->nBlueSum)/(pTree->nPixelCount));
			prgb[*pIndex].rgbReserved =	0;
			pTree->nColorIndex = *pIndex;
			(*pIndex)++;
		} else {
			for	(int i=0; i<8; i++)	{
				if (pTree->pChild[i] !=	NULL)
					GetPaletteColors (pTree->pChild[i],	prgb, pIndex);
			}
		}
	}
}

/////////////////////////////////////////////////////////////////////////////
UINT CQuantizer::GetColorCount ()
{
	return m_nLeafCount;
}

/////////////////////////////////////////////////////////////////////////////
void CQuantizer::SetColorTable (RGBQUAD* prgb)
{
	UINT nIndex	= 0;
	GetPaletteColors (m_pTree, prgb, &nIndex);
}

/////////////////////////////////////////////////////////////////////////////
BYTE CQuantizer::GetPixelIndex(long x, long y, int nbit, long effwdt, BYTE *pimage)
{
	if (nbit==8)
	{
		return pimage[y*effwdt + x];
	} else 
	{
		BYTE pos;
		BYTE iDst= pimage[y*effwdt + (x*nbit >> 3)];
		if (nbit==4)
		{
			pos = 4*(1-x%2);
			iDst &= (0x0F<<pos);
			return iDst >> pos;
		} else if (nbit==1)
		{
			pos = 7-x%8;
			iDst &= (0x01<<pos);
			return iDst >> pos;
		}
	}
	return 0;
}

CRGBQuantizer::CRGBQuantizer(UINT nMaxColors, UINT nColorBits)
:CQuantizer(nMaxColors,nColorBits)
{
}

BOOL CRGBQuantizer::ProcessImageRGB(BYTE *pRGBData,UINT nWidth,UINT nHeight)
{
	BYTE r,g,b;

	for	(int i=0; i<nHeight;	i++) 
	{
		for	(int j=0; j<nWidth; j++)	
		{
			r =	*pRGBData++;
			g =	*pRGBData++;
			b =	*pRGBData++;
			AddColor (&m_pTree,	r, g, b, m_nColorBits, 0, &m_nLeafCount,
				m_pReducibleNodes);
			while (m_nLeafCount	> m_nMaxColors)
				ReduceTree (m_nColorBits, &m_nLeafCount, m_pReducibleNodes);
		}
	}

	return TRUE;
}

CBitmapQuantizer::CBitmapQuantizer(UINT nMaxColors, UINT nColorBits)
:CQuantizer(nMaxColors,nColorBits)
{
}

BOOL CBitmapQuantizer::ProcessImageBitmap (HBITMAP hBmp)
{
	BITMAP bm;
	PBITMAPINFO bmpInf;

	if(GetObject(hBmp,sizeof(bm),&bm)==0)
		return FALSE;

	int nPaletteSize=0;

	if(bm.bmBitsPixel<16)
		nPaletteSize=(int)pow(2.0,bm.bmBitsPixel);

	bmpInf=(PBITMAPINFO)LocalAlloc(LPTR,sizeof(BITMAPINFOHEADER)+
		sizeof(RGBQUAD)*nPaletteSize+
		(bm.bmWidth+7)/8*bm.bmHeight*bm.bmBitsPixel);

	BYTE* buf = ((BYTE*)bmpInf)+
		sizeof(BITMAPINFOHEADER)+
		sizeof(RGBQUAD)*nPaletteSize;
	
	//-----------------------------------------------
	bmpInf->bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
	bmpInf->bmiHeader.biWidth = bm.bmWidth;
	bmpInf->bmiHeader.biHeight = bm.bmHeight;
	bmpInf->bmiHeader.biPlanes = bm.bmPlanes;
	bmpInf->bmiHeader.biBitCount = bm.bmBitsPixel;
	bmpInf->bmiHeader.biCompression = BI_RGB;
	bmpInf->bmiHeader.biSizeImage = (bm.bmWidth+7)/8*bm.bmHeight*bm.bmBitsPixel;
	bmpInf->bmiHeader.biXPelsPerMeter = 0;
	bmpInf->bmiHeader.biYPelsPerMeter = 0;
	bmpInf->bmiHeader.biClrUsed = 0;
	bmpInf->bmiHeader.biClrImportant = 0;
	//-----------------------------------------------

	HDC hDC = ::GetWindowDC(NULL);

	if(!::GetDIBits(hDC,hBmp,0,(UINT)bm.bmHeight,buf,bmpInf,DIB_RGB_COLORS))
	{
		::ReleaseDC(NULL,hDC);
		LocalFree(bmpInf);
		return FALSE;
	}

	::ReleaseDC(NULL,hDC);

	bmpInf->bmiHeader.biSize = sizeof(BITMAPINFOHEADER)+sizeof(RGBQUAD)*nPaletteSize;
	BOOL bRet = ProcessImage(bmpInf);

	LocalFree(bmpInf);

	return bRet;
}