#include <memory>

#include "tropcm.h"

// TnzCore includes
#include "traster.h"

// STD includes
#include <limits>

//#define UNIT_TEST                                             // Enables unit testing at program startup

//************************************************************************
//    Rationale
//************************************************************************

/*!
  \file                   tdistancetransform.cpp
  \brief    This file implements an O(rows * cols) 2-dimensional distance
            transform algorithm with customizable action on squared pixel
            distance from the closest pixel.
*/

//************************************************************************
//    Local namespace  stuff
//************************************************************************

namespace
{

/*!
  \brief    Given 2 parabolas with (minimal) height at centers \p a and \p b
            and centers separated by distance \p d, returns the min between
            \p d and the value \p x satisfying <TT>a + x^2 == b + (x - d)^2</TT>.
*/

unsigned int takeoverDist(unsigned int a, unsigned int b, unsigned int d)
{
	// The actual formula is: x = (h^2 + b - a) / 2h. It simplifies as follows
	// using integers only.

	// NOTE: It can be proven that with integer division, x/ab == (x/a)/b.
	return (b < a) ? d : std::max((d + (b - a) / d + 1) / 2, d); // Note the +1 to get the ceil
}

//--------------------------------------------------------------

template <typename Pix, typename IsInsideFunc>
void initializeDT(const TRasterPT<Pix> &ras, const TRasterPT<unsigned int> &dtRas,
				  IsInsideFunc isInside)
{
	assert(ras->getLx() == dtRas->getLx() && ras->getLy() == dtRas->getLy());

	static const unsigned int uiMax =					// Due to the above takeoverDist, for
		(std::numeric_limits<unsigned int>::max)() - 2; // d == 1

	int lx = ras->getLx(), ly = ras->getLy();
	for (int y = 0; y != ly; ++y) {
		Pix *pix = ras->pixels(y), *rowEnd = pix + lx;
		unsigned int *dt = dtRas->pixels(y);

		for (; pix != rowEnd; ++pix, ++dt) {
			assert(*dt == 0u);

			if (!isInside(*pix))
				*dt = uiMax;
		}
	}
}

//--------------------------------------------------------------

template <typename Pix, typename OutFunc>
void expand(int lineLength, int linesCount,
			Pix *buf, int incrPix,
			int incrLine,
			unsigned int *dtBuf, int dtIncrPix,
			int dtIncrLine,
			OutFunc outFunc)
{
	struct locals {
		static void copyLine(unsigned int *dst, unsigned int *src, unsigned int *srcEnd,
							 int srcStride)
		{
			for (; src != srcEnd; src += srcStride, ++dst)
				*dst = *src;
		}

		static void buildRange(unsigned int *dtRef, unsigned int *dtLineEnd,
							   unsigned int *&dtEnd, unsigned int *&dtNewRef)
		{
			unsigned int d = 1, dNew = 0,						   // dNew at 0 to provide a consistent dtNewRef
				dMax = (std::numeric_limits<unsigned int>::max)(); // at the end - should not matter though
			unsigned int *dt = dtRef + 1;

			for (; d <= dMax && dt != dtLineEnd; ++d, ++dt) // Pick larger intervals if possible
			{
				unsigned int newDMax = ::takeoverDist(*dtRef, *dt, d); //
				if (newDMax <= dMax) {
					dNew = d;
					dMax = newDMax;
				}
			}

			dtEnd = dtRef + std::min(d, dMax); // Could end the line before (dMax < d)
			dtNewRef = dtRef + dNew;
		}
	}; // locals

	// Allocate a buffer equivalent to a dt line. It will store the original
	// dt values. Final dt values will be written directly on the dt raster.
	// This is necessary since read and write intervals overlap.
	std::unique_ptr<unsigned[]> dtOriginalLine(new unsigned[lineLength]);

	unsigned int *odtLineStart = dtOriginalLine.get(),
				 *odtLineEnd = odtLineStart + lineLength;

	// Process each line
	for (int l = 0; l != linesCount; ++l) {
		unsigned int *dtLineStart = dtBuf + dtIncrLine * l,	// Using dtBuf to track colors from now on,
			*dtLineEnd = dtLineStart + dtIncrPix * lineLength, // it already embeds colorFunc's output due
				*dt = dtLineStart,							   // to the way it was initialized.
					*odtRef = odtLineStart;

		Pix *lineStart = buf + incrLine * l,
			*pix = lineStart;

		// Make a copy of the original dt values
		locals::copyLine(dtOriginalLine.get(), dtLineStart, dtLineEnd, dtIncrPix);

		// Expand a colored pixel along the line
		while (dt != dtLineEnd) {
			// The line is subdivided in consecutive ranges associated to the same
			// half-parabola - process one

			// Build a half-parabola range
			unsigned int *dtEnd, *odtNewRef;
			locals::buildRange(odtRef, odtLineEnd, dtEnd, odtNewRef);

			assert(odtLineStart <= odtNewRef && odtNewRef <= odtLineEnd);
			assert(odtLineStart <= dtEnd && dtEnd <= odtLineEnd);

			dtEnd = dtLineStart + dtIncrPix * (dtEnd - odtLineStart); // Convert dtEnd to the dt raster buffer

			// Process the range
			Pix *ref = lineStart + incrPix * (odtRef - odtLineStart);

			unsigned int d = (pix - ref) / incrPix;
			for (; dt != dtEnd; ++d, dt += dtIncrPix, pix += incrPix)
				outFunc(*pix, *ref, *dt = *odtRef + sq(d));

			odtRef = odtNewRef;
		}
	}
}

//--------------------------------------------------------------

/*!
  \brief    Performs an O(rows * cols) distance transform on the specified
            raster image.

  \details  The algorithm relies on the separability of the 2D DT into 2
            passes (by rows and columns) of 1-dimensional DTs.

            The 1D DT sums a pre-existing (from the previous DT step if any)
            DT result with the one currently calculated.

  \warning  Templace parameter OutFunc is supposed to satisfy \a transitivity
            upon comparison of its output - so, if \p b is the output of \p a,
            and \p c is the output of \p b, then \p c is the same as the output
            of \p a.

  \todo     Accept a different output raster - but preserve the case where
            (srcRas == dstRas).
*/

template <typename Pix, typename IsInsideFunc, typename OutFunc>
void distanceTransform(const TRasterPT<Pix> &ras, IsInsideFunc isInside, OutFunc outFunc)
{
	int lx = ras->getLx(), ly = ras->getLy();

	// Allocate a suitable temporary raster holding the (squared) distance transform
	// built from the specified color function
	TRasterPT<unsigned int> dtRas(lx, ly); // Summed squared distances will be limited to
										   // 2 billions. This is generally suitable.
	::initializeDT(ras, dtRas, isInside);  // The raster is binarized directly into the
										   // auxiliary dtRas. Pixels in the set to expand
										   // will have value 0, the others a suitable high value.
	expand(lx, ly, ras->pixels(0), 1, ras->getWrap(),
		   dtRas->pixels(0), 1, dtRas->getWrap(), outFunc);
	expand(lx, ly, ras->pixels(0) + lx - 1, -1, ras->getWrap(),
		   dtRas->pixels(0) + lx - 1, -1, dtRas->getWrap(), outFunc);

	expand(ly, lx, ras->pixels(0), ras->getWrap(), 1,
		   dtRas->pixels(0), dtRas->getWrap(), 1, outFunc);
	expand(ly, lx, ras->pixels(ly - 1), -ras->getWrap(), 1,
		   dtRas->pixels(ly - 1), -dtRas->getWrap(), 1, outFunc);
}
}

//************************************************************************
//    Local Functors
//************************************************************************

/*
  Using functors here just to be absolutely sure that calls are not
  callbacks.
*/

namespace
{

struct SomePaint {
	inline bool operator()(const TPixelCM32 &pix) const
	{
		return (pix.getTone() != 0) || (pix.getPaint() != 0);
	}
};

struct CopyPaint {
	inline void operator()(TPixelCM32 &out, const TPixelCM32 &in, unsigned int) const
	{
		out.setPaint(in.getPaint());
	}
};
}

//************************************************************************
//    API functions
//************************************************************************

void TRop::expandPaint(const TRasterCM32P &rasCM)
{
	distanceTransform(rasCM, SomePaint(), CopyPaint());
}

//************************************************************************
//    Unit testing
//************************************************************************

#if defined UNIT_TEST && !defined NDEBUG

namespace
{

void assertEqualBufs(const TRasterT<unsigned int> &a, const TRasterT<unsigned int> &b)
{
	for (int y = 0; y != a.getLy(); ++y) {
		for (int x = 0; x != a.getLx(); ++x)
			assert(a.pixels(y)[x] == b.pixels(y)[x]);
	}
}

struct Selector {
	inline bool operator()(unsigned int val) const
	{
		return val;
	}
};

struct OutputDT {
	inline void operator()(unsigned int &out, const unsigned int &in, unsigned int d2) const
	{
		out = d2;
	}
};

struct DTTest {
	DTTest()
	{
		unsigned int imgBuf[] = {
			0, 0, 1, 0, 0, 0,
			0, 0, 1, 0, 0, 0,
			0, 1, 1, 1, 1, 0,
			0, 1, 0, 0, 0, 0,
			0, 0, 0, 0, 0, 0,
		};

		unsigned int dtBuf[] = {
			4, 1, 0, 1, 4, 5,
			2, 1, 0, 1, 1, 2,
			1, 0, 0, 0, 0, 1,
			1, 0, 1, 1, 1, 2,
			2, 1, 2, 4, 4, 5,
		};

		TRasterPT<unsigned int> imgRas(6, 5, 6, imgBuf, false),
			dtRas(6, 5, 6, dtBuf, false);

		distanceTransform(imgRas, Selector(), OutputDT());
		assertEqualBufs(*imgRas, *dtRas);
	}
} dtTest;

} // namespace

#endif // UNIT_TEST && !NDEBUG