// ------------------------------------------------------------
// clob.c
//
// A color quantizer using binary-split technique
// From:
// "Color Quantization of Images", Orchard, Bouman
// IEEE Trans. on Sig. Proc., vol. 39, no. 12, pp. 2677-2690, Dec. 1991.
//
// 10.03.2005, J.Ahrensfeld
// ------------------------------------------------------------
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <sys/time.h>
#include "bintree.h"
#include "mutil.h"
#include "lapack_cwrap.h"
#include "clob.h"

// ------------------------------------------------------------
void CodebookAlloc(codebook_t *pObj, int N_max, int M_max)
{
	int i;

	pObj->pColor = (color_vector_t*)malloc(M_max*sizeof(color_vector_t));
	pObj->size_C = (int*)malloc(M_max*sizeof(int));
	memset(pObj->size_C, 0, M_max*sizeof(int));
	pObj->pC = (int*)malloc(N_max*sizeof(int));
	pObj->ppC = (int**)malloc(M_max*sizeof(int*));

	pObj->ppC[0] = pObj->pC;;
	for (i=1; i < M_max; i++)
		pObj->ppC[i] = NULL;

	pObj->max_colors = M_max;
	pObj->max_samples = N_max;
	pObj->num_colors = 0;

}

// ------------------------------------------------------------
void CodebookFree(codebook_t *pObj)
{
	int i;

	if(!pObj)
		return;

	free(pObj->pC);
	free(pObj->ppC);
	free(pObj->pColor);
	free(pObj->size_C);
}

// ------------------------------------------------------------
void ColorQuant_Init(colorquant_t *pCQ, int N_max, int M_max)
{
	M_max++;
	pCQ->N_max = N_max;
	pCQ->M_max = M_max;

	pCQ->pRootData = (node_data_t*)malloc(sizeof(node_data_t));
	memset(pCQ->pRootData, 0, sizeof(node_data_t));

	pCQ->pData = (node_data_t*)malloc(2*M_max*sizeof(node_data_t));
	memset(pCQ->pData, 0, 2*M_max*sizeof(node_data_t));

	pCQ->pTemp = (double*)malloc(N_max*sizeof(double));
	pCQ->pXs = (color_vector_t*)malloc(N_max*sizeof(color_vector_t));

}

// ------------------------------------------------------------
void ColorQuant_Free(colorquant_t *pCQ)
{
	int m;

	FreeNodeData(pCQ->pRootData);
	for (m=0; m < 2*pCQ->M_max; m++)
	{
		FreeNodeData(&pCQ->pData[m]);
	}
	free(pCQ->pRootData);
	pCQ->pRootData = NULL;
	free(pCQ->pData);
	pCQ->pData = NULL;
	free(pCQ->pXs);
	free(pCQ->pTemp);
}

// ------------------------------------------------------------
void FreeNodeData(node_data_t *pObj)
{
	if(pObj)
	{
		if(pObj->pC)
		{
			free(pObj->pC);
			pObj->pC = NULL;
		}
		if(pObj->pExpr_left)
		{
			free(pObj->pExpr_left);
			pObj->pExpr_left = NULL;
		}
	}
}

// ------------------------------------------------------------
void ColorNodeStat(colorquant_t *pCQ, node_data_t *pObj, color_vector_t *pXs)
{
	int i, j, s, info;

	if (pObj->N == 0)
	{
		fprintf(stderr, "ColorNodeStat(): N is zero!\n");
		return;
	}

	// Calc node statistics
	// R = sum_s(xs_s*xs_s') ; mit s = Element aus Menge C
	memset(pObj->R, 0, ROW_DIM*ROW_DIM*sizeof(double));
	dsyrk('U', 'N', ROW_DIM, pObj->N, 1.0, (double*)pXs, ROW_DIM, 1.0, (double*)pObj->R, ROW_DIM);
//	PrintMatrix("R=",(double*)pObj->R, ROW_DIM, ROW_DIM);

	// m = sum_s(xs_s)
	memset(pObj->m, 0, ROW_DIM*sizeof(double));
	for(i=0; i < pObj->N; i++)
	{
		pObj->m[0] += pXs[i][0];
		pObj->m[1] += pXs[i][1];
		pObj->m[2] += pXs[i][2];

	}
//	PrintMatrix("m=",(double*)pObj->m, ROW_DIM, 1);

}

// ------------------------------------------------------------
double ColorLambda(colorquant_t *pCQ, node_data_t *pObj, color_vector_t *pXs)
{

	int i, s, info, lwork;
	color_vector_t d = {0};
	double *pWork;
	double *pDiff, Ls;

	if (pObj->N == 0)
	{
		fprintf(stderr, "ColorLambda(): N is zero!\n");
		pObj->lambda = 0;
		return 0;
	}

	pObj->pExpr_left = (double*)malloc(pObj->N*sizeof(double));

	// q = m/N;
	memcpy(pObj->q, pObj->m, sizeof(color_vector_t));
	dscal (ROW_DIM, 1.0/pObj->N, pObj->q, 1);
//	PrintMatrix("q=",(double*)pObj->q, ROW_DIM, 1);

	// cov = R - 1/N*m*m'
	memcpy(pObj->cov, pObj->R, ROW_DIM*sizeof(color_vector_t));
	dsyr('U', ROW_DIM, -1.0/pObj->N, pObj->m, 1, (double*)pObj->cov, ROW_DIM);
//	PrintMatrix("Cov=",(double*)pObj->cov, ROW_DIM, ROW_DIM);

	// Query dsyev() for optimal work buffer size
	pWork = pCQ->pTemp;
	pWork[0] = 0;
	info = dsyev('V', 'U', ROW_DIM, (double*)NULL, ROW_DIM, NULL, pWork, -1);
	lwork = (int)pWork[0];

	// d = eig(cov)
	info = dsyev('V', 'U', ROW_DIM, (double*)pObj->cov, ROW_DIM, d, pWork, lwork);
	if (info != 0)
	{
		fprintf (stderr, "Eigenvalue solver failure with error %d\n", info);
		printf("N=%d\n", pObj->N);
		PrintMatrix("Cov=",(double*)pObj->cov, ROW_DIM, ROW_DIM);
	}

//	PrintMatrix("d=",(double*)d, ROW_DIM, 1);
//	PrintMatrix("V=",(double*)pObj->cov, ROW_DIM, ROW_DIM);

	// Form unit-vector e in direction of max. element in d
	dcopy (ROW_DIM, d, 1, pObj->e, 1);
	GetPrincipleEigenVector((double*)pObj->cov, (double*)pObj->e, ROW_DIM);
//	PrintMatrix("e=",(double*)pObj->e, ROW_DIM, 1);

	// -------------------------
	// Compute Lambda

	// Expr. left: dot(xs,e) = e'*xs_s = xs_s'*e
	// Also used later in ColorSplit()
	dgemv('T', ROW_DIM, pObj->N, 1.0, (double*)pXs, ROW_DIM, pObj->e, 1, 0.0, pObj->pExpr_left, 1);

	// Expr. right: dot(q,e) = e'*q = q'*e
	// Also used later in ColorSplit()
	pObj->expr_right = ddot(ROW_DIM, pObj->q, 1, pObj->e, 1);

	// lambda = sum_s((xs_s - q)'*e).^2);
	pObj->lambda = 0;
	for (i=0; i < pObj->N; i++)
	{
		// xs_s'*e - q'*e
		Ls = pObj->pExpr_left[i] - pObj->expr_right;
		pObj->lambda += Ls*Ls;
	}
	return pObj->lambda;
}

// ------------------------------------------------------------
void ColorNodeSplit(colorquant_t *pCQ, node_data_t *pObj, node_data_t *pData2n, node_data_t *pData2n1)
{
	int N2n, N2n1, s, i;

	memset(pData2n, 0, sizeof(node_data_t));
	memset(pData2n1, 0, sizeof(node_data_t));

	pData2n->pC = (int*)malloc(pObj->N*sizeof(int));
	if(!pData2n->pC)
	{
		fprintf(stderr, "pData2n->pC = NULL for N=%d\n", pObj->N);
	}

	pData2n1->pC = (int*)malloc(pObj->N*sizeof(int));
	if(!pData2n1->pC)
	{
		fprintf(stderr, "pData2n1->pC = NULL for N=%d\n", pObj->N);
	}

	// Split into two new sets
	N2n = 0;
	N2n1 = 0;
	for (i=0; i < pObj->N; i++)
	{
		s = pObj->pC[i];

		if(pObj->pExpr_left[i] > pObj->expr_right)
		{
			pData2n1->pC[N2n1] = s;
			N2n1++;
		}
		else
		{
			pData2n->pC[N2n] = s;
			N2n++;
		}
	}

	// Compute statistics of two new nodes
	// Node(2n)
	pData2n->N = N2n;
	for (i=0; i < N2n; i++)
	{
		s = pData2n->pC[i];
		pCQ->pXs[i][0] = pCQ->pX[s][0];
		pCQ->pXs[i][1] = pCQ->pX[s][1];
		pCQ->pXs[i][2] = pCQ->pX[s][2];
	}

	ColorNodeStat(pCQ, pData2n, pCQ->pXs);
	ColorLambda(pCQ, pData2n, pCQ->pXs);

	// Node(2n+1)
	pData2n1->N = N2n1;
	for(i=0; i < ROW_DIM; i++)
	{
		pData2n1->R[i][0] = pObj->R[i][0] - pData2n->R[i][0];
		pData2n1->R[i][1] = pObj->R[i][1] - pData2n->R[i][1];
		pData2n1->R[i][2] = pObj->R[i][2] - pData2n->R[i][2];
	}

	for(i=0; i < ROW_DIM; i++)
		pData2n1->m[i] = pObj->m[i] - pData2n->m[i];

	for (i=0; i < N2n1; i++)
	{
		s = pData2n1->pC[i];
		pCQ->pXs[i][0] = pCQ->pX[s][0];
		pCQ->pXs[i][1] = pCQ->pX[s][1];
		pCQ->pXs[i][2] = pCQ->pX[s][2];
	}
	ColorLambda(pCQ, pData2n1, pCQ->pXs);

}

// ------------------------------------------------------------
int ColorQuant(colorquant_t *pCQ, codebook_t *pCB, color_vector_t *pX, int K, int N, int M, double lambda_min)
{
	int i, m, *pC, bst_state, s, sum_pixel, codebook_size, size, *ptr, *ptr_next;
	node_data_t *pData, *pData2n, *pData2n1;
	node_t *node_max, *node_temp;
	bintree_t bst;

	if (N == 0)
	{
		fprintf(stderr, "ColorQuant(): N is zero!\n");
		return 0;
	}

	pCQ->nRow = K;
	pCQ->nCol = N;

	pCQ->pRootData->N = N;
	pCQ->pX = (color_vector_t*)pX;

	pCQ->pRootData->pC = (int*)malloc(N*sizeof(int));

	// Create root data set
	for(i=0; i < N; i++)
	{
		pCQ->pRootData->pC[i] = i;				// store indices of vector boundaries
	}

	ColorNodeStat(pCQ, pCQ->pRootData, pCQ->pX);
	ColorLambda(pCQ, pCQ->pRootData, pCQ->pX);

	// Create root-node
	bintree_init(&bst, 2*N);

	bintree_insert(&bst, pCQ->pRootData->lambda, pCQ->pRootData);

	for (m=0; m < M-1; m++)
	{
		node_max = bintree_GetNodeMax(&bst);
		if (!node_max)
		{
			fprintf(stderr,"node_max = NULL!\n");
			break;
		}

		if (node_max->value < lambda_min)
			break;

		pData2n = &pCQ->pData[2*m];
		pData2n1 = &pCQ->pData[2*m+1];
		ColorNodeSplit(pCQ, (node_data_t*)node_max->pData, pData2n, pData2n1);
		if (pData2n->N > 0)
		{
			bintree_insert(&bst, pData2n->lambda, pData2n);
		}
		else
		{
			FreeNodeData(pData2n);
		}

		if (pData2n1->lambda > 0)
		{
			bintree_insert(&bst, pData2n1->lambda, pData2n1);
		}
		else
		{
			FreeNodeData(pData2n1);
		}

//		if ((pData2n->lambda == 0) || (pData2n1->lambda == 0))
//		{
//			continue;
//		}

		node_temp = bintree_node_remove(&bst, node_max);
		if (node_temp == bst.pRoot)
		{
			break;
		}
		FreeNodeData((node_data_t*)node_max->pData);

		if (!bst.pRoot)
			break;

		bst_state = bintree_isBST(bst.pRoot);
		if (!bst_state)
		{
			fprintf(stderr,"Binary Tree is not BST!\n");
		}
	}

//	bintree_print(pCQ->root);
//	printf("\n");
	sum_pixel = 0;
	codebook_size = 0;
	pCB->ppC[0] = pCB->pC;
	do
	{
		node_max = bintree_GetNodeMax(&bst);
		if (!node_max)
		{
			fprintf(stderr,"node_max = NULL!\n");
		}
//		printf("Lambda=%5.5g\n",node_max->value);

		size = ((node_data_t*)node_max->pData)->N;
		pCB->size_C[codebook_size] = size;
		ptr = pCB->ppC[codebook_size];
		ptr_next = ptr + size;
		sum_pixel += size;
		memcpy(ptr,((node_data_t*)node_max->pData)->pC, size*sizeof(int));
		memcpy(pCB->pColor[codebook_size],((node_data_t*)(node_max->pData))->q, sizeof(color_vector_t));
		if ((codebook_size+1) < pCB->max_colors)
			pCB->ppC[codebook_size+1] = (int*)(ptr_next);

		FreeNodeData((node_data_t*)(node_max->pData));
		node_temp = bintree_node_remove(&bst, node_max);
		codebook_size++;

	} while(node_temp);
	pCB->num_colors = codebook_size;

//	printf("Num. pixels   = %d\n",sum_pixel);
//	printf("Codebook size = %d\n",codebook_size);
//	PrintMatrix("Codebook=",(double*)pCB->pColor, ROW_DIM, codebook_size);

	FreeNodeData(pCQ->pRootData);
	bintree_free(&bst);
	return codebook_size;
}