#pragma once

#ifndef TCG_IMAGE_OPS_HPP
#define TCG_IMAGE_OPS_HPP

// tcg includes
#include "../image_ops.h"
#include "../pixel.h"
#include "../unique_ptr.h"

// STD includes
#include <assert.h>
#include <memory>

namespace tcg
{
namespace image_ops
{

//**************************************************************************
//    Flat Blur  implementation
//**************************************************************************

template <typename PtrIn, typename PtrOut, typename PixSum>
void _flatFilterLine(PtrIn in, PtrOut out, int lx, int addIn, int addOut,
					 int radius, PixSum *sums)
{
	typedef typename std::iterator_traits<PtrIn> trIn;
	typedef typename trIn::value_type pixIn_type;

	struct locals {
		inline static void fillSums(int lx, int radius, PtrIn lineIn, PixSum *pixsum, int addIn, int addSum)
		{
			using namespace pixel_ops;

			PtrIn newPixIn = lineIn;

			PixSum sum;

			// Build up to radius pixels adding new pixels only
			int x, xNew, xEnd = tmin(radius, lx);

			for (xNew = 0; xNew != xEnd; ++xNew, newPixIn += addIn)
				sum = sum + *newPixIn;

			// Store sums AND add new pixels. Observe that the current pixel value is decreased
			// BEFORE storage. This makes this function "strict" in calculating the sums.
			for (x = 0; xNew != lx; ++x, ++xNew, lineIn += addIn, newPixIn += addIn, pixsum += addSum) {
				sum = (sum + *newPixIn) - *lineIn;
				*pixsum = *pixsum + sum;
			}

			// Finally, without adding new pixels
			for (; x != lx; ++x, lineIn += addIn, pixsum += addSum) {
				sum = sum - *lineIn;
				*pixsum = *pixsum + sum;
			}
		}
	};

	memset(sums, 0, lx * sizeof(PixSum));

	// Add *strict* halves of the convolution sums
	locals::fillSums(lx, radius, in, sums, addIn, 1);								// Right part of the convolution
	locals::fillSums(lx, radius, in + addIn * (lx - 1), sums + lx - 1, -addIn, -1); // Left   ...

	// Normalize sums
	int diameter = 2 * radius + 1;

	for (int x = 0; x != lx; ++x, in += addIn, out += addOut, ++sums) {
		using namespace pixel_ops;

		pixIn_type pix(*in);
		pixel_ops::assign(pix, (*sums + pix) / diameter); // Note that *in is added, due *strict*
														  // halving the sums above
		*out = pix;
	}
}

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

template <typename PtrIn, typename PtrOut, typename PixSum, typename SelectorFunc>
void _flatFilterLine(PtrIn in, PtrOut out, int lx, int addIn, int addOut,
					 SelectorFunc selector, int radius, PixSum *sums, int *counts)
{
	typedef typename std::iterator_traits<PtrIn> trIn;
	typedef typename trIn::value_type pixIn_type;

	struct locals {
		inline static void fillSums(int lx, int radius, PtrIn lineIn, PixSum *pixsum, int *pixcount,
									int addIn, int addSum, SelectorFunc selector)
		{
			using namespace pixel_ops;

			PtrIn newPixIn = lineIn;

			PixSum sum;
			int count = 0;

			// Build up to radius pixels adding new pixels only
			int x, xNew, xEnd = tmin(radius, lx);

			for (xNew = 0; xNew != xEnd; ++xNew, newPixIn += addIn) {
				const pixIn_type &npi = *newPixIn;

				if (selector(npi))
					sum = sum + npi, ++count;
			}

			// Store sums AND add new pixels. Observe that the current pixel value is decreased
			// BEFORE storage. This makes this function "strict" in calculating the sums.
			for (x = 0; xNew != lx; ++x, ++xNew, lineIn += addIn, newPixIn += addIn,
				pixsum += addSum, pixcount += addSum) {
				const pixIn_type &npi = *newPixIn, &li = *lineIn;

				if (selector(npi))
					sum = sum + npi, ++count;
				if (selector(li))
					sum = sum - li, --count;

				*pixsum = *pixsum + sum, *pixcount += count;
			}

			// Finally, without adding new pixels
			for (; x != lx; ++x, lineIn += addIn, pixsum += addSum, pixcount += addSum) {
				const pixIn_type &li = *lineIn;

				if (selector(li))
					sum = sum - li, --count;

				*pixsum = *pixsum + sum, *pixcount += count;
			}
		}
	};

	memset(sums, 0, lx * sizeof(PixSum));
	memset(counts, 0, lx * sizeof(int));

	// Add *strict* halves of the convolution sums
	locals::fillSums(lx, radius, in, sums, counts, addIn, 1, selector);
	locals::fillSums(lx, radius, in + addIn * (lx - 1), sums + lx - 1, counts + lx - 1,
					 -addIn, -1, selector);

	// Normalize sums
	for (int x = 0; x != lx; ++x, in += addIn, out += addOut, ++sums, ++counts) {
		using namespace pixel_ops;

		assert(*counts >= 0);

		pixIn_type pix(*in);
		pixel_ops::assign(pix, (*sums + pix) / (*counts + 1));

		*out = pix;
	}
}

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

template <typename ImgIn, typename ImgOut, typename Scalar>
void blurRows(const ImgIn &imgIn, ImgOut &imgOut, int radius, Scalar)
{
	typedef typename image_traits<ImgIn>::pixel_ptr_type PtrIn;
	typedef typename image_traits<ImgOut>::pixel_ptr_type PtrOut;
	typedef tcg::Pixel<Scalar, image_traits<ImgIn>::pixel_category> PixSum;

	// Validate parameters
	int inLx = image_traits<ImgIn>::width(imgIn), inLy = image_traits<ImgIn>::height(imgIn);
	int outLx = image_traits<ImgOut>::width(imgOut), outLy = image_traits<ImgOut>::height(imgOut);

	assert(inLx == outLx && inLy == outLy);

	int inWrap = image_traits<ImgIn>::wrap(imgIn), outWrap = image_traits<ImgOut>::wrap(imgOut);

	// Allocate an intermediate line of pixels to store sum values
	tcg::unique_ptr<PixSum[]> sums(new PixSum[inLx]);

	// Filter rows
	for (int y = 0; y != inLy; ++y) {
		PtrIn lineIn = image_traits<ImgIn>::pixel(imgIn, 0, y);
		PtrOut lineOut = image_traits<ImgOut>::pixel(imgOut, 0, y);

		_flatFilterLine(lineIn, lineOut, inLx, 1, 1, radius, sums.get());
	}
}

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

template <typename ImgIn, typename ImgOut, typename SelectorFunc, typename Scalar>
void blurRows(const ImgIn &imgIn, ImgOut &imgOut, int radius, SelectorFunc selector, Scalar)
{
	typedef typename image_traits<ImgIn>::pixel_ptr_type PtrIn;
	typedef typename image_traits<ImgOut>::pixel_ptr_type PtrOut;
	typedef tcg::Pixel<Scalar, image_traits<ImgIn>::pixel_category> PixSum;

	// Validate parameters
	int inLx = image_traits<ImgIn>::width(imgIn), inLy = image_traits<ImgIn>::height(imgIn);
	int outLx = image_traits<ImgOut>::width(imgOut), outLy = image_traits<ImgOut>::height(imgOut);

	assert(inLx == outLx && inLy == outLy);

	int inWrap = image_traits<ImgIn>::wrap(imgIn), outWrap = image_traits<ImgOut>::wrap(imgOut);

	// Allocate an intermediate line of pixels to store sum values
	tcg::unique_ptr<PixSum[]> sums(new PixSum[inLx]);
	tcg::unique_ptr<int[]> counts(new int[inLx]);

	// Filter rows
	for (int y = 0; y != inLy; ++y) {
		PtrIn lineIn = image_traits<ImgIn>::pixel(imgIn, 0, y);
		PtrOut lineOut = image_traits<ImgOut>::pixel(imgOut, 0, y);

		_flatFilterLine(lineIn, lineOut, inLx, 1, 1, selector, radius, sums.get(), counts.get());
	}
}

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

template <typename ImgIn, typename ImgOut, typename Scalar>
void blurCols(const ImgIn &imgIn, ImgOut &imgOut, int radius, Scalar)
{
	typedef typename image_traits<ImgIn>::pixel_ptr_type PtrIn;
	typedef typename image_traits<ImgOut>::pixel_ptr_type PtrOut;
	typedef tcg::Pixel<Scalar, typename image_traits<ImgIn>::pixel_category> PixSum;

	// Validate parameters
	int inLx = image_traits<ImgIn>::width(imgIn), inLy = image_traits<ImgIn>::height(imgIn);
	int outLx = image_traits<ImgOut>::width(imgOut), outLy = image_traits<ImgOut>::height(imgOut);

	assert(inLx == outLx && inLy == outLy);

	int inWrap = image_traits<ImgIn>::wrap(imgIn), outWrap = image_traits<ImgOut>::wrap(imgOut);

	// Allocate an intermediate line of pixels to store sum values
	tcg::unique_ptr<PixSum[]> sums(new PixSum[inLy]);

	// Filter columns
	for (int x = 0; x != inLx; ++x) {
		PtrIn lineIn = image_traits<ImgIn>::pixel(imgIn, x, 0);
		PtrOut lineOut = image_traits<ImgOut>::pixel(imgOut, x, 0);

		_flatFilterLine(lineIn, lineOut, inLy, inWrap, outWrap, radius, sums.get());
	}
}

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

template <typename ImgIn, typename ImgOut, typename SelectorFunc, typename Scalar>
void blurCols(const ImgIn &imgIn, ImgOut &imgOut, int radius, SelectorFunc selector, Scalar)
{
	typedef typename image_traits<ImgIn>::pixel_ptr_type PtrIn;
	typedef typename image_traits<ImgOut>::pixel_ptr_type PtrOut;
	typedef tcg::Pixel<Scalar, typename image_traits<ImgIn>::pixel_category> PixSum;

	// Validate parameters
	int inLx = image_traits<ImgIn>::width(imgIn), inLy = image_traits<ImgIn>::height(imgIn);
	int outLx = image_traits<ImgOut>::width(imgOut), outLy = image_traits<ImgOut>::height(imgOut);

	assert(inLx == outLx && inLy == outLy);

	int inWrap = image_traits<ImgIn>::wrap(imgIn), outWrap = image_traits<ImgOut>::wrap(imgOut);

	// Allocate an intermediate line of pixels to store sum values
	tcg::unique_ptr<PixSum[]> sums(new PixSum[inLy]);
	tcg::unique_ptr<int[]> counts(new int[inLy]);

	// Filter columns
	for (int x = 0; x != inLx; ++x) {
		PtrIn lineIn = image_traits<ImgIn>::pixel(imgIn, x, 0);
		PtrOut lineOut = image_traits<ImgOut>::pixel(imgOut, x, 0);

		_flatFilterLine(lineIn, lineOut, inLy, inWrap, outWrap, selector, radius, sums.get(), counts.get());
	}
}

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

template <typename ImgIn, typename ImgOut, typename Scalar>
void blur(const ImgIn &imgIn, ImgOut &imgOut, int radius, Scalar)
{
	blurRows(imgIn, imgOut, radius, Scalar);
	blurCols(imgOut, imgOut, radius, Scalar);
}

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

template <typename ImgIn, typename ImgOut, typename SelectorFunc, typename Scalar>
void blur(const ImgIn &imgIn, ImgOut &imgOut, int radius, SelectorFunc selector, Scalar)
{
	blurRows(imgIn, imgOut, radius, selector, Scalar);
	blurCols(imgOut, imgOut, radius, selector, Scalar);
}
}
} // namespace tcg::image_ops

#endif // TCG_IMAGE_OPS_HPP