#include "morphology.h"
#include "morphologicalflagger.h"
#include "siroperator.h"

#include "../../util/logger.h"

#include <stack>
#include <iostream>

size_t
	Morphology::BROADBAND_SEGMENT = 1,
	Morphology::LINE_SEGMENT = 2,
	Morphology::BLOB_SEGMENT = 3;

void Morphology::SegmentByMaxLength(const Mask2D* mask, SegmentedImagePtr output)
{
	int **lengthWidthValues = new int*[mask->Height()];
	for(size_t y=0;y<mask->Height();++y)
		lengthWidthValues[y] = new int[mask->Width()];
	
	calculateOpenings(mask, lengthWidthValues);

	for(size_t y=0;y<mask->Height();++y)
	{
		for(size_t x=0;x<mask->Width();++x)
			output->SetValue(x,y,0);
	}
	for(size_t y=0;y<mask->Height();++y)
	{
		for(size_t x=0;x<mask->Width();++x)
		{
			if(mask->Value(x, y) && output->Value(x,y) == 0)
			{
				floodFill(mask, output, lengthWidthValues, x, y, output->NewSegmentValue());
			}
		}
	}
		
	for(size_t y=0;y<mask->Height();++y)
		delete[] lengthWidthValues[y];
	delete[] lengthWidthValues;
}

void Morphology::SegmentByLengthRatio(const Mask2D* mask, SegmentedImagePtr output)
{
	Mask2DPtr maskCopy(new Mask2D(*mask));
	
	Mask2DPtr matrices[3];
	for(size_t i=0;i<3;++i)
		matrices[i] = Mask2D::CreateUnsetMaskPtr(mask->Width(), mask->Height());
	
	int
		**hCounts = new int*[mask->Height()],
		**vCounts = new int*[mask->Height()];
	for(size_t y=0;y<mask->Height();++y)
	{
		hCounts[y] = new int[mask->Width()];
		vCounts[y] = new int[mask->Width()];
	}
	
	// Calculate convolved counts
	calculateHorizontalCounts(maskCopy.get(), hCounts);
	calculateVerticalCounts(maskCopy.get(), vCounts);
	
	calculateOpenings(maskCopy.get(), matrices, hCounts, vCounts);


	for(size_t y=0;y<mask->Height();++y)
	{
		for(size_t x=0;x<mask->Width();++x)
			output->SetValue(x, y, 0);
	}
	MorphologicalFlagger::DilateFlags(matrices[0].get(), _hLineEnlarging, 0);
	MorphologicalFlagger::DilateFlags(matrices[2].get(), 0, _vLineEnlarging);
	SIROperator::OperateHorizontally(*matrices[0], _hDensityEnlargeRatio);
	SIROperator::OperateVertically(*matrices[2], _vDensityEnlargeRatio);

	// Calculate counts again with new matrices
	calculateHorizontalCounts(matrices[0].get(), hCounts);
	calculateVerticalCounts(matrices[2].get(), vCounts);

	for(size_t z=0;z<3;z+=2)
	{
		for(size_t y=0;y<mask->Height();++y)
		{
			for(size_t x=0;x<mask->Width();++x)
			{
				if(matrices[z]->Value(x, y) && output->Value(x, y)==0)
				{
					floodFill(mask, output, matrices, x, y, z, output->NewSegmentValue(), hCounts, vCounts);
				}
			}
		}
	}
	
	for(size_t y=0;y<mask->Height();++y)
	{
		delete[] hCounts[y];
		delete[] vCounts[y];
	}
	delete[] hCounts;
	delete[] vCounts;
}

void Morphology::calculateHorizontalCounts(const Mask2D* mask, int **values)
{
	for(size_t y=0;y<mask->Height();++y)
	{
		size_t length = 0;

		for(size_t x=0;x<mask->Width();++x)
		{
			if(mask->Value(x, y))
			{
				++length;
			} else if(length > 0) {
				for(size_t i=x-length;i<x;++i)
				{
					values[y][i] = length;
				}
				length = 0;
				values[y][x] = 0;
			} else {
				values[y][x] = 0;
			}
		}
		for(size_t i=mask->Width()-length;i<mask->Width();++i)
		{
			values[y][i] = length;
		}
	}
}

void Morphology::calculateVerticalCounts(const Mask2D* mask, int **values)
{
	for(size_t x=0;x<mask->Width();++x)
	{
		size_t length = 0;

		for(size_t y=0;y<mask->Height();++y)
		{
			if(mask->Value(x, y))
			{
				++length;
			} else if(length > 0) {
				for(size_t i=y-length;i<y;++i)
				{
					values[i][x] = length;
				}
				length = 0;
				values[y][x] = 0;
			} else {
				values[y][x] = 0;
			}
		}
		for(size_t i=mask->Height()-length;i<mask->Height();++i)
		{
			values[i][x] = length;
		}
	}
}

void Morphology::calculateOpenings(const Mask2D* mask, int **values)
{
	for(size_t y=0;y<mask->Height();++y)
	{
		size_t length = 0;

		for(size_t x=0;x<mask->Width();++x)
		{
			if(mask->Value(x, y))
			{
				++length;
			} else if(length > 0) {
				for(size_t i=x-length;i<x;++i)
				{
					values[y][i] = length;
				}
				length = 0;
				values[y][x] = 0;
			} else {
				values[y][x] = 0;
			}
		}
		if(length > 0) {
			for(size_t i=mask->Width()-length;i<mask->Width();++i)
			{
				values[y][i] = length;
			}
		}
	}

	for(size_t x=0;x<mask->Width();++x)
	{
		size_t length = 0;

		for(size_t y=0;y<mask->Height();++y)
		{
			if(mask->Value(x, y))
			{
				++length;
			} else if(length > 0) {
				for(size_t i=y-length;i<y;++i)
				{
					if(values[i][x] < (int) length)
						values[i][x] = -(int) length;
				}
				length = 0;
			}
		}
		if(length > 0) {
			for(size_t i=mask->Height()-length;i<mask->Height();++i)
			{
				if(values[i][x] < (int) length)
					values[i][x] = -(int) length;
			}
		}
	}
}

void Morphology::calculateOpenings(const Mask2D* mask, Mask2DPtr *values, int **hCounts, int **vCounts)
{
	//const int zThreshold = 5;
	
	for(size_t y=0;y<mask->Height();++y)
	{
		for(size_t x=0;x<mask->Width();++x)
		{
			bool v = mask->Value(x, y);
			values[0]->SetValue(x, y, v && (hCounts[y][x] > vCounts[y][x]));
			values[1]->SetValue(x, y, v && false);
			//values[1]->SetValue(x, y, v && (abs(hCounts[y][x] - vCounts[y][x]) < zThreshold));
			values[2]->SetValue(x, y, v && (hCounts[y][x] <= vCounts[y][x]));
		}
	}
}

struct MorphologyPoint2D { size_t x, y; };
struct MorphologyPoint3D { size_t x, y, z; };

void Morphology::floodFill(const Mask2D* mask, SegmentedImagePtr output, const int *const *lengthWidthValues, size_t x, size_t y, size_t value)
{
	std::stack<MorphologyPoint2D> points;
	MorphologyPoint2D startPoint;
	startPoint.x = x;
	startPoint.y = y;
	points.push(startPoint);
	do {
		MorphologyPoint2D p = points.top();
		points.pop();
		output->SetValue(p.x, p.y, value);
		int z = lengthWidthValues[p.y][p.x];
		if(p.x > 0 && output->Value(p.x-1, p.y) == 0 && mask->Value(p.x-1,p.y))
		{
			int zl = lengthWidthValues[p.y][p.x-1];
			if((zl > 0 && z > 0) || (zl < 0 && z < 0))
			{
				MorphologyPoint2D newP;
				newP.x = p.x-1; newP.y = p.y;
				points.push(newP);
			}
		}
		if(p.x < mask->Width()-1 && output->Value(p.x+1, p.y)==0 && mask->Value(p.x+1,p.y))
		{
			int zr = lengthWidthValues[p.y][p.x+1];
			if((zr > 0 && z > 0) || (zr < 0 && z < 0))
			{
				MorphologyPoint2D newP;
				newP.x = p.x+1; newP.y = p.y;
				points.push(newP);
			}
		}
		if(p.y > 0 && output->Value(p.x, p.y-1)==0 && mask->Value(p.x,p.y-1))
		{
			int zt = lengthWidthValues[p.y-1][p.x];
			if((zt > 0 && z > 0) || (zt < 0 && z < 0))
			{
				MorphologyPoint2D newP;
				newP.x = p.x; newP.y = p.y-1;
				points.push(newP);
			}
		}
		if(p.y < mask->Height()-1 && output->Value(p.x, p.y+1)==0 && mask->Value(p.x,p.y+1))
		{
			int zb = lengthWidthValues[p.y+1][p.x];
			if((zb > 0 && z > 0) || (zb < 0 && z < 0))
			{
				MorphologyPoint2D newP;
				newP.x = p.x; newP.y = p.y+1;
				points.push(newP);
			}
		}
	} while(points.size() != 0);
}

void Morphology::floodFill(const Mask2D* mask, SegmentedImagePtr output, Mask2DPtr *matrices, size_t x, size_t y, size_t z, size_t value, int **hCounts, int **vCounts)
{
	std::stack<MorphologyPoint3D> points;
	MorphologyPoint3D startPoint;
	startPoint.x = x;
	startPoint.y = y;
	startPoint.z = z;
	points.push(startPoint);
	do {
		MorphologyPoint3D p = points.top();
		points.pop();
		if(mask->Value(p.x, p.y))
		{
			if(output->Value(p.x, p.y) == 0)
			{
				output->SetValue(p.x, p.y, value);
			} else {
				// now we need to decide whether to change this sample to the new segment or not
				if(hCounts[p.y][p.x] < vCounts[p.y][p.x] && p.z == 2)
					output->SetValue(p.x, p.y, value);
			}
		}
		Mask2DPtr matrix = matrices[p.z];
		matrix->SetValue(p.x, p.y, false);
		if((p.z == 0 || p.z == 2) && matrices[1]->Value(p.x,p.y))
		{
			MorphologyPoint3D newP;
			newP.x = p.x; newP.y = p.y; newP.z = 1;
			points.push(newP);
		}
		if(p.x > 0 && matrix->Value(p.x-1,p.y))
		{
			MorphologyPoint3D newP;
			newP.x = p.x-1; newP.y = p.y; newP.z = p.z;
			points.push(newP);
		}
		if(p.x < mask->Width()-1 && matrix->Value(p.x+1,p.y))
		{
			MorphologyPoint3D newP;
			newP.x = p.x+1; newP.y = p.y; newP.z = p.z;
			points.push(newP);
		}
		if(p.y > 0 && matrix->Value(p.x,p.y-1))
		{
			MorphologyPoint3D newP;
			newP.x = p.x; newP.y = p.y-1; newP.z = p.z;
			points.push(newP);
		}
		if(p.y < mask->Height()-1 && matrix->Value(p.x,p.y+1))
		{
			MorphologyPoint3D newP;
			newP.x = p.x; newP.y = p.y+1; newP.z = p.z;
			points.push(newP);
		}
	} while(points.size() != 0);
}

void Morphology::Cluster(SegmentedImagePtr segmentedImage)
{
	std::map<size_t,SegmentInfo> segments = createSegmentMap(segmentedImage);
	Logger::Debug << "Segments before clustering: " << segments.size();

	for(std::map<size_t,SegmentInfo>::iterator i=segments.begin();i!=segments.end();++i)
	{
		SegmentInfo &info1 = i->second;
		for(std::map<size_t,SegmentInfo>::iterator j=segments.begin();j!=segments.end();++j)
		{
			if(info1.segment != j->second.segment && !(i->second.mark) && !(j->second.mark))
			{
				SegmentInfo &info2 = j->second;
				size_t hDist = info1.HorizontalDistance(info2);
				size_t vDist = info1.VerticalDistance(info2);
	
				// The MERGE criteria
				bool cluster = false;
				//int minDist = hDist > vDist ? vDist : hDist;
				int maxDist = hDist > vDist ? hDist : vDist;
				//int maxCount = info1.count > info2.count ? info1.count : info2.count;
				int minCount = info1.count > info2.count ? info2.count : info1.count;
				int maxWidth = info1.width > info2.width ? info1.width : info2.width;
				int maxHeight = info1.height > info2.height ? info1.height : info2.height;
				int minHeight = info1.height > info2.height ? info2.height : info1.height;
				//int lDist = abs((int) info1.left - (int) info2.left);
				//int rDist = abs((int) info1.right - (int) info2.right);
				//int tDist = abs((int) info1.top - (int) info2.top);
				//int bDist = abs((int) info1.bottom - (int) info2.bottom);
				int widthDist = abs((int) info1.width - (int) info2.width);
				int heightDist = abs((int) info1.height - (int) info2.height);
				//double x1Mean = (double) info1.xTotal / info1.count; 
				//double x2Mean = (double) info2.xTotal / info2.count;
				//double xMeanDist = fabs(x1Mean - x2Mean); 
				double y1Mean = (double) info1.yTotal / info1.count; 
				double y2Mean = (double) info2.yTotal / info2.count;
				double yMeanDist = fabs(y1Mean - y2Mean); 
	
				bool remove1 = false, remove2 = false;
				
				// Cluster large segments with very small segments that are close together (probably noise
				// from the continuous transmitter)
				bool
					noiseH1 = maxDist <= 1 && info2.count > (info1.count*20) && info2.width>info1.width*8 && info1.height < 16 && info1.width<segmentedImage->Width()/10,
					noiseH2 =  maxDist <= 1 && info1.count > (info2.count*20) && info1.width>info2.width*8 && info2.height < 16 && info2.width<segmentedImage->Width()/10;
				cluster = cluster || noiseH1 || noiseH2;
				remove1 = remove1 || noiseH1;
				remove2 = remove2 || noiseH2;

				bool
					noiseV1 = maxDist <= 1 && info2.count > (info1.count*20) && info2.height>info1.height*8 && info1.height < 16 && info1.width<segmentedImage->Width()/10,
					noiseV2 =  maxDist <= 1 && info1.count > (info2.count*20) && info1.height>info2.height*8 && info2.height < 16 && info2.width<segmentedImage->Width()/10;
				cluster = cluster || noiseV1 || noiseV2;
				remove1 = remove1 || noiseV1;
				remove2 = remove2 || noiseV2;
	
				// Cluster same-shaped segments that are in the same channels
				cluster = cluster || (vDist == 0 && yMeanDist*8 <= (maxHeight+minHeight) && widthDist <= (maxWidth / 4 + 2) && heightDist <= (maxHeight / 4 + 2) && maxDist < minCount*32);

				if(cluster)
				{
					size_t oldSegment = info2.segment;
					segmentedImage->MergeSegments(info1.segment, oldSegment);
					for(std::map<size_t,SegmentInfo>::iterator i=segments.begin();i!=segments.end();++i)
					{
						SegmentInfo &info = i->second;
						if(info.segment == oldSegment) info.segment=info1.segment;
					}
				}
				if(remove1)
					info1.mark = true;
				if(remove2)
					info2.mark = true;
			}
		}
	}
}

std::map<size_t,Morphology::SegmentInfo> Morphology::createSegmentMap(SegmentedImageCPtr segmentedImage) const
{
	std::map<size_t,SegmentInfo> segments;
	for(size_t y=0;y<segmentedImage->Height();++y)
	{
		for(size_t x=0;x<segmentedImage->Width();++x)
		{
			size_t segmentValue = segmentedImage->Value(x,y);
			if(segmentValue != 0)
			{
				if(segments.count(segmentValue) == 0)
				{
					SegmentInfo segment;
					segment.segment = segmentValue;
					segment.left = x;
					segment.right = x+1;
					segment.top = y;
					segment.bottom = y+1;
					segment.AddPoint(x,y);
					segments.insert(std::map<size_t,SegmentInfo>::value_type(segmentValue, segment));
				} else {
					SegmentInfo &segment = segments.find(segmentValue)->second;
					segment.AddPoint(x,y);
				}
			}
		}
	}

	for(std::map<size_t,SegmentInfo>::iterator i=segments.begin();i!=segments.end();++i)
	{
		SegmentInfo &info = i->second;
		info.width = info.right - info.left;
		info.height = info.bottom - info.top;
	}
	return segments;
}

void Morphology::RemoveSmallSegments(SegmentedImagePtr segmentedImage, size_t thresholdLevel)
{
	std::map<size_t,SegmentInfo> segments = createSegmentMap(segmentedImage);
	size_t removedSegments = 0;

	for(std::map<size_t,SegmentInfo>::iterator i=segments.begin();i!=segments.end();++i)
	{
		const SegmentInfo &segment = i->second;
		if(segment.count <= thresholdLevel)
		{
			++removedSegments;
			segmentedImage->RemoveSegment(segment.segment, segment.left, segment.right, segment.top, segment.bottom);
		}
	}
	Logger::Debug << "Removed " << removedSegments << " segments of size " << thresholdLevel << " or smaller.\n";
}

void Morphology::Classify(SegmentedImagePtr segmentedImage)
{
	std::map<size_t,SegmentInfo> segments = createSegmentMap(segmentedImage);

	for(std::map<size_t,SegmentInfo>::iterator i=segments.begin();i!=segments.end();++i)
	{
		SegmentInfo &info = i->second;
		if(info.width > info.height * 10)
			segmentedImage->MergeSegments(LINE_SEGMENT, info.segment);
		else if(info.height > info.width * 10)
			segmentedImage->MergeSegments(BROADBAND_SEGMENT, info.segment);
		else
			segmentedImage->MergeSegments(BLOB_SEGMENT, info.segment);
	}
}
