//  Segmentation.cpp -- simple image segmentation using anisotropic diffusion
//
//
// DESCRIPTION
//
// Simple color segmenter, based on Larry Zitnik's description.  It has two stages:
//
// 1. Anisotropic diffusion to sharpen the edges and reduce the noise:
// Average each pixel with its most similar 3 (of 8) consecutive
// neighbors, unless average difference of all 8 neighbors is within 
// 'difThresh', in which case average with all 8.  Repeat 'iter' times.
//
// 2. Using threshold 'segThresh' for (Euclidean) color distance, segment the image
// using the "union find" algorithm (similar to CX 336 HW 2).  Keep a
// table of segments that contain the average color and the number of
// pixels in the segment.  Go through the image in scanline order.  For
// each new pixel, compare it to the segments of the pixel above and the
// pixel to the left.  If it's too different from both, create a new
// segment.  If it's similar enough to one but not the other, add it to
// that.  If it's similar enough to both, but the segments themself are
// not similar enough, add it to the closer one.  If the segments and the
// pixels are all similar enough, merge all three (using a "union tree"
// to record that the two segments are now considered one).
//
// Copyright © Daniel Scharstein and Timothy Bahls 2004.
//
///////////////////////////////////////////////////////////////////////////

#include <vector>
#include <stdlib.h>
#include <math.h>
#include <assert.h>

#include "Image.h"
#include "ImageIO.h"
#include "Error.h"
#include "Color.h"
#include "Convert.h"
#include "SegXYT.h"
//#include "ShowImage.h"
#include "Utils.h"

// defined in Main.cpp
extern bool verbose;
extern bool visualize;


// anisotropically diffuse dif into dif2
void aniDiffImage(CByteImage *difPtr, CByteImage *dif2Ptr, int difThresh, int nFrames)
{
	difThresh = 8 * difThresh * difThresh; // make threshold compatible with squared diff's
	CShape sh = difPtr[0].Shape();
	//dif2.ReAllocate(sh);
	int w = sh.width, h = sh.height, nB = sh.nBands, x, y, b, k;
	assert(nB <= 4);
	int B = __min(nB, 3);
	// diffuse interior
	for(int f=0; f<nFrames; f++){
		CByteImage dif=difPtr[f];
		CByteImage dif2=difPtr[f];
		for (y = 1; y < h-1; y++) {
			uchar *sm = (uchar *)dif.PixelAddress(1, y-1, 0);
			uchar *s  = (uchar *)dif.PixelAddress(1, y,   0);
			uchar *sp = (uchar *)dif.PixelAddress(1, y+1, 0);
			uchar *d  = (uchar *)dif2.PixelAddress(1, y,   0);
			for (x = 1; x < w-1; x++, sm += nB, s += nB, sp += nB, d += nB) {
				int a[10][3];
				for (b = 0; b < B; b++) {
					a[0][b] = sm[b-nB]; 
					a[1][b] = sm[b]; 
					a[2][b] = sm[b+nB]; 
					a[3][b] = s[b+nB]; 
					a[4][b] = sp[b+nB]; 
					a[5][b] = sp[b]; 
					a[6][b] = sp[b-nB]; 
					a[7][b] = s[b-nB]; 
					a[8][b] = a[0][b];
					a[9][b] = a[1][b];
				}
				int diff[10];
				int totdiff = 0;		// total difference of all 8 neighbors
				// ************* TODO: should be MAX, not sum??
				for (k = 0; k < 8; k++) {
					diff[k] = 0;
					for (b = 0; b < B; b++) {
						int di = a[k][b] - s[b];
						diff[k] += di*di;		// SQUARED DIFF
						//diff[k] += abs(di);	// ABS DIFF
					}
					totdiff += diff[k];
				}
				if (totdiff <= difThresh) { // all neighbors close: average all together
					for (b = 0; b < B; b++) {
						int db = s[b];
						for (k = 0; k < 8; k++) 
							db += a[k][b];
						d[b] = db / 9;
					}
				} else { // find 3 most similar consecutive neighbors and average those
					diff[8] = diff[0];
					diff[9] = diff[1];
					int cursum[3], bestsum[3], curdif, bestdif;
					// sum first three (0-2)
					curdif = bestdif = diff[0] + diff[1] + diff[2];
					for (b = 0; b < B; b++)
						cursum[b] = bestsum[b] = a[0][b] + a[1][b] + a[2][b];
					// sum remaining 7 triplets (1-3 .. 7-9) by adding next, subtracting prev
					for (int k=0; k<7; k++) {
						curdif += (diff[k+3] - diff[k]);
						for (b = 0; b < B; b++)
								cursum[b] += (a[k+3][b] - a[k][b]);
						if (curdif < bestdif) {
							
							bestdif = curdif;
							for (b = 0; b < B; b++)
								bestsum[b] = cursum[b];
						}
					}
					// now form average of center pixel and best three consecutive neighbors
					for (b = 0; b < B; b++) {
						d[b] = (bestsum[b] + s[b]) / 4;
					}
				}
			}
		}

		// handle border
		// for now, just use 75% of intensity of interior pixel next to it
		for (x = 0; x < w; x++) {
			for (b = 0; b < B; b++) {
				dif2.Pixel(x, 0,   b) = (dif.Pixel(x, 0,   b) + 3*dif.Pixel(x, 1,   b))/4;
				dif2.Pixel(x, h-1, b) = (dif.Pixel(x, h-1, b) + 3*dif.Pixel(x, h-2, b))/4;
			}
		}
		for (y = 0; y < h; y++) {
			for (b = 0; b < B; b++) {
				dif2.Pixel(0,   y, b) = (dif.Pixel(0,   y, b) + 3*dif.Pixel(1,   y, b))/4;
				dif2.Pixel(w-1, y, b) = (dif.Pixel(w-1, y, b) + 3*dif.Pixel(w-2, y, b))/4;
			}
		}
	}
}



float similarity(float *a, float *b) {
	float d0 = a[0] - b[0];
	float d1 = a[1] - b[1];
	float d2 = a[2] - b[2];
	return (float)sqrt(d0*d0 + d1*d1 + d2*d2);
}

float similarity(float *a, uchar *u) {
	float b[3] = {(float)u[0], (float)u[1], (float)u[2]};
	return similarity(a, b);
}


void update_seg(std::vector<CSegXYT> &segment, int &s) {
	int s0 = s;
	while (s != segment[s].root) {
		s = segment[s].root;
		segment[s0].root = s;
	}
}

#define DEBUGSEG 0

// segment image
void segmentImage(CByteImage *difPtr, std::vector<CSegXYT> &segment, CIntImage *&segnumPtr,
				  int thresh, int nFrames)
{
	CShape sh = difPtr[0].Shape();
	if (sh.nBands < 4)
		throw CError("Only color images supported");
	sh.nBands = 1;
	int w = sh.width, h = sh.height, x, y, b;
	segment.reserve(1000);
	CSegXYT s0(0, 1000, 1000, 1000, 0); // dummy elt with very different color
	segment.push_back(s0);
	int curseg = 0;
	assert(segment.size()-1 == curseg);
	int nsegments = 0;

	for(int f=0; f<nFrames;f++){
		CIntImage segnum=segnumPtr[f];
		CByteImage dif=difPtr[f];
//segnum.ReAllocate(sh);
		segnumPtr[f].ReAllocate(sh);
		for (y = 0; y < h; y++) {
			for (x = 0; x < w; x++) {
				
				int s1 = (y > 0) ? segnumPtr[f].Pixel(x, y-1, 0) : 0;
				int s2 = (x > 0) ? segnumPtr[f].Pixel(x-1, y, 0) : 0;
				update_seg(segment, s1);
				update_seg(segment, s2);
				assert(s1 >= 0 && s1 <= curseg);
				assert(s2 >= 0 && s2 <= curseg);
				uchar *pixcol = (uchar *)dif.PixelAddress(x, y, 0);
				float d1, d2;
				d1 = similarity(segment[s1].col, pixcol);
				if (s1 == s2) { // same segment: no need to consider s2
					d2 = thresh + 1.0f;
				} else {
					d2 = similarity(segment[s2].col, pixcol);
					// make sure we're not merging two segments that are not similar enough
					float d12 = similarity(segment[s1].col, segment[s2].col);
					if (d12 > thresh) { // don't merge, but make sure pixel sticks with closer
						if (d1 < d2)
							d2 = thresh + 1.0f;
						else
							d1 = thresh + 1.0f;
					}
				}
				if (d1 <= thresh) {
					if (d2 <= thresh) {
						// both similar to pixel and to each other: merge!
						// union s1, s2, and pixel; make s1 new root
						nsegments--;
						int ns1 = segment[s1].nPix;
						int ns2 = segment[s2].nPix;
						for (b = 0; b < 3; b++)
							segment[s1].col[b] = 
							(segment[s1].col[b] * ns1 + 
							segment[s2].col[b] * ns2 + pixcol[b]) / (ns1 + ns2 + 1.0f);
						segment[s1].nPix += (ns2 + 1);
						segment[s2].root = s1;
						segnumPtr[f].Pixel(x, y, 0) = s1;

						if (DEBUGSEG) {
							printf("(%d,%d)->segs %d+%d (%d %d %d)", x, y, s1, s2,
								(int)segment[s1].col[0], (int)segment[s1].col[1], (int)segment[s1].col[2]);
							printf("      ns1=%d ns2=%d\n", ns1, ns2);
						}
					} else {
						// add pixel to s1
						int ns = segment[s1].nPix;
						for (b = 0; b < 3; b++)
							segment[s1].col[b] = (segment[s1].col[b] * ns + pixcol[b]) / (ns+1.0f);
						segment[s1].nPix++;
						segnumPtr[f].Pixel(x, y, 0) = s1;
						if (DEBUGSEG) {
							printf("(%d,%d)->seg1 %d (%d %d %d)", x, y, s1,
								(int)segment[s1].col[0], (int)segment[s1].col[1], (int)segment[s1].col[2]);
							printf("      ns=%d\n", ns);
						}
					}
				} else {
					if (d2 <= thresh) {
						// add pixel to s2
						int ns = segment[s2].nPix;
						for (b = 0; b < 3; b++)
							segment[s2].col[b] = (segment[s2].col[b] * ns + pixcol[b]) / (ns+1.0f);
						segment[s2].nPix++;
						segnumPtr[f].Pixel(x, y, 0) = s2;
						if (DEBUGSEG) {
							printf("(%d,%d)->seg2 %d (%d %d %d)", x, y, s2,
								(int)segment[s2].col[0], (int)segment[s2].col[1], (int)segment[s2].col[2]);
							printf("      ns=%d\n", ns);
						}
					} else {
						// start new segment
						curseg++;
						nsegments++;
						CSegXYT s(1, pixcol[0], pixcol[1], pixcol[2], curseg);
						segment.push_back(s);
						segnumPtr[f].Pixel(x, y, 0) = curseg;
						assert(segment[curseg].root == curseg);
						if (DEBUGSEG) {
							printf("(%d,%d)->new seg %d (%d %d %d)\n", x, y, curseg,
								(int)segment[curseg].col[0], (int)segment[curseg].col[1], (int)segment[curseg].col[2]);
						}
					}
				}
				
				//update and check similarity
				int s0=segnumPtr[f].Pixel(x, y, 0);
				int s3 = (f > 0) ? segnumPtr[f-1].Pixel(x, y, 0) : 0; 
				update_seg(segment, s3);
				update_seg(segment, s0);

				float d03 = similarity(segment[s0].col, segment[s3].col);
				//if(x<50&&f>0)
					//printf("Similar=%f\n",d3);


				



				if(d03<=thresh&&s0!=s3) {

					


					nsegments--;
					int ns0 = segment[s0].nPix;
					int ns3 = segment[s3].nPix;
					for (b = 0; b < 3; b++){
						segment[s3].col[b] =
							segment[s0].col[b]=
						(segment[s0].col[b] * ns0 + 
						segment[s3].col[b] * ns3) / (ns0 + ns3);
					}
					segment[s3].nPix += ns0;
					segment[s0].root = s3;
					segnumPtr[f].Pixel(x, y, 0) = s3;
					//printf("Yes f=%d x=%d y=%d\n",f,x,y);
				}else{
					//printf("No f=%d x=%d y=%d\n",f,x,y);
				}

			}
			
		}
	}


}

void update_bbox(CSegXYT &seg, int x, int y)
{
	seg.xmin = __min(seg.xmin, x);
	seg.xmax = __max(seg.xmax, x);
	seg.ymin = __min(seg.ymin, y);
	seg.ymax = __max(seg.ymax, y);
}

// renumber segments to remove those that got merged from segment vector
void compactSegments(std::vector<CSegXYT> &segment, CIntImage *segnumPtr, 
					 int startindex, bool verbose, int nFrames)
{

	CShape sh = segnumPtr[0].Shape();
	int w = sh.width, h = sh.height, i;

	// clean up segment numbers
	std::vector<CSegXYT> newseg;
	int snum = 0;

	for (i = startindex; i < (int)segment.size(); i++) {
		int k = i; // update_seg modifies k!
		update_seg(segment, k);
		if (segment[i].root == i) {
			segment[i].num = snum;
			newseg.push_back(segment[i]);
			newseg[snum].root = snum;
			// initialize bounding box to be empty
			newseg[snum].xmin = w;
			newseg[snum].xmax = -1;
			newseg[snum].ymin = h;
			newseg[snum].ymax = -1;
			snum++;
		}
	}

	// update segment number image
for(int f=0; f<nFrames;f++){
		CIntImage segnum=segnumPtr[f];
	for (int y = 0; y < h; y++) {
		int *sn = (int *)segnum.PixelAddress(0, y, 0);
		for (int x = 0; x < w; x++) {
			int snn=segnum.Pixel(x, y, 0);
			CSegXYT seg=segment[snn];
			int r=seg.root;



			//int r = segment[sn[x]].root;
			update_bbox(segment[r], x, y);
			sn[x] = segment[r].num;
		}
	}
}
	// copy new segment vector back
	segment.resize(snum);
	for (i = 0; i < snum; i++) {
		segment[i] = newseg[i];
	}

	if (verbose)
		fprintf(stderr, "%d segments found\n", snum-1);//disregard dummy segment
}


void mergeSimilarSegments(CByteImage *difPtr, std::vector<CSegXYT> &segment, CIntImage *segnumPtr, 
						  int mergeThresh, int nFrames, int minSize)
{

	CShape sh = difPtr[0].Shape();
	int w = sh.width, h = sh.height, nB = sh.nBands, x, y, b;
	int nseg = (int)segment.size();
	CShape sh2(nseg, nseg, 1);
	CIntImage bordercount(sh2);
	CIntImage borderssd(sh2);
	bordercount.FillPixels(0);
	borderssd.FillPixels(0);
	// first step: consider all pairs of adjacent pixels
	// if they lie in different segments, increase the segments' bordercount and borderssd
	for(int f=0; f<nFrames;f++){
		CIntImage segnum=segnumPtr[f];
		CByteImage dif=difPtr[f];
		for (y = 0; y < h-1; y++) {
		int *sn0 = (int *)segnum.PixelAddress(0, y, 0);
		int *sn1 = (int *)segnum.PixelAddress(0, y+1, 0);
		uchar *d0 = (uchar *)dif.PixelAddress(0, y, 0);
		uchar *d1 = (uchar *)dif.PixelAddress(0, y+1, 0);
		for (x = 0; x < w; x++, d0 += nB, d1 += nB) {
			
			// vertical difference
			int s0 = sn0[x];
			int s1 = sn1[x];
			if (s0 != s1) {
				bordercount.Pixel(s0, s1, 0)++;
				bordercount.Pixel(s1, s0, 0)++;
				//printf("at line 4\n");
				int ssd = 0;
				for (b = 0; b < nB; b++) {
					int d = (int)d0[b] - (int)d1[b];
					ssd += d*d;
				}
				borderssd.Pixel(s0, s1, 0) += ssd;
				borderssd.Pixel(s1, s0, 0) += ssd;
			}
			// horizontal difference
			if(x>=w-1)
				goto time;
			s0 = sn0[x];
			s1 = sn0[x+1];
			
			if (s0 != s1) {
				
				bordercount.Pixel(s0, s1, 0)++;
				bordercount.Pixel(s1, s0, 0)++;
				int ssd = 0;
				for (b = 0; b < 3; b++) {
					int d = (int)d0[b] - (int)d0[b + nB];
					ssd += d*d;
				}
				borderssd.Pixel(s0, s1, 0) += ssd;
				borderssd.Pixel(s1, s0, 0) += ssd;
			}
time:
			if(f>=nFrames-1)
				continue;
			s0 = sn0[x];
			s1 = segnumPtr[f+1].Pixel(x,y,0);
			uchar *d2= &difPtr[f+1].Pixel(x,y,0);
			if (s0 != s1) {
				
				bordercount.Pixel(s0, s1, 0)++;
				bordercount.Pixel(s1, s0, 0)++;
				int ssd = 0;
				for (b = 0; b < 3; b++) {
					int d = (int)d0[b] - (int)d2[b];
					ssd += d*d;
				}
				borderssd.Pixel(s0, s1, 0) += ssd;
				borderssd.Pixel(s1, s0, 0) += ssd;
			}

			



		}
		
	}
	}

	//second step: merge small polygons to whatever is the closest segment


	for (int k1 = 1; k1 < nseg; k1++) {
		int s1 = k1;
		update_seg(segment, s1); // find root of segment
		int ns1 = segment[s1].nPix;
		if(ns1>minSize)//segment big enough on its own right
			continue;
		float minThresh=99999.9f;//really big
		int threshIndex=-1;
		for (int k0 = 1; k0 < nseg; k0++) {
			int bcnt = bordercount.Pixel(k0, k1, 0);
			if (bcnt == 0) // segments don't share a border
				continue;
			int s0 = k0;
			update_seg(segment, s0);// find root of segment
			if (s0 == s1)	// already merged earlier
				continue;
			int ns0 = segment[s0].nPix;
			if(ns0<=minSize)//too small to merge to
				continue;
			int bssd = borderssd.Pixel(k0, k1, 0);
			float thresh= ((float)bssd) / ((float)bcnt);
			if(thresh>minThresh)//already found better match
				continue;
			minThresh=thresh;
			threshIndex=k0;
		}
		if(threshIndex!=-1){
			int ns0 = segment[threshIndex].nPix;
			for (b = 0; b < 3; b++)
				segment[threshIndex].col[b] = (segment[threshIndex].col[b] * ns0 + segment[s1].col[b] * ns1)
									 / (ns0 + ns1);
			segment[threshIndex].nPix += ns1;
			segment[s1].root = threshIndex;
			if (0)
				printf(" merging segs %d, %d (%d %d %d) ns0=%d ns1=%d\n", threshIndex, s1,
					(int)segment[threshIndex].col[0], (int)segment[threshIndex].col[1], (int)segment[threshIndex].col[2],
					threshIndex, ns1);
		}else{
			printf("Not merged!\n");
		}
	}

	// third step: if the average intensity differences accross border are small
	// then merge segments
	mergeThresh *= mergeThresh; // square threshold to make compatible with SSD
	for (int k1 = 1; k1 < nseg; k1++) {
		for (int k0 = 0; k0 < k1; k0++) { // only look at ordered pairs
			int bcnt = bordercount.Pixel(k0, k1, 0);
			if (bcnt == 0) // segments don't share a border
				continue;
			int bssd = borderssd.Pixel(k0, k1, 0);
			float currThresh=(float)mergeThresh;
			

			
			int s0 = k0;
			int s1 = k1;
			update_seg(segment, s0); // find root of each segment
			update_seg(segment, s1);
			if (s0 == s1)	// already merged earlier
				continue;
			int ns0 = segment[s0].nPix;
			int ns1 = segment[s1].nPix;
			int minSize=__min(ns0,ns1);
			if(minSize<100&&minSize<(bcnt*6)){
				currThresh=(bssd /bcnt)+1.0f;
			}
			if (bssd / bcnt > (int)currThresh)
				continue; // segments too different

			for (b = 0; b < 3; b++)
				segment[s0].col[b] = (segment[s0].col[b] * ns0 + segment[s1].col[b] * ns1)
									 / (ns0 + ns1);
			segment[s0].nPix += ns1;
			segment[s1].root = s0;
			if (0)
				printf(" merging segs %d, %d (%d %d %d) ns0=%d ns1=%d\n", s0, s1,
					(int)segment[s0].col[0], (int)segment[s0].col[1], (int)segment[s0].col[2],
					ns0, ns1);
		}
	}
}


// create color image from segments
void colorSegments(std::vector<CSegXYT> &segment, CIntImage *segnumPtr, CByteImage *colsegPtr, 
				   int colorByIndex, int nFrames)
{
	CShape sh = segnumPtr[0].Shape();
	int w = sh.width, h = sh.height;
	sh.nBands = 4;
	

for(int f=0; f<nFrames;f++){

	if(colsegPtr[f].Shape()!=sh){
		printf("Necesary!");
		colsegPtr[f].ReAllocate(sh);
	}


		CIntImage segnum=segnumPtr[f];
		CByteImage colseg=colsegPtr[f];

	for (int y = 0; y < h; y++) {
		for (int x = 0; x < w; x++) {
			int s = segnum.Pixel(x, y, 0);
			// don't need these two lines if we always run "compactSegments" first
			//update_seg(segment, s);
			//segnum.Pixel(x, y, 0) = s;
			if (colorByIndex) { // use color index
				CCol *cs = (CCol *)colseg.PixelAddress(x, y, 0);
				cs[0] = CCol(s+1);
			} else { // use average color
				uchar *cs = (uchar *)colseg.PixelAddress(x, y, 0);

				float *col = segment[s].col;
				for (int b = 0; b < 3; b++)
					cs[b] = (int)col[b];
				cs[3] = 255; // full alpha
			}
		}
	}
}
}


void diffuseAndSegmentImage(CByteImage *im, CByteImage *colseg, int iters, int difThresh, 
							int segThresh, int mergeThresh, int colorByIndex, int nFrames, int minSize, int output)
{
	CShape sh = im[0].Shape();
	CByteImage *dif=new CByteImage[nFrames];
	CByteImage *dif2=new CByteImage[nFrames];
	for(int f=0; f<nFrames; f++){
		dif[f]=CByteImage(sh);
		dif2[f]=CByteImage(sh);
		CopyPixels(im[f], dif[f]);
		colseg[f].ReAllocate(sh);
	}

	// anisotropic diffusion
	if (iters > 0) {
		if (verbose) fprintf(stderr, "diffusing %d times with difThresh %d", iters, difThresh);
		for (int i = 0; i < iters; i += 2) {
			aniDiffImage(dif, dif2, i>0 ? difThresh : 0, nFrames); // don't use 8-averaging in first iter
			if (verbose) fprintf(stderr, ".");
			aniDiffImage(dif2, dif, difThresh, nFrames);
			if (verbose) fprintf(stderr, ".");
		}
		if (verbose) fprintf(stderr, "\n");
		if(output){
		for(int f=0; f<nFrames; f++){
			char fname[200];
			sprintf(fname, "anidiff%02d.ppm", f);
			WriteImage(dif[f], fname);	
		}
		}
	}
	// actual segmentation
	CIntImage * segnum=new CIntImage[nFrames];
	for(int f=0; f<nFrames; f++){
		segnum[f].ReAllocate(sh);
	}
	std::vector<CSegXYT> segment;
	if (verbose) fprintf(stderr, "segmenting, segThres = %d\n", segThresh);
	segmentImage(dif, segment, segnum, segThresh, nFrames);
	printf("Segmenting complete\n");
	

	// renumber segments
	int startindex = 1; // skip dummy segment 0
	compactSegments(segment, segnum, false, verbose, nFrames);
	colorSegments(segment, segnum, colseg, colorByIndex,nFrames);
	if(output){
		for(int f=0; f<nFrames; f++){
			char fname[200];
			sprintf(fname, "segmentPre%02d.ppm", f);
			WriteImage(colseg[f], fname);	
		}
	}

	// merge similar segments based on average color difference on their border
	if (verbose) fprintf(stderr, "merging, mergeThres = %d\n", mergeThresh);
	mergeSimilarSegments(dif, segment, segnum, mergeThresh, nFrames, minSize);

	// renumber segments again
    startindex = 0; // this time, start at 0
	compactSegments(segment, segnum, startindex, verbose,nFrames);

	// color segments
	colorSegments(segment, segnum, colseg, colorByIndex,nFrames);
}