#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): LOD0(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
	{
#if 1           // original solution using the precomputed subSamples arrays
		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;
					}
				}
				j0 = 0;
			}
			else
			{
				tile.offset += subSamples[k].size();
			}
		}
		tile.samplesOverall += tile.samples;
		return tile.samples > 0;

#else           // subsamples are drawn on the fly, no need for the precomputed subSamples arrays
                tile.antialiasing = true;

                float step = stepTable[LOD0 + tile.LOD];
                int invStep = 1.0f / step;
                int samplesInLOD = (tile.LOD == 0 ? invStep*invStep : 3 * invStep*invStep / 4);  // first LOD: complete grid; subsequent LODs: only 3/4

		int i0 = tile.offset / samplesInLOD;
                int j0 = tile.offset % samplesInLOD;

                int dx0;
                int dy0;

                if (tile.LOD == 0) {  // first LOD: complete grid
                    dx0 = j0 % invStep;
                    dy0 = j0 / invStep;
                }
                else {  // subsequent LODs: only 3/4, pattern of samples repeats every two rows
                    int x2 = j0 % (3 * invStep / 2);  // index within the row pair
                    int y2 = j0 / (3 * invStep / 2);  // index of the row pair

                    if (x2 < invStep / 2) {  // even row: only odd samples
                        dx0 = 2 * x2 + 1;
                        dy0 = 2 * y2;
                    }
                    else {  // odd row: all samples
                        dx0 = x2 - invStep / 2;
                        dy0 = 2 * y2 + 1;
                    }
                }

                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)
                        {
                            float ystep = dy0 * step;
                            for (int dy = dy0; dy < invStep; dy++, ystep+=step) {
                                float xstep = dx0 * step;
                                for (int dx = dx0; dx < invStep; dx++, xstep+=step) {
                                    if (tile.LOD == 0 || (dy & 1 == 1) || (dx & 1 == 1)) {
                                        Vector2 &pixel = pixels[tile.samples++];
                                        pixel.x = x + xstep;
                                        pixel.y = y + ystep;
                                        ++tile.offset;
                                        if (tile.samples >= MAX_SAMPLES)
                                        {
                                            tile.samplesOverall += tile.samples;
                                                return true;
                                        }
                                    }
                                }
                                dx0 = 0;
                            }
                            dy0 = 0;
                        }
                        else
                        {
                            tile.offset += samplesInLOD;
                        }
                }
                tile.samplesOverall += tile.samples;
                return tile.samples > 0;
#endif
	}
}

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;
}




InfiniteUniformSampler::InfiniteUniformSampler(int step, bool halfNow) {
	half = halfNow;
	GenerateSamples(step);
}

InfiniteUniformSampler::~InfiniteUniformSampler()
{
}


void InfiniteUniformSampler::GenerateSamples(int stepNow)
{
	step = stepNow;

	samples.clear();
	samples.resize(1);
	for(int x = 0; x < TILE_SIZE; x+=step)
		for (int y = 0; y < TILE_SIZE; y += step) {
			samples[0].push_back(Vector2(x, y));
		}
}



bool InfiniteUniformSampler::SamplePixels(Vector2 pixels[], Tile & tile)
{
	tile.samples = 0;
	// pixel sampling
	if (/*tile.LOD < samples.size()*/1)
	{
		int q = tile.LOD;
		int even = q % 2;
		if (half)
			q /= 2;

		int n = log(3 * q + 1) / log(4); //last summed exponent in q
		int sum = (1 - (1 << (n * 2))) / (-3); //last sum
		int rem = q - sum;

		int bigNum = 1 << (n * 2);
		rem = (78791 * rem) % bigNum; //randomize pass

		int smallNum = sqrt(bigNum);
		int x = rem % smallNum;
		int y = rem / smallNum % smallNum;

		float baseVal = 1.0f / (1 << (n + 1));
		float offVal = 1.0f / (1 << (n));
		float xF = baseVal + offVal*x;
		float yF = baseVal + offVal*y;
		xF *= step;
		yF *= step;

		int k = tile.LOD % samples.size();
		if (half && !even) //skip first offset when odd
			tile.offset++;
		while (tile.offset < samples[k].size() && tile.samples < MAX_SAMPLES)
		{
			Vector2 pixel = samples[k][tile.offset++];
			pixel.x += tile.min.x + xF;
			pixel.y += tile.min.y + yF;
			if (pixel.x < tile.max.x && pixel.y < tile.max.y)
				pixels[tile.samples++] = pixel;
			if (half) //add one more offset when using half the samples
				tile.offset++;
		}
		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 InfiniteUniformSampler::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;
}