#include <i3d/image3d.h>
#include <i3d/regions.h>
#include <i3d/morphology.h>
#include <set>


template <class VT>
bool reduceMarkers(const i3d::Image3d<VT>& traIn, i3d::Image3d<VT>& traOut, const int erosionSize)
{
	//find set of labels
	std::set<int> setOfMarkers;
	std::map<int,int> sizesOfMarkers;
	const VT* inP = traIn.GetFirstVoxelAddr();
	const VT* const inPL = inP + traIn.GetImageSize();
	while (inP != inPL)
	{
		if (*inP > 0)
		{
			setOfMarkers.insert(*inP);
			++sizesOfMarkers[*inP];
		}
		++inP;
	}

	//erode the whole image
	i3d::ErosionO(traIn,traOut,erosionSize);

	//image for individual labels/markers
	i3d::Image3d<VT> tmpImg;
	tmpImg.CopyMetaData(traOut);
	//
	//for the CCA around it
	i3d::LabeledImage3d<size_t,i3d::GRAY16> lblImg;

	//running pointers...
	VT* tmpP, * outP;
	VT* const outPL = traOut.GetFirstVoxelAddr() + traOut.GetImageSize();

	//process label by label and keep only its largest connected component
	for (int label : setOfMarkers)
	{
		//copy out the label from traOut, CCA on it, remove non-largest CC in the traOut image
		//std::cout << "found label: " << label << " (volume=" << sizesOfMarkers[label] << ")\n";
		if (sizesOfMarkers[label] < 1000)
		{
			std::cout << "label " << label << ": too small (" << sizesOfMarkers[label] << ") to be reduced\n";
			//have to restore it again (because the out image is already eroded)
			inP  = traIn.GetFirstVoxelAddr();
			outP = traOut.GetFirstVoxelAddr();
			while (inP != inPL)
			{
				*outP = *inP == label ? label : *outP;
				++inP; ++outP;
			}
			continue;
		}

		int originalLabelAfterErosionSize = 0;

		//copy out the single label from traOut to tmpImg
		tmpP = tmpImg.GetFirstVoxelAddr();
		outP = traOut.GetFirstVoxelAddr();
		while (outP != outPL)
		{
			if (*outP == label)
			{
				*tmpP = 1;
				++originalLabelAfterErosionSize;
			}
			else *tmpP = 0;
			++outP; ++tmpP;
		}

		//CCA on it
		lblImg.CreateForegroundRegionsFF(tmpImg);

		//find largest CC
		int largestSize = 0;
		size_t largestLabel = 0;
		for (const auto& c : lblImg.components)
		{
			//std::cout << "considering comp: " << c.first << " of volume " << c.second.volume << "\n";
			if (c.second.volume > largestSize)
			{
				largestSize = c.second.volume;
				largestLabel = c.first;
			}
		}

		//DEBUG
		std::cout << "label " << label << ": "
		          << largestSize << " px is the largest component (of label "
		          << largestLabel << " from " << lblImg.components.size() << " components), ";

		//remove non-largest CC in the traOut image
		int howMuchWasRemoved = 0;
		size_t* lblP = lblImg.GetFirstVoxelAddr();
		outP = traOut.GetFirstVoxelAddr();
		while (outP != outPL)
		{
			//in the territory of the original marker but not in its largest component?
			if (*outP == label && *lblP != largestLabel)
			{
				*outP = 0;
				++howMuchWasRemoved;
			}
			++lblP; ++outP;
		}

		//break the whole thing if we experience a label that got diminished by this erosion
		int howMuchIsLeft = originalLabelAfterErosionSize - howMuchWasRemoved;
		if (howMuchIsLeft != lblImg.components[largestLabel].volume) std::cout << "ERROR! ";
		std::cout << howMuchWasRemoved << " px was removed, and " << howMuchIsLeft << " px is left\n";
		if (howMuchIsLeft < 9) return false;
	}

	return true;
}


int main(int argc, char** argv)
{
	if (argc != 3)
	{
		std::cout << "params: TRAmarkers_G16.tif  participantRes_G16.tif\n";
		return -1;
	}

	i3d::Image3d<i3d::GRAY16> resImg(argv[2]);
	i3d::Image3d<i3d::GRAY16> traImg(argv[1]);
	i3d::Image3d<i3d::GRAY16> reducedTraImg;

	char fn[128];
	int erosionRound = 1;
	while ( reduceMarkers(traImg,reducedTraImg,erosionRound) )
	{
		sprintf(fn,"/temp/reduced_%03d.tif",erosionRound);
		std::cout << "managed and saving: " << fn << "\n";
		reducedTraImg.SaveImage(fn);

		//TODO: test against resImg

		++erosionRound;
	}

	std::cout << "returning with max legal erosion of " << erosionRound-1 << "\n";
	return erosionRound-1;
}
