/*
 * pbrt source code Copyright(c) 1998-2004 Matt Pharr and Greg Humphreys
 *
 * All Rights Reserved.
 * For educational use only; commercial use expressly forbidden.
 * NO WARRANTY, express or implied, for this software.
 * (See file License.txt for complete license)
 */

// stratified.cpp*
#include "sampling.h"
#include "paramset.h"
#include "film.h"

// MultijitteredSampler Declarations
class MultijitteredSampler : public Sampler {
public:
	// MultijitteredSampler Public Methods
	MultijitteredSampler(int xstart, int xend,
	                  int ystart, int yend,
					  int xs, int ys, bool jitter);
	int RoundSize(int size) const {
		return size;
	}
	~MultijitteredSampler() {
		FreeAligned(imageSamples);
	}
	bool GetNextSample(Sample *sample);
	void MultijitteredSample2D(float *samp, int nx, int ny, bool jitter);
private:
	// MultijitteredSampler Private Data
	int xPixelSamples, yPixelSamples;
	bool jitterSamples;
	int xPos, yPos;
	float *imageSamples, *lensSamples, *timeSamples;
	int samplePos;
};
// MultijitteredSampler Method Definitions
MultijitteredSampler::MultijitteredSampler(int xstart, int xend,
		int ystart, int yend, int xs, int ys, bool jitter)
	: Sampler(xstart, xend, ystart, yend, xs * ys) {
	jitterSamples = jitter;
	xPos = xPixelStart;
	yPos = yPixelStart;
	xPixelSamples = xs;
	yPixelSamples = ys;
	// Allocate storage for a pixel's worth of stratified samples
	imageSamples = (float *)AllocAligned(5 * xPixelSamples *
		yPixelSamples * sizeof(float));
	lensSamples =
		imageSamples + 2 * xPixelSamples * yPixelSamples;
	timeSamples =
		lensSamples + 2 * xPixelSamples * yPixelSamples;
	// Generate stratified camera samples for (_xPos_,_yPos_)
	MultijitteredSample2D(imageSamples,
		xPixelSamples, yPixelSamples,
		jitterSamples);
	MultijitteredSample2D(lensSamples,
		xPixelSamples, yPixelSamples,
		jitterSamples);
	//1D samples already satisfy the latin hypercube property, so the stratified 1D samples are fine
	StratifiedSample1D(timeSamples, xPixelSamples*yPixelSamples,
		jitterSamples);
	// Shift stratified image samples to pixel coordinates
	for (int o = 0; o < 2 * xPixelSamples * yPixelSamples; o += 2) {
		imageSamples[o]   += xPos;
		imageSamples[o+1] += yPos;
	}
	// Decorrelate sample dimensions
	Shuffle(lensSamples, xPixelSamples*yPixelSamples, 2);
	Shuffle(timeSamples, xPixelSamples*yPixelSamples, 1);
	samplePos = 0;
}

void MultijitteredSampler::MultijitteredSample2D(float *samp, int nx, int ny,
								   bool jitter) {
	float subcell_width = 1.f/(nx*ny);
	//initialize canonical pattern:
	for(int i = 0; i < ny; i++) {
		for(int j = 0; j < nx; j++) {
			float jx = jitter ? RandomFloat() : 0.5f;
			float jy = jitter ? RandomFloat() : 0.5f;
			*samp++ = (i*nx + j + jx)*subcell_width;
			*samp++ = (j*ny + i + jy)*subcell_width;
		}
	}
	//shuffle!
	Shuffle(samp, xPixelSamples*yPixelSamples, 2);
	/*for(int i = 0; i < ny; i++) {
		for(int j = 0; j < nx-1; j++) {
			int k = RandomUInt()%(nx-j) + j;
			swap(samp[i*nx+j],samp[i*nx+k]);
			int l = RandomUInt()%(ny-i) + i;
			swap(samp[i*nx+j],samp[l*nx+j]);
		}
	}*/
}

bool MultijitteredSampler::GetNextSample(Sample *sample) {
	// Compute new set of samples if needed for next pixel
	if (samplePos == xPixelSamples * yPixelSamples) {
		// Advance to next pixel for stratified sampling
		if (++xPos == xPixelEnd) {
			xPos = xPixelStart;
			++yPos;
		}
		if (yPos == yPixelEnd)
			return false;
		// Generate stratified camera samples for (_xPos_,_yPos_)
		MultijitteredSample2D(imageSamples,
			xPixelSamples, yPixelSamples,
			jitterSamples);
		MultijitteredSample2D(lensSamples,
			xPixelSamples, yPixelSamples,
			jitterSamples);
		StratifiedSample1D(timeSamples, xPixelSamples*yPixelSamples,
			jitterSamples); //no need to change 1D sampling.
		// Shift stratified image samples to pixel coordinates
		for (int o = 0;
		     o < 2 * xPixelSamples * yPixelSamples;
			 o += 2) {
			imageSamples[o]   += xPos;
			imageSamples[o+1] += yPos;
		}
		// Decorrelate sample dimensions
		Shuffle(lensSamples, xPixelSamples*yPixelSamples, 2);
		Shuffle(timeSamples, xPixelSamples*yPixelSamples, 1);
		samplePos = 0;
	}
	// Return next _MultijitteredSampler_ sample point
	sample->imageX = imageSamples[2*samplePos];
	sample->imageY = imageSamples[2*samplePos+1];
	sample->lensU = lensSamples[2*samplePos];
	sample->lensV = lensSamples[2*samplePos+1];
	sample->time = timeSamples[samplePos];
	// Generate stratified samples for integrators
	for (u_int i = 0; i < sample->n1D.size(); ++i)
		LatinHypercube(sample->oneD[i], sample->n1D[i], 1);
	for (u_int i = 0; i < sample->n2D.size(); ++i)
		LatinHypercube(sample->twoD[i], sample->n2D[i], 2);
	++samplePos;
	return true;
}

extern "C" DLLEXPORT Sampler *CreateSampler(const ParamSet &params, const Film *film) {
	bool jitter = params.FindOneBool("jitter", true);
	// Initialize common sampler parameters
	int xstart, xend, ystart, yend;
	film->GetSampleExtent(&xstart, &xend, &ystart, &yend);
	int xsamp = params.FindOneInt("xsamples", 2);
	int ysamp = params.FindOneInt("ysamples", 2);
	return new MultijitteredSampler(xstart, xend, ystart, yend, xsamp, ysamp,
		jitter);
}