// tcg includes

#include "tcg/tcg_misc.h"

#include "trop.h"

/*! \file terodilate.cpp

This file contains an implementation of a greyscale (ie per-channel) erode/dilate

morphological operator, following the van Herk/Gil-Werman O(row*cols) algorithm.

An extension with circular structuring element is attempted - unfortunately I could

not retrieve a copy of Miyataka's paper about that, which seemingly claimed

O(rows * cols) too. The implemented algorithm is a sub-optimal O(rows*cols*radius).

*/

//********************************************************

//    Auxiliary  functions

//********************************************************

namespace

{

template <typename pix=""></typename>

void copyMatte(const TRasterPT<pix> &src, const TRasterPT<typename pix::channel=""> &matte)</typename></pix>

{

	typedef typename Pix::Channel Chan;

	int y, lx = src->getLx(), ly = src->getLy();

	for (y = 0; y != ly; ++y) {

		Pix *s, *sBegin = src->pixels(y), *sEnd = sBegin + lx;

		Chan *m, *mBegin = matte->pixels(y);

		for (s = sBegin, m = mBegin; s != sEnd; ++s, ++m)

			*m = s->m;

	}

}

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

template <typename pix=""></typename>

void copyChannels_erode(const TRasterPT<pix> &src, const TRasterPT<typename pix::channel=""> &matte,</typename></pix>

						const TRasterPT<pix> &dst)</pix>

{

	typedef typename Pix::Channel Chan;

	// Just assemble src and matte, remembering to depremultiply src pixels before

	// applying the new matte

	double fac;

	int y, lx = src->getLx(), ly = src->getLy();

	for (y = 0; y != ly; ++y) {

		const Pix *s, *sBegin = src->pixels(y), *sEnd = sBegin + lx;

		Pix *d, *dBegin = dst->pixels(y);

		Chan *m, *mBegin = matte->pixels(y);

		for (s = sBegin, d = dBegin, m = mBegin; s != sEnd; ++s, ++d, ++m) {

			fac = double(*m) / double(s->m);

			d->r = fac * s->r, d->g = fac * s->g, d->b = fac * s->b, d->m = *m;

		}

	}

}

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

template <typename pix=""></typename>

void copyChannels_dilate(const TRasterPT<pix> &src, const TRasterPT<typename pix::channel=""> &matte,</typename></pix>

						 const TRasterPT<pix> &dst)</pix>

{

	typedef typename Pix::Channel Chan;

	// Trickier - since src is presumably premultiplied, increasing its pixels' alpha by direct

	// substitution would expose the excessive RGB discretization of pixels with a low matte value.

	// So, let's just put the pixels on a black background. It should do fine.

	double max = Pix::maxChannelValue;

	int y, lx = src->getLx(), ly = src->getLy();

	for (y = 0; y != ly; ++y) {

		const Pix *s, *sBegin = src->pixels(y), *sEnd = sBegin + lx;

		Pix *d, *dBegin = dst->pixels(y);

		const Chan *m, *mBegin = matte->pixels(y);

		for (s = sBegin, d = dBegin, m = mBegin; s != sEnd; ++s, ++d, ++m) {

			*d = *s;

			d->m = s->m + (1.0 - s->m / max) * *m;

		}

	}

}

} // namespace

//********************************************************

//    EroDilate  algorithms

//********************************************************

namespace

{

template <typename chan=""></typename>

struct MaxFunc {

	inline Chan operator()(const Chan &a, const Chan &b) { return tmax(a, b); }

};

template <typename chan=""></typename>

struct MinFunc {

	inline Chan operator()(const Chan &a, const Chan &b) { return tmin(a, b); }

};

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

// NOTE: src and dst must be NOT OVERLAPPING (eg src != dst)

template <typename chan,="" func="" typename=""></typename>

void erodilate_row(int len, const Chan *src, int sIncr, Chan *dst, int dIncr,

				   int rad, double radR, Func func)

{

	assert(rad >= 0);

	// Segment the row of specified length into wCount windows of max wSize elements

	int w, wSize = 2 * rad + 1, wCount = len / wSize + 1;

	int swIncr = wSize * sIncr, srIncr = rad * sIncr;

	int dwIncr = wSize * dIncr, drIncr = rad * dIncr;

	const Chan *s, *sEnd = src + len * sIncr;

	Chan *d, *dEnd = dst + len * dIncr;

	double one_radR = (1.0 - radR);

	for (w = 0; w != wCount; ++w) {

		Chan *dwBegin = dst + w * dwIncr, *dwEnd = tmin(dwBegin + dwIncr, dEnd);

		// Compute prefixes

		const Chan *swBegin = src + tmax(w * swIncr - srIncr - sIncr, 0),

				   *swEnd = src + tmin(w * swIncr + srIncr + sIncr, len * sIncr);

		s = swEnd - sIncr, d = dst + ((s - src) / sIncr) * dIncr + drIncr; // d already decremented by dIncr

		Chan val = *s, oldVal;

		for (s -= sIncr; (d >= dEnd) && (s >= swBegin); s -= sIncr, d -= dIncr) // s decremented here

		{

			assert(s >= src);

			assert(s < sEnd);

			assert((s - src) % sIncr == 0);

			assert(d >= dst);

			assert((d - dst) % dIncr == 0);

			val = func(oldVal = val, *s);

		}

		for (; s >= swBegin; s -= sIncr, d -= dIncr) {

			assert(s >= src);

			assert(s < sEnd);

			assert((s - src) % sIncr == 0);

			assert(d >= dst);

			assert(d < dEnd);

			assert((d - dst) % dIncr == 0);

			val = func(oldVal = val, *s);

			*d = (oldVal == val) ? val : one_radR * oldVal + radR * val;

		}

		for (d = tmin(d, dEnd - dIncr); d >= dwBegin; d -= dIncr) {

			assert(d >= dst);

			assert(d < dEnd);

			assert((d - dst) % dIncr == 0);

			val = func(oldVal = val, 0);

			*d = (oldVal == val) ? val : one_radR * oldVal + radR * val;

		}

		// Compute suffixes

		swBegin = src + w * swIncr + srIncr, swEnd = tmin(swBegin + swIncr + sIncr, sEnd);

		if (swBegin >= swEnd)

			continue;

		s = swBegin, d = dwBegin;

		val = *s;

		for (s += sIncr; (s < swEnd); s += sIncr, d += dIncr) {

			assert(s >= src);

			assert(s < sEnd);

			assert((s - src) % sIncr == 0);

			assert(d >= dst);

			assert(d < dEnd);

			assert((d - dst) % dIncr == 0);

			val = func(oldVal = val, *s);

			*d = func(*d, (oldVal == val) ? val : one_radR * oldVal + radR * val);

		}

		for (; d < dwEnd; d += dIncr) {

			assert(d >= dst);

			assert(d < dEnd);

			assert((d - dst) % dIncr == 0);

			val = func(oldVal = val, 0);

			*d = func(*d, (oldVal == val) ? val : one_radR * oldVal + radR * val);

		}

	}

}

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

template <typename chan="" pix,="" typename=""></typename>

void erodilate_chan(const TRasterPT<pix> &src, const TRasterPT<chan> &dst, double radius, bool dilate)</chan></pix>

{

	assert(radius > 0.0);

	int radI = tfloor(radius);

	double radR = radius - radI;

	// Using a temporary raster to keep intermediate results. This allows us to

	// perform a cache-friendly iteration in the separable/square kernel case

	int x, y, lx = src->getLx(), ly = src->getLy();

	// Peform rows erodilation

	TRasterPT<chan> temp(ly, lx); // Notice transposition plz</chan>

	{

		if (dilate)

			for (y = 0; y != ly; ++y)

				::erodilate_row(lx, &src->pixels(y)->m, 4, temp->pixels(0) + y, ly, radI, radR, MaxFunc<chan>());</chan>

		else

			for (y = 0; y != ly; ++y)

				::erodilate_row(lx, &src->pixels(y)->m, 4, temp->pixels(0) + y, ly, radI, radR, MinFunc<chan>());</chan>

	}

	// Perform columns erodilation

	{

		if (dilate)

			for (x = 0; x != lx; ++x)

				::erodilate_row(ly, temp->pixels(x), 1, dst->pixels(0) + x, dst->getWrap(), radI, radR, MaxFunc<chan>());</chan>

		else

			for (x = 0; x != lx; ++x)

				::erodilate_row(ly, temp->pixels(x), 1, dst->pixels(0) + x, dst->getWrap(), radI, radR, MinFunc<chan>());</chan>

	}

}

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

template <typename pix=""></typename>

void rect_erodilate(const TRasterPT<pix> &src, const TRasterPT<pix> &dst, double radius)</pix></pix>

{

	typedef typename Pix::Channel Chan;

	if (radius == 0.0) {

		// No-op case

		TRop::copy(dst, src);

		return;

	}

	bool dilate = (radius >= 0.0);

	// Perform columns erodilation

	TRasterPT<chan> temp(src->getLx(), src->getLy());</chan>

	::erodilate_chan(src, temp, fabs(radius), dilate);

	// Remember that we have just calculated the matte values. We still have to apply them to the old RGB

	// values, which requires depremultiplying from source matte and premultiplying with the new one.

	if (dilate)

		::copyChannels_dilate(src, temp, dst);

	else

		::copyChannels_erode(src, temp, dst);

}

} // namespace

//********************************************************

//    EroDilate  round algorithm

//********************************************************

namespace

{

template <typename chan,="" func="" typename=""></typename>

void erodilate_quarters(int lx, int ly,

						Chan *src, int sIncrX, int sIncrY,

						Chan *dst, int dIncrX, int dIncrY,

						double radius, double shift, Func func)

{

	double sqRadius = sq(radius);

	double squareHeight = radius / tcg::consts::sqrt2;

	int squareHeightI = tfloor(squareHeight);

	// For every arc point

	int arcY;

	for (arcY = -squareHeightI; arcY <= squareHeightI; ++arcY) {

		// Calculate x and weights

		double sqArcY = sq(arcY);

		assert(sqRadius >= sqArcY);

		double x = shift + sqrt(sqRadius - sqArcY) - squareHeight;

		int arcX = tfloor(x);

		double w = x - arcX, one_w = 1.0 - w;

		// Build dst area influenced by the arc point. Func with 0 outside that.

		TRect bounds(0, 0, lx, ly);

		TRect dRect(bounds * (bounds + TPoint(-arcX, -arcY)));

		TRect sRect(bounds * (bounds + TPoint(arcX, arcY)));

		int sy, dy;

		// Func with 0 before dRect.y0

		for (dy = 0; dy < dRect.y0; ++dy) {

			Chan *d, *dBegin = dst + dy * dIncrY, *dEnd = dBegin + lx * dIncrX;

			for (d = dBegin; d != dEnd; d += dIncrX) {

				//assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);

				*d = func(*d, 0);

			}

		}

		// Func with 0 after dRect.y1

		for (dy = dRect.y1; dy < ly; ++dy) {

			Chan *d, *dBegin = dst + dy * dIncrY, *dEnd = dBegin + lx * dIncrX;

			for (d = dBegin; d != dEnd; d += dIncrX) {

				//assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);

				*d = func(*d, 0);

			}

		}

		// For every dst pixel in the area, Func with the corresponding pixel in src

		for (dy = dRect.y0, sy = sRect.y0; dy != dRect.y1; ++dy, ++sy) {

			Chan *d, *dLine = dst + dy * dIncrY, *dBegin = dLine + dRect.x0 * dIncrX;

			Chan *s, *sLine = src + sy * sIncrY, *sBegin = sLine + sRect.x0 * sIncrX, *sEnd = sLine + sRect.x1 * sIncrX;

			Chan *sLast = sEnd - sIncrX; // sLast would lerp with sEnd

			for (d = dBegin, s = sBegin; s != sLast; d += dIncrX, s += sIncrX) // hence we stop before it

			{

				//assert(s >= src); assert(s < sEnd); assert((s-src) % sIncrX == 0);

				//assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);

				*d = func(*d, *s * one_w + *(s + sIncrX) * w);

			}

			//assert(s >= src); assert(s < sEnd); assert((s-src) % sIncrX == 0);

			//assert(d >= dst); assert(d < dEnd); assert((d-dst) % dIncrX == 0);

			*d = func(*d, *s * one_w); // lerp sLast with 0

		}

	}

}

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

template <typename pix=""></typename>

void circular_erodilate(const TRasterPT<pix> &src, const TRasterPT<pix> &dst, double radius)</pix></pix>

{

	typedef typename Pix::Channel Chan;

	if (radius == 0.0) {

		// No-op case

		TRop::copy(dst, src);

		return;

	}

	// Ok, the idea is: consider the maximal embedded square in our circular structuring element.

	// Erodilating by it consists in the consecutive erodilation by rows and columns with the same

	// 'square' radius. Now, it's easy to see that the square could be 'bent' so that one of its

	// edges matches that of a 1/4 of the circle's edge, while remaining inside the circle.

	// Erodilating by the bent square can be achieved by erodilating first by rows or column for

	// the square edge radius, followed by perpendicular erodilationg with a fourth of our

	// circumference. Sum the 4 erodilations needed to complete the circumference - and it's done.

	// NOTE: Unfortunately, the above decomposition has lots of intersections among the pieces - yet

	// it's simple enough and removes an O(radius) from the naive algorithm. Could be done better?

	// First, build the various erodilation data

	bool dilate = (radius >= 0.0);

	radius = fabs(radius);

	double inner_square_diameter = radius * tcg::consts::sqrt2;

	double shift = 0.25 * inner_square_diameter; // Shift of the bent square SE needed to avoid

												 // touching the circumference on the other side

	double row_filter_radius = 0.5 * (inner_square_diameter - shift);

	double cseShift = 0.5 * shift; // circumference structuring element shift

	int lx = src->getLx(), ly = src->getLy();

	TRasterPT<chan> temp1(lx, ly), temp2(lx, ly);</chan>

	int radI = tfloor(row_filter_radius);

	double radR = row_filter_radius - radI;

	if (dilate) {

		temp2->fill(0); // Initialize with a Func-neutral value

		if (row_filter_radius > 0.0)

			for (int y = 0; y != ly; ++y)

				::erodilate_row(lx, &src->pixels(y)->m, 4, temp1->pixels(y), 1, radI, radR, MaxFunc<chan>());</chan>

		else

			::copyMatte(src, temp1);

		::erodilate_quarters(lx, ly, temp1->pixels(0), 1, lx, temp2->pixels(0), 1, lx, radius, cseShift, MaxFunc<chan>());</chan>

		::erodilate_quarters(lx, ly, temp1->pixels(0) + lx - 1, -1, lx, temp2->pixels(0) + lx - 1, -1, lx, radius, cseShift, MaxFunc<chan>());</chan>

		if (row_filter_radius > 0.0)

			for (int x = 0; x != lx; ++x)

				::erodilate_row(ly, &src->pixels(0)[x].m, 4 * src->getWrap(), temp1->pixels(0) + x, lx, radI, radR, MaxFunc<chan>());</chan>

		else

			::copyMatte(src, temp1);

		::erodilate_quarters(ly, lx, temp1->pixels(0), lx, 1, temp2->pixels(0), lx, 1, radius, cseShift, MaxFunc<chan>());</chan>

		::erodilate_quarters(ly, lx, temp1->pixels(0) + lx * ly - 1, -lx, -1, temp2->pixels(0) + lx * ly - 1, -lx, -1, radius, cseShift, MaxFunc<chan>());</chan>

	} else {

		temp2->fill((std::numeric_limits<chan>::max)()); // Initialize with a Func-neutral value</chan>

		if (row_filter_radius > 0.0)

			for (int y = 0; y != ly; ++y)

				::erodilate_row(lx, &src->pixels(y)->m, 4, temp1->pixels(y), 1, radI, radR, MinFunc<chan>());</chan>

		else

			::copyMatte(src, temp1);

		::erodilate_quarters(lx, ly, temp1->pixels(0), 1, lx, temp2->pixels(0), 1, lx, radius, cseShift, MinFunc<chan>());</chan>

		::erodilate_quarters(lx, ly, temp1->pixels(0) + lx - 1, -1, lx, temp2->pixels(0) + lx - 1, -1, lx, radius, cseShift, MinFunc<chan>());</chan>

		if (row_filter_radius > 0.0)

			for (int x = 0; x != lx; ++x)

				::erodilate_row(ly, &src->pixels(0)[x].m, 4 * src->getWrap(), temp1->pixels(0) + x, lx, radI, radR, MinFunc<chan>());</chan>

		else

			::copyMatte(src, temp1);

		::erodilate_quarters(ly, lx, temp1->pixels(0), lx, 1, temp2->pixels(0), lx, 1, radius, cseShift, MinFunc<chan>());</chan>

		::erodilate_quarters(ly, lx, temp1->pixels(0) + lx * ly - 1, -lx, -1, temp2->pixels(0) + lx * ly - 1, -lx, -1, radius, cseShift, MinFunc<chan>());</chan>

	}

	// Remember that we have just calculated the matte values. We still have to apply them to the old RGB

	// values, which requires depremultiplying from source matte and premultiplying with the new one.

	if (dilate)

		::copyChannels_dilate(src, temp2, dst);

	else

		::copyChannels_erode(src, temp2, dst);

}

} // namespace

//********************************************************

//    EroDilate  main functions

//********************************************************

void TRop::erodilate(const TRasterP &src, const TRasterP &dst,

					 double radius, ErodilateMaskType type)

{

	assert(src->getSize() == dst->getSize());

	src->lock(), dst->lock();

	if ((TRaster32P)src && (TRaster32P)dst)

		switch (type) {

		case ED_rectangular:

			::rect_erodilate<tpixel32>(src, dst, radius);</tpixel32>

			CASE ED_circular : ::circular_erodilate<tpixel32>(src, dst, radius);</tpixel32>

		DEFAULT:

			assert(!"Unknown mask type");

		}

	else if ((TRaster64P)src && (TRaster64P)dst)

		switch (type) {

		case ED_rectangular:

			::rect_erodilate<tpixel64>(src, dst, radius);</tpixel64>

			CASE ED_circular : ::circular_erodilate<tpixel64>(src, dst, radius);</tpixel64>

		DEFAULT:

			assert(!"Unknown mask type");

		}

	else

		assert(!"Unsupported raster type!");

	src->unlock(), dst->unlock();

}