#include "UniformSampler.h"

using namespace VRUT;

const float UniformSampler::stepTable[] = {8.0f, 4.0f, 2.0f, 1.0f, 0.5f, 0.25f, 0.125f, 0.0625f, 0.03125f};

void UniformSampler::GenerateSamples(int LOD0)
{
	samples.clear();
	samples.resize(std::max(UNIFORM_PIXEL_LODS - LOD0, 0));
	if (samples.size() == 0)
		return;
	bool sampled[TILE_SIZE * TILE_SIZE];
	memset(sampled, 0, sizeof(bool) * TILE_SIZE * TILE_SIZE);
	for (int k = 0; k < UNIFORM_PIXEL_LODS; ++k)
	{
		int step = stepTable[k];
		int subDiv = 1;
		if (step == 2)
			subDiv = 2;
		if (step == 1)
			subDiv = 4;
		int subTileSize = TILE_SIZE / subDiv;
		int y0 = 0;

		//Iterate over subtiles
		for (int i = 0; i < subDiv; ++i)
		{
			int x0 = 0;

			for (int j = 0; j < subDiv; ++j)
			{
				//Move around subtile
				for (int y = y0; y < y0 + subTileSize; y += step)
				{
					for (int x = x0; x < x0 + subTileSize; x += step)
					{
						int index = y * TILE_SIZE + x;
						if (!sampled[index])
						{
							sampled[index] = true;
							//All samples, up to LOD0, are stored in samples[0]
							//Higher LODs are stored subsequently
							if (k <= LOD0)
								samples[0].push_back(Vector2(x,y));
							else
								samples[k-LOD0].push_back(Vector2(x,y));
						}
					}
				}
				x0 += subTileSize;
			}
			y0 += subTileSize;
		}
	}
}

void UniformSampler::GenerateSubSamples(int LOD0)
{
	subSamples.clear();
	subSamples.resize(std::max(UNIFORM_SUBPIXEL_LODS - std::max(LOD0 - UNIFORM_PIXEL_LODS, 0), 0));
	if (subSamples.size() == 0)
		return;
	bool sampled[TILE_SUBSIZE * TILE_SUBSIZE];
	memset(sampled, 0, sizeof(bool) * TILE_SUBSIZE * TILE_SUBSIZE);
	// when x and y in the loops below are integers (i.e., zeros),
	// the corresponding sample is already present in samples[0];
	// but samples[0] only exists when sampling the integer
	// coordinates separately in GenerateSamples()
	if (LOD0 < UNIFORM_PIXEL_LODS)
		sampled[0] = true;
	LOD0 = std::max(LOD0 - UNIFORM_PIXEL_LODS, 0);
	for (int k = 0; k < UNIFORM_SUBPIXEL_LODS; ++k)
	{
		float step = stepTable[k + UNIFORM_PIXEL_LODS];
		for (float y = 0.0f; y < 1.0f; y += step)
		{
			for (float x = 0.0f; x < 1.0f; x += step)
			{
				int index = int(y * TILE_SUBSIZE) * TILE_SUBSIZE + int(x * TILE_SUBSIZE);
				if (!sampled[index])
				{
					sampled[index] = true;
					//All subsamples, up to LOD0, are stored in subSamples[0]
					//Higher LODs are stored subsequently
					if (k <= LOD0)
						subSamples[0].push_back(Vector2(x,y));
					else
						subSamples[k-LOD0].push_back(Vector2(x,y));
				}
			}
		}
	}
}

UniformSampler::UniformSampler(int LOD0)
{
	GenerateSamples(LOD0);
	GenerateSubSamples(LOD0);
}

UniformSampler::~UniformSampler()
{
}

bool UniformSampler::SamplePixels(Vector2 pixels[], Tile & tile)
{
	tile.samples = 0;
	// pixel sampling
	if (tile.LOD < samples.size())
	{
		int k = tile.LOD;
		while (tile.offset < samples[k].size() && tile.samples < MAX_SAMPLES)
		{
			Vector2 pixel = samples[k][tile.offset++];
			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
	{
		size_t k = tile.LOD - 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 = IntMod2k(i, TILE_SIZE) + tile.min.x;
			int y = (i >> LOG_TILE_SIZE) + tile.min.y;
			if (x < tile.max.x && y < tile.max.y)
			{
				for (size_t j = j0; j < subSamples[k].size(); ++j)
				{
					Vector2 &pixel = pixels[tile.samples++];
					pixel = subSamples[k][j];
					pixel.x += x;
					pixel.y += y;
					++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 UniformSampler::GenPixMapIsNecessary(std::map<uint64_t, unsigned int>& m, const Tile* tile, int numPixels)
{
    m.clear();

    const unsigned int LOD = tile->LOD;
    const unsigned int processingRounds = tile->skip;

    if (tile->antialiasing || numPixels==tile->Size() || (numPixels<=0 && (LOD+processingRounds)>samples.size()))
        return false;

    //called with valid Skip value in tile
    if (numPixels<=0)
    {
        if ((LOD + processingRounds) >= samples.size())
            return false;

        for (unsigned int i = LOD; i < (LOD + processingRounds); ++i)
        {
            for (auto sample : samples[i])
            {
                const int px = sample.x + (unsigned int)(tile->min.x);
                const int py = sample.y + (unsigned int)(tile->min.y);

                if (px < tile->max.x && py < tile->max.y)
                {
                    const int index = m.size();
                    m[Vec2iToU64(px, py)] = index;
                }
            }
        }
    }
    //Called without valid Skip value, but with valid number of pixels
    //This is the case when reconstructing the pixel data and final pixel count is known
    else
    {
        //LOD in this case is the NEXT that should be processed in the (possible) following iteration
        //There is no knowledge as to how many processing rounds did the tile undergo (previous skip value unknown)
        //So going in reverse - starting from last LOD & pixel and going towards LOD 0, decrementing pixel count
        for(int i=LOD-1; i>=0 && numPixels > 0; --i)
        {
            for (int j = samples[i].size()-1; j >= 0 && numPixels > 0; --j)
            {
                const int px = samples[i][j].x + (unsigned int)(tile->min.x);
                const int py = samples[i][j].y + (unsigned int)(tile->min.y);

                if (px < tile->max.x && py < tile->max.y)
                {
                    const int index = --numPixels;
                    m[Vec2iToU64(px, py)] = index;
                }
            }
        }
    }

    return true;
}