#include "HaltonSampler.h"

using namespace VRUT;

HaltonSampler::HaltonSampler()
{
}

void HaltonSampler::GenerateSamples(const std::vector<unsigned int> & samplesCounts, bool subSampling, bool sort)
{
	// Halton sampling
	Vector2 haltonSamples[TILE_SIZE * TILE_SIZE];
	bool sampled[TILE_SIZE * TILE_SIZE];
	memset(sampled, 0, sizeof(bool) * TILE_SIZE * TILE_SIZE);
	int i = 0;
	int j = 0;
	Halton halton;
	while (i < TILE_SIZE * TILE_SIZE)
	{
		Vector2 sample = TILE_SIZE * halton.Sample2D(j++, 2, 3);
		int x = (int)sample.x;
		int y = (int)sample.y;
		int index = y * TILE_SIZE + x;
		if (!sampled[index])
		{
			sampled[index] = true;
			haltonSamples[i++] = sample;
		}
	}
	// sort along Hilbert curve
	Hilbert2D hilbert;
	HilbertComparator comparator;
	std::vector<HilbertPair> pairs;
	if (subSampling)
	{
		subSamples.clear();
		subSamples.resize(samplesCounts.size());
	}
	else 
	{
		samples.clear();
		samples.resize(samplesCounts.size());
	}
	int offset = 0;
	float scale = 1.0f / TILE_SIZE;
	for (size_t i = 0; i < samplesCounts.size(); ++i)
	{
		for (int j = offset; j < __min(offset + int(samplesCounts[i]), TILE_SIZE * TILE_SIZE); ++j)
		{
			int code = hilbert.xy2d(TILE_SIZE, int(haltonSamples[j].x), int(haltonSamples[j].y));
			pairs.push_back(HilbertPair(haltonSamples[j], code));
		}
		if (sort)
			std::sort(pairs.begin(), pairs.end(), comparator);
		for (auto &pair : pairs)
		{
			if (subSampling)
				subSamples[i].push_back(scale * pair.first);
			else
				samples[i].push_back(pair.first);
		}
		offset += pairs.size();
		pairs.clear();
	}
}

HaltonSampler::HaltonSampler(const std::vector<unsigned int> & samplesCounts, const std::vector<unsigned int> & subSamplesCounts, bool sort)
{
	GenerateSamples(samplesCounts, false, sort);
	GenerateSamples(subSamplesCounts, true, sort);
}

InfiniteHaltonSampler::InfiniteHaltonSampler(const std::vector<unsigned int> & samplesCounts)
{
	GenerateSamples(samplesCounts, false, false);
	//k = -1;
	//lastk = -1;
}

HaltonSampler::~HaltonSampler()
{
}

bool HaltonSampler::SamplePixels(Vector2 pixels[], Tile & tile)
{
	tile.samples = 0;
	// pixel sampling
	int k = tile.LOD;
	if (k < samples.size())
	{
		while (tile.offset < samples[k].size() && tile.samples < MAX_SAMPLES)
		{
			Vector2 pixel = samples[k][tile.offset++];
			pixel.x = floorf(pixel.x) + 0.5f;
			pixel.y = floorf(pixel.y) + 0.5f;
			pixel.x += tile.min.x;
			pixel.y += tile.min.y;
			if (pixel.x < tile.max.x && pixel.y < tile.max.y)
				pixels[tile.samples++] = pixel;
		}
		tile.samplesOverall += tile.samples;
		return tile.samples > 0;
	}
	// antialiasing
	else
	{
		Hilbert2D hilbert;
		k -= samples.size();
		if (k >= subSamples.size())
			return false;
		tile.antialiasing = true;
		int i0 = tile.offset / subSamples[k].size();
		int j0 = tile.offset % subSamples[k].size();
		for (int i = i0; i < TILE_SIZE * TILE_SIZE; ++i)
		{
			int x, y;
			hilbert.d2xy(TILE_SIZE, i, x, y);
			x += tile.min.x;
			y += tile.min.y;
			if (x < tile.max.x && y < tile.max.y)
			{
				for (size_t j = j0; j < subSamples[k].size(); ++j)
				{
					Vector2 pixel = subSamples[k][j];
					pixel.x += x;
					pixel.y += y;
					pixels[tile.samples++] = pixel;
					tile.offset++;
					if (tile.samples >= MAX_SAMPLES)
					{
						tile.samplesOverall += tile.samples;
						return true;
					}
				}
			}
			else
			{
				tile.offset += subSamples[k].size();
			}
		}
		tile.samplesOverall += tile.samples;
		return tile.samples > 0;
	}
}

//static int k = -1;
//static int lastk = -1;
bool InfiniteHaltonSampler::SamplePixels(Vector2 pixels[], Tile & tile)
{
	uint seed = chrono::system_clock::now().time_since_epoch().count();
	default_random_engine gen(seed);
	uniform_real_distribution<double> dist(0.0, 1.0);

	tile.samples = 0;
	// pixel sampling
	//int k = tile.LOD;
	/*int k0 = tile.LOD;
	if (k0 != lastk) {
		k++;
		k %= samples.size();
		lastk = k0;
	}*/
	int k = tile.LOD % samples.size();
	if (k < samples.size())
	{
		while (tile.offset < samples[k].size() && tile.samples < MAX_SAMPLES)
		{
			Vector2 pixel = samples[k][tile.offset++];
			pixel.x = floorf(pixel.x) + 0.5f;
			pixel.y = floorf(pixel.y) + 0.5f;
			pixel.x += tile.min.x;
			pixel.y += tile.min.y;
			//pixel.x += (dist(gen)*2 - 1)*1;
			//pixel.y += (dist(gen)*2 - 1)*1;
			if (pixel.x < tile.max.x && pixel.y < tile.max.y)
				pixels[tile.samples++] = pixel;
		}
		tile.samplesOverall += tile.samples;
		return tile.samples > 0;
	}
	// antialiasing
	else
	{
		Hilbert2D hilbert;
		k -= samples.size();
		if (k >= subSamples.size())
			return false;
		tile.antialiasing = true;
		int i0 = tile.offset / subSamples[k].size();
		int j0 = tile.offset % subSamples[k].size();
		for (int i = i0; i < TILE_SIZE * TILE_SIZE; ++i)
		{
			int x, y;
			hilbert.d2xy(TILE_SIZE, i, x, y);
			x += tile.min.x;
			y += tile.min.y;
			if (x < tile.max.x && y < tile.max.y)
			{
				for (size_t j = j0; j < subSamples[k].size(); ++j)
				{
					Vector2 pixel = subSamples[k][j];
					pixel.x += x;
					pixel.y += y;
					pixels[tile.samples++] = pixel;
					tile.offset++;
					if (tile.samples >= MAX_SAMPLES)
					{
						tile.samplesOverall += tile.samples;
						return true;
					}
				}
			}
			else
			{
				tile.offset += subSamples[k].size();
			}
		}
		tile.samplesOverall += tile.samples;
		return tile.samples > 0;
	}
}

bool HaltonSampler::GenPixMapIsNecessary(std::map<uint64_t, unsigned int>& m, const Tile* tile, int numPixels)
{
    return false;
}
