
#include <stdlib.h> 
#include <stdio.h> 
#include <algorithm> //for use of sort or min_element
#include <vector> //for use of vector
#include "vision_utilities.h"

using namespace std;

const int MAXSIZE = 511;//the maximum Colors of the image (after first pass), also the maximum number of the objects (labels) in the image

int findLabel(int x);//find the root of the label in the UFS
void unionLabel(int label1, int label2);//union the two labels that are equivalent
int pre[MAXSIZE+1];//the linked list used in the UFS

/*
THIS PROGRAM IS USED TO LABEL DISCONNECTED REGIONS IN AN IMAGE.
THE EQUIVALENCE TABLE IS IMPLEMENTED BY USING UFS (UNION-FIND SET).
*/

/*
Every UFS contains the nodes that are unioned (in this case, labels that are "equivalent").
UFS has two main operations: find and union.
"find" can find the root of a node in its UFS (in this case, find a label to "represent" all labels in a row of the equivalence table).
"union" can add a node to another node (and thus its UFS) (in this case, add a label to a row of the equivalence table).
int findLabel(int x): function that acts as "find".
void unionLabel(int label1, int label2): function that acts as "union".
*/


int p2(char *fname_input, char *fname_output)
{
	
	Image *im_input = new Image;//input image
	Image *im_output = new Image;//output image
	readImage(im_input, fname_input);
	int NRows = getNRows(im_input);
	int NCols = getNCols(im_input);
	setSize(im_output,NRows,NCols);
	setColors(im_output,MAXSIZE);//temporarily set Colors of the output image to its maximum

	int label = 0;//initialize the label
	for( int i=0; i <= MAXSIZE; ++i ) pre[i] = -1;
	/*
		At the beginning, assign -1 to all the nodes in the linked list, 
		so that later for root node x, pre[x] = -num (num is the number of nodes in the corresponding UFS); 
		for non-root node y, pre[y] > 0, pre[y] is the father of y.
	*/

	vector<int> *eqTable = new vector<int>[MAXSIZE+1];//equivalence table
	/*
		eqTable[l] has all the labels that are equivalent to label "l"
		for example, eqTabel[1][0] == 1, eqTabel[1][1] == 3, eqTabel[1][2] == 4,
		that means equivalence table: 1==1,3,4
	*/
	eqTable[0].push_back(0);

	vector<int> tempLabel;//NOT consecutive natural numbers: like 1,2,70...
	/*
		tempLabel contains all the labels that are used after joining the equivalent labels,
		but are not consecutive natural numbers.
	*/
	tempLabel.push_back(1);

	int finalLabel;//consecutive natural numbers: 1,2,3,4...

	int **labelArray = new int*[NRows];
	for (int i = 0; i<NRows; i++)
	{
		labelArray[i] = new int[NCols];
	}
	/*
		labelArray is a 2-D array that contains the label at each pixel of input image
		labelArray[i][j] == label at pixel[i][j]
	*/

	/*======FIRST PASS======*/
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j<NCols; j++)
		{
			int pixel = getPixel(im_input,i,j);//gray-level of the input image at (i,j)

			int label_nw = labelArray[(i-1)>0?(i-1):0][(j-1)>0?(j-1):0];//label at (i-1,j-1)
			int label_w = labelArray[i][(j-1)>0?(j-1):0];//label at (i,j-1)
			int label_n = labelArray[(i-1)>0?(i-1):0][j];//label at (i-1,j)

			if (pixel==0) {labelArray[i][j] = 0;}
			else if (label_nw!=0) {labelArray[i][j] = label_nw;}
			else if ((label_w==0)&&(label_n==0)) {labelArray[i][j] = ++label;}
			else if ((label_w==0)&&(label_n!=0)) {labelArray[i][j] = label_n;}
			else if ((label_w!=0)&&(label_n==0)) {labelArray[i][j] = label_w;}
			else if (label_w==label_n) {labelArray[i][j] = label_n;}
			else {labelArray[i][j] = label_n; unionLabel(label_n,label_w);}
		}//end for
	}//end for

	int labelsFirstPass = label;
	printf("The number of labels after first pass is:\n%d \n",labelsFirstPass);

	/*======ADJUSTING THE EQUIVALENCE TABLE======*/
	for (int i = 1; i <= labelsFirstPass; i++)
	{
		eqTable[findLabel(i)].push_back(i);
	}
	
	/*======SECOND PASS======*/
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			int assignedLabel = *min_element( eqTable[findLabel(labelArray[i][j])].begin(), eqTable[findLabel(labelArray[i][j])].end() );
			//findLabel(labelArray[i][j]) is the root of the label in the UFS, 
			//eqTable[findLabel(labelArray[i][j])] is the UFS in which the label lies, 
			//min_element is used to find the minimum element in the UFS.
			//note that the assignedLabel here is not consecutive natural numbers,
			//we will then adjust them to finalLabel.

			vector<int>::iterator iter = tempLabel.begin();
			bool usedLabel = false;
			if (assignedLabel==0) 
			{
				finalLabel = 0;
			}
			else
			{
				finalLabel = 1;
				while (iter != tempLabel.end())
				{
					if (*iter == assignedLabel) {usedLabel = true; break;}	
					iter++;
					finalLabel++;
				}
			
				if (usedLabel == false) tempLabel.push_back(assignedLabel);
			}
			
			setPixel(im_output,
			i,
			j,
			finalLabel);
		}
	}

	setColors(im_output, int(tempLabel.size()));
	printf("The equivalence table is:\n");
	vector<int>::iterator tempLabelIter = tempLabel.begin();
	while (tempLabelIter!=tempLabel.end())
	{
		printf("%d: ",*tempLabelIter);
		vector<int>::iterator eqTableIter = eqTable[findLabel(*tempLabelIter)].begin();
		while (eqTableIter!=eqTable[findLabel(*tempLabelIter)].end())
		{
			printf("%d ",*eqTableIter);
			eqTableIter++;
		}
		printf("\n\n");
		tempLabelIter++;
	}
	printf("The number of labels after second pass is:\n%d \n",int(tempLabel.size()));

	writeImage(im_output, fname_output);
	printf("U got it\n");
	return 0;
}

int findLabel(int x)
{
	int p = x;
	while( pre[p] > 0 ) p = pre[p];//find the root of x
	while( x != p )
	{
		int temp = pre[x]; pre[x] = p; x = temp;//path compression
	}
	return x;
}

void unionLabel(int label1, int label2)
{
	int a = findLabel(label1); int b = findLabel(label2);
	if( a == b ) return ;
	if( pre[a] < pre[b] )
	{
		pre[a] += pre[b]; pre[b] = a;
	}
	else 
	{
		pre[b] += pre[a]; pre[a] = b;
	}
}

/*======MAIN======*/
int main(int argc, char **argv) 
{ 
	if (argc <= 2){ 
		printf("Usage: %s <input image>  <output image>\n", argv[0]); 
		return 0; 
	} 
	char *inputFilename = argv[1]; 
	char *outputFilename = argv[2]; 
	p2(inputFilename,outputFilename); 
	return 0; 
}
