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


template <class VT>
void replaceMarkers(const i3d::Image3d<VT>& traIn, i3d::Image3d<VT>& traOut, const long radius)
{
	//find set of labels (and their sizes)
	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;
	}

	//allocate output image
	traOut.CopyMetaData(traIn);
	traOut.GetVoxelData() = 0;

	//some shortcuts...
	const int xLine = traOut.GetSizeX();
	const int yLine = traOut.GetSizeY();
	const int zLine = traOut.GetSizeZ();
	VT* outP;

	//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] < 500)
		{
			std::cout << "label " << label << ": too small (" << sizesOfMarkers[label] << ") to be replaced\n";
			//keep the marker as it is
			inP  = traIn.GetFirstVoxelAddr();
			outP = traOut.GetFirstVoxelAddr();
			while (inP != inPL)
			{
				*outP = *inP == label ? label : *outP;
				++inP; ++outP;
			}
		}
		else
		{
			//replace the label with its geom centre
			//current sweeping pos
			long x=0, y=0, z=0;

			//cumulated centre
			long xC=0, yC=0, zC=0;
			long cnt = 0;

			inP  = traIn.GetFirstVoxelAddr();
			while (inP != inPL)
			{
				if (*inP == label)
				{
					xC += x;
					yC += y;
					zC += z;
					++cnt;
				}

				++inP;
				++x;
				if (x == xLine)
				{
					x = 0;
					++y;
					if (y == yLine)
					{
						y = 0;
						++z;
					}
				}
			}

			//finalize the centre coord
			xC /= cnt;
			yC /= cnt;
			zC /= cnt;
			std::cout << "label " << label << ": new marker at [" << xC << "," << yC << "," << zC << "]\n";

			//draw the marker a circle of the given radius, intersected with the original marker,
			//so that the new marker is never outside the original marker
			const long radiusSq = radius*radius;

			inP  = traIn.GetFirstVoxelAddr();
			outP = traOut.GetFirstVoxelAddr();
			for (z = -radius; z <= +radius; ++z)
			for (y = -radius; y <= +radius; ++y)
			for (x = -radius; x <= +radius; ++x)
			{
				//fits into circle?
				if ((x*x + y*y + z*z) > radiusSq) continue;
				//std::cout << "want to draw at " << (x+xC) << "," << (y+yC) << "," << (z+zC) << "\n";

				//fits into the image?
				if ((x+xC) < 0 || (x+xC) >= xLine
				 || (y+yC) < 0 || (y+yC) >= yLine
				 || (z+zC) < 0 || (z+zC) >= zLine) continue;
				//std::cout << "passed\n";

				long offset = ((z+zC)*yLine + y+yC)*xLine + x+xC;
				if (inP[offset] == label) outP[offset] = label;
			}
		}
	}
}


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];
	replaceMarkers(traImg,reducedTraImg,5);

	//find last '/'
	int FSpos = 0;
	int i = 0;
	while (argv[1][i] != 0)
	{
		if (argv[1][i] == '/') FSpos = i;
		++i;
	}

	sprintf(fn,"/temp/");
	int j = 6;
	while (FSpos < i) fn[j++] = argv[1][FSpos++];
	fn[j] = 0;

	std::cout << "managed and saving: " << fn << "\n";
	reducedTraImg.SaveImage(fn);

	return 0;
}
