/***************************************************************************
 *   Copyright (C) 1998-2009 by authors (see AUTHORS.txt )                 *
 *                                                                         *
 *   This file is part of LuxRender.                                       *
 *                                                                         *
 *   Lux Renderer is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 3 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   Lux Renderer is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>. *
 *                                                                         *
 *   This project is based on PBRT ; see http://www.pbrt.org               *
 *   Lux Renderer website : http://www.luxrender.net                       *
 ***************************************************************************/

// initial metropolis light transport sample integrator
// by radiance

// TODO: add scaling of output image samples

// erpt.cpp*
#include "erpt.h"
#include "dynload.h"
#include "scene.h"
#include "error.h"

using namespace lux;

#define SAMPLE_FLOATS 6

// mutate a value in the range [0-1]
static float mutate(const float x, const float randomValue)
{
	static const float s1 = 1.f / 1024.f, s2 = 1.f / 16.f;
	float dx = s2 * powf(s1 / s2, fabsf(2.f * randomValue - 1.f));
	if (randomValue < 0.5f) {
		float x1 = x + dx;
		return (x1 > 1.f) ? x1 - 1.f : x1;
	} else {
		float x1 = x - dx;
		return (x1 < 0.f) ? x1 + 1.f : x1;
	}
}

// mutate a value in the range [min-max]
static float mutateScaled(const float x, const float randomValue, const float mini, const float maxi, const float range)
{
	float dx = range * expf(-logf(2.f * range) * fabsf(2.f * randomValue - 1.f));
	if (randomValue < 0.5f) {
		float x1 = x + dx;
		return (x1 > maxi) ? x1 - maxi + mini : x1;
	} else {
		float x1 = x - dx;
		return (x1 < mini) ? x1 - mini + maxi : x1;
	}
}

// Metropolis method definitions
ERPTSampler::ERPTSampler(u_int totMutations, float rng,
	Sampler *sampler) :
 Sampler(sampler->xPixelStart, sampler->xPixelEnd,
	sampler->yPixelStart, sampler->yPixelEnd, sampler->samplesPerPixel),
 normalSamples(0), totalSamples(0), totalTimes(0), totalMutations(totMutations),
 range(rng), baseSampler(sampler),
 baseImage(NULL), sampleImage(NULL), currentImage(NULL),
 baseTimeImage(NULL), timeImage(NULL), currentTimeImage(NULL), offset(NULL),
 numChains(0), chain(0), mutation(~0U), stamp(0),
 baseLY(0.f), quantum(0.f), weight(0.f), LY(0.f), alpha(0.f),
 totalLY(0.), sampleCount(0.)
{
}

ERPTSampler::~ERPTSampler() {
	FreeAligned(sampleImage);
	FreeAligned(baseImage);
	FreeAligned(timeImage);
	FreeAligned(baseTimeImage);
	delete baseSampler;
}

// Copy
ERPTSampler* ERPTSampler::clone() const
{
	ERPTSampler *newSampler = new ERPTSampler(*this);
	newSampler->baseSampler = baseSampler->clone();
	newSampler->totalSamples = 0;
	newSampler->sampleImage = NULL;
	return newSampler;
}

static void initERPT(ERPTSampler *sampler, const Sample *sample)
{
	u_int i;
	sampler->normalSamples = SAMPLE_FLOATS;
	for (i = 0; i < sample->n1D.size(); ++i)
		sampler->normalSamples += sample->n1D[i];
	for (i = 0; i < sample->n2D.size(); ++i)
		sampler->normalSamples += 2 * sample->n2D[i];
	sampler->totalSamples = sampler->normalSamples;
	sampler->offset = new u_int[sample->nxD.size()];
	sampler->totalTimes = 0;
	for (i = 0; i < sample->nxD.size(); ++i) {
		sampler->offset[i] = sampler->totalSamples;
		sampler->totalTimes += sample->nxD[i];
		sampler->totalSamples += sample->dxD[i] * sample->nxD[i];
	}
	sampler->sampleImage = AllocAligned<float>(sampler->totalSamples);
	sampler->baseImage = AllocAligned<float>(sampler->totalSamples);
	sampler->timeImage = AllocAligned<int>(sampler->totalTimes);
	sampler->baseTimeImage = AllocAligned<int>(sampler->totalTimes);
	sampler->baseSampler->SetTsPack(sampler->tspack);
	sampler->baseSampler->SetFilm(sampler->film);
	sampler->mutation = ~0U;

	// Fetch first contribution buffer from pool
	sampler->contribBuffer = sampler->film->scene->contribPool->Next(NULL);
}

// interface for new ray/samples from scene
bool ERPTSampler::GetNextSample(Sample *sample, u_int *use_pos)
{
	sample->sampler = this;
	if (sampleImage == NULL) {
		initERPT(this, sample);
	}

	if (mutation == ~0U) {
		// Dade - we are at a valid checkpoint where we can stop the
		// rendering. Check if we have enough samples per pixel in the film.
		if (film->enoughSamplePerPixel)
			return false;

		const bool ret = baseSampler->GetNextSample(sample, use_pos);
		sample->sampler = this;
		for (u_int i = 0; i < totalTimes; ++i)
			sample->timexD[0][i] = -1;
		sample->stamp = 0;
		currentImage = baseImage;
		currentTimeImage = baseTimeImage;
		stamp = 0;
		return ret;
	} else {
		if (mutation == 0) {
			// *** new chain ***
			for (u_int i = 0; i < totalTimes; ++i)
				sample->timexD[0][i] = baseTimeImage[i];
			sample->stamp = 0;
			currentImage = baseImage;
			currentTimeImage = baseTimeImage;
			stamp = 0;
		}
		// *** small mutation ***
		// mutate current sample
		sample->imageX = mutateScaled(currentImage[0], tspack->rng->floatValue(), xPixelStart, xPixelEnd, range);
		sample->imageY = mutateScaled(currentImage[1], tspack->rng->floatValue(), yPixelStart, yPixelEnd, range);
		sample->lensU = mutate(currentImage[2], tspack->rng->floatValue());
		sample->lensV = mutate(currentImage[3], tspack->rng->floatValue());
		sample->time = mutate(currentImage[4], tspack->rng->floatValue());
		sample->wavelengths = mutate(currentImage[5], tspack->rng->floatValue());
		for (u_int i = SAMPLE_FLOATS; i < normalSamples; ++i)
				sample->oneD[0][i - SAMPLE_FLOATS] = mutate(currentImage[i], tspack->rng->floatValue());
		++(sample->stamp);
	}

	return true;
}

float *ERPTSampler::GetLazyValues(Sample *sample, u_int num, u_int pos)
{
	const u_int size = sample->dxD[num];
	float *data = sample->xD[num] + pos * size;
	const int stampLimit = sample->stamp;
	if (sample->timexD[num][pos] != stampLimit) {
		if (sample->timexD[num][pos] == -1) {
			baseSampler->GetLazyValues(sample, num, pos);
			sample->timexD[num][pos] = 0;
		} else {
			const u_int start = offset[num] + pos * size;
			float *image = currentImage + start;
			for (u_int i = 0; i < size; ++i)
				data[i] = image[i];
		}
		for (int &time = sample->timexD[num][pos]; time < stampLimit; ++time) {
			for (u_int i = 0; i < size; ++i)
				data[i] = mutate(data[i], tspack->rng->floatValue());
		}
	}
	return data;
}

// interface for adding/accepting a new image sample.
void ERPTSampler::AddSample(const Sample &sample)
{
	vector<Contribution> &newContributions(sample.contributions);
	float newLY = 0.0f;
	for(u_int i = 0; i < newContributions.size(); ++i)
		newLY += newContributions[i].color.Y();
	if (mutation == 0U || mutation == ~0U) {
		if (weight > 0.f) {
			// Add accumulated contribution of previous reference sample
			weight *= quantum / LY;
			if (!isinf(weight) && LY > 0.f) {
				for(u_int i = 0; i < oldContributions.size(); ++i) {
					// Radiance - added new use of contributionpool/buffers
					if(&oldContributions && !contribBuffer->Add(&oldContributions[i], weight)) {
						contribBuffer = film->scene->contribPool->Next(contribBuffer);
						contribBuffer->Add(&oldContributions[i], weight);
					}
				}
			}
			weight = 0.f;
		}
		if (mutation == ~0U) {
			contribBuffer->AddSampleCount(1.f);
			++sampleCount;
			if (!(newLY > 0.f)) {
				newContributions.clear();
				return;
			}
			totalLY += newLY;
			const float meanIntensity = totalLY > 0. ? static_cast<float>(totalLY / sampleCount) : 1.f;
			// calculate the number of chains on a new seed
			quantum = newLY / meanIntensity;
			numChains = max(1U, Floor2UInt(quantum + .5f));
			if (numChains > 100) {
				printf("%d chains -> %d\n", numChains, totalMutations);
				numChains = totalMutations;
			}
			quantum /= (numChains * totalSamples);
			baseLY = newLY;
			baseContributions = newContributions;
			baseImage[0] = sample.imageX;
			baseImage[1] = sample.imageY;
			baseImage[2] = sample.lensU;
			baseImage[3] = sample.lensV;
			baseImage[4] = sample.time;
			baseImage[5] = sample.wavelengths;
			for (u_int i = SAMPLE_FLOATS; i < totalSamples; ++i)
				baseImage[i] = sample.oneD[0][i - SAMPLE_FLOATS];
			for (u_int i = 0 ; i < totalTimes; ++i)
				baseTimeImage[i] = sample.timexD[0][i];
			mutation = 0;
			newContributions.clear();
			return;
		}
		LY = baseLY;
		oldContributions = baseContributions;
	}
	// calculate accept probability from old and new image sample
	float accProb;
	if (LY > 0.f)
		accProb = min(1.f, newLY / LY);
	else
		accProb = 1.f;
	float newWeight = accProb;
	weight += 1.f - accProb;

	// try accepting of the new sample
	if (accProb == 1.f || tspack->rng->floatValue() < accProb) {
		// Add accumulated contribution of previous reference sample
		weight *= quantum / LY;
		if (!isinf(weight) && LY > 0.f) {
			for(u_int i = 0; i < oldContributions.size(); ++i) {
				// Radiance - added new use of contributionpool/buffers
				if(&oldContributions && !contribBuffer->Add(&oldContributions[i], weight)) {
					contribBuffer = film->scene->contribPool->Next(contribBuffer);
					contribBuffer->Add(&oldContributions[i], weight);
				}
			}
		}
		weight = newWeight;
		LY = newLY;
		oldContributions = newContributions;

		// Save new contributions for reference
		sampleImage[0] = sample.imageX;
		sampleImage[1] = sample.imageY;
		sampleImage[2] = sample.lensU;
		sampleImage[3] = sample.lensV;
		sampleImage[4] = sample.time;
		sampleImage[5] = sample.wavelengths;
		for (u_int i = SAMPLE_FLOATS; i < totalSamples; ++i)
			sampleImage[i] = sample.oneD[0][i - SAMPLE_FLOATS];
		for (u_int i = 0 ; i < totalTimes; ++i)
			timeImage[i] = sample.timexD[0][i];
		stamp = sample.stamp;
		currentImage = sampleImage;
		currentTimeImage = timeImage;
	} else {
		// Add contribution of new sample before rejecting it
		newWeight *= quantum / newLY;
		if (!isinf(newWeight) && newLY > 0.f) {
			for(u_int i = 0; i < newContributions.size(); ++i) {
				// Radiance - added new use of contributionpool/buffers
				if(!contribBuffer->Add(&newContributions[i], newWeight)) {
					contribBuffer = film->scene->contribPool->Next(contribBuffer);
					contribBuffer->Add(&newContributions[i], newWeight);
				}
			}
		}

		// Restart from previous reference
		for (u_int i = 0; i < totalTimes; ++i)
			sample.timexD[0][i] = currentTimeImage[i];
		sample.stamp = stamp;
	}
	if (++mutation >= totalMutations) {
		mutation = 0;
		if (++chain >= numChains) {
			chain = 0;
			mutation = ~0U;
		}
	}
	newContributions.clear();
}

Sampler* ERPTSampler::CreateSampler(const ParamSet &params, const Film *film)
{
	int xStart, xEnd, yStart, yEnd;
	film->GetSampleExtent(&xStart, &xEnd, &yStart, &yEnd);
	int totMutations = params.FindOneInt("chainlength", 100);	// number of mutations from a given seed
	float range = params.FindOneFloat("mutationrange", (xEnd - xStart + yEnd - yStart) / 50.f);	// maximum distance in pixel for a small mutation
	string base = params.FindOneString("basesampler", "random");	// sampler for new chain seed
	Sampler *sampler = MakeSampler(base, params, film);
	if (sampler == NULL) {
		luxError(LUX_SYSTEM, LUX_SEVERE, "ERPTSampler: Could not obtain a valid sampler");
		return NULL;
	}
	return new ERPTSampler(max(totMutations, 0), range, sampler);
}

static DynamicLoader::RegisterSampler<ERPTSampler> r("erpt");