/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2014 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba 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/>.
*/

#include <mitsuba/render/irrcache.h>
#include <mitsuba/core/statistics.h>

MTS_NAMESPACE_BEGIN

HemisphereSampler::HemisphereSampler(uint32_t M, uint32_t N) : m_M(M), m_N(N) {
	m_entries = new SampleEntry[m_M*m_N];
	m_uk = new Vector[m_N];
	m_vk = new Vector[m_N];
	m_vkMinus = new Vector[m_N];
	m_random = new Random();
}

HemisphereSampler::~HemisphereSampler() {
	delete[] m_entries;
	delete[] m_uk;
	delete[] m_vk;
	delete[] m_vkMinus;
}

void HemisphereSampler::generateDirections(const Intersection &its) {
	for (uint32_t j=0; j<m_M; j++) {
		for (uint32_t k=0; k<m_N; k++) {
			SampleEntry &entry = m_entries[j*m_N + k];
			Point2 sample(m_random->nextFloat(), m_random->nextFloat());

			/* Sample uniformly wrt. projected solid angles */
			Float sinTheta2 = (j+sample.x)/m_M,
			      cosTheta = math::safe_sqrt(1 - sinTheta2),
			      sinTheta = std::sqrt(sinTheta2),
			      sinPhi, cosPhi;

			math::sincos(2*M_PI*(k+sample.y)/m_N, &sinPhi, &cosPhi);
			entry.d  = its.shFrame.toWorld(Vector(
				sinTheta*cosPhi, sinTheta*sinPhi, cosTheta));
			entry.cosTheta = cosTheta;
			entry.sinTheta = sinTheta;
			entry.dist = -1;
		}
	}

	/* Precompute planar vectors - see "Practical Global Illumination" by Jaroslav Krivanek
	   and Pascal Gautron for more details on this notation */
	for (uint32_t k=0; k<m_N; k++) {
		Float phi     =  2*M_PI*(k+.5f)/m_N,
			  vk      =  phi - M_PI/2,
			  vkMinus = (2*M_PI*k)/m_N + M_PI/2;

		/* v_k plane vectors (centered) */
		m_vk[k] = its.shFrame.toWorld(Vector(std::cos(vk), std::sin(vk), 0));

		/* v_{k-} (positioned at the start of the cell intervals) */
		m_vkMinus[k] = its.shFrame.toWorld(Vector(std::cos(vkMinus), std::sin(vkMinus), 0));

		/* u_k plane vectors (centered) */
		m_uk[k] = its.shFrame.toWorld(Vector(std::cos(phi), std::sin(phi), 0));
	}
}

void HemisphereSampler::process(const Intersection &its) {
	for (int i=0; i<3; ++i) {
		m_rGrad[i] = Spectrum(0.0f);
		m_tGrad[i] = Spectrum(0.0f);
	}
	m_E = Spectrum(0.0f);
	m_hMean = 0;
	m_hMin = std::numeric_limits<Float>::infinity();
	m_hMinRestricted = std::numeric_limits<Float>::infinity();

	Float invDists = 0;
	for (uint32_t j=0; j<m_M; j++) {
		const Float cosThetaMinus = std::sqrt(1-j/(Float)m_M),
					sinThetaMinus = std::sqrt(j/(Float)m_M),
					cosTheta      = std::sqrt(1-(j+.5f)/m_M),
					sinTheta      = std::sqrt((j+.5f)/m_M),
					cosThetaPlus  = std::sqrt(1-(j+1)/(Float)m_M),
					cosThetaDiff  = cosThetaMinus - cosThetaPlus,
					tanTheta      = sinTheta / cosTheta;
		for (uint32_t k=0; k<m_N; k++) {
			const SampleEntry &entry = m_entries[j*m_N + k];

			/* Rotational gradient - \pi/(MN) * \sum_{k=0}^{N-1}(v_k \sum_{j=0}^{M-1}) \tan\theta_j * L_{jk}) */
			for (int i=0; i<3; ++i)
				m_rGrad[i] += entry.L * (-tanTheta * m_vk[k][i]);

			if (j>1) {
				/* Gradient in the u_k-direction */
				const SampleEntry &other = m_entries[(j-1)*m_N + k];
				const Float minDist = std::min(entry.dist, other.dist);
				if (minDist > 0) {
					const Float factor = (2*M_PI*cosThetaMinus*cosThetaMinus*sinThetaMinus)/(m_N*minDist);
					const Spectrum spec = (entry.L - other.L) * factor;
					for (int i=0; i<3; ++i)
						m_tGrad[i] += spec * m_uk[k][i];
				}
			}

			int kPrev = k-1;
			if (kPrev < 0)
				kPrev = m_N-1;

			/* Gradient in the v_k-direction */
			const SampleEntry &other = m_entries[j*m_N + kPrev];
			const Float minDist = std::min(entry.dist, other.dist);
			if (minDist > 0) {
				const Spectrum spec = (entry.L - other.L) * (cosTheta * cosThetaDiff
						/ (minDist * sinTheta));
				for (int i=0; i<3; ++i)
					m_tGrad[i] += spec * m_vkMinus[k][i];
			}

			if (entry.dist > 0) {
				invDists += 1.0f / entry.dist;
				m_hMin = std::min(m_hMin, entry.dist);
				/* Discard rays close to the tangent plane */
				if (entry.cosTheta > 0.173f) // at least 10deg
					m_hMinRestricted = std::min(m_hMinRestricted, entry.dist);
			}
			m_E += entry.L;
		}
	}
	if (invDists > 0)
		m_hMean = (m_M*m_N) / invDists;
	for (int i=0; i<3; ++i)
		m_rGrad[i] *= M_PI/(m_M*m_N);
	m_E *= M_PI / (m_M*m_N);
}

/* First pass of neighbor clamping */
struct clamp_self_functor {
	clamp_self_functor(const Point &p, Float &R0) : p(p), R0(R0) {
	}

	void operator()(IrradianceCache::Record *sample) {
		Float distance = (p - sample->p).length();
		R0 = std::min(R0, sample->originalR0 + distance);
	}

	Point p;
	Float &R0;
};

/* Second pass of neighbor clamping */
struct clamp_neighbors_functor {
	clamp_neighbors_functor(const Point &p, Float R0) : p(p), R0(R0) {
	}

	void operator()(IrradianceCache::Record *sample) {
		Float distance = (p - sample->p).length();
		Float distanceLimit = R0 + distance;

		if (sample->originalR0 > distanceLimit) {
			/* Update valid range and clamp back into the
			   permitted interval */
			sample->originalR0 = distanceLimit;
			sample->R0 = std::min(sample->R0_max,
				std::max(sample->R0_min, sample->originalR0));
		}
	}

	Point p;
	Float R0;
};

/* Irradiance interpolation functor */
struct irr_interp_functor {
	irr_interp_functor(const Intersection &its, Float kappa, bool gradients) : its(its),
		kappa(kappa), weightSum(0), gradients(gradients), E(0.0f) {
	}

	void operator()(const IrradianceCache::Record *sample) {
		Float weight = sample->getWeight(its.p, its.shFrame.n, kappa);

		if (weight == 0)
			return;

		Spectrum extrapolated = sample->E;
		if (gradients) {
			Vector crossN = cross(sample->n, its.shFrame.n);
			Vector diff = its.p - sample->p;

			for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
				for (int j=0; j<3; ++j) {
					extrapolated[i] +=
						crossN[j] * sample->rGrad[j][i]
						+ diff[j] * sample->tGrad[j][i];
				}
				extrapolated[i] = std::max(extrapolated[i], (Float) 0);
			}
		}
		E += weight * extrapolated;

		weightSum += weight;
	}

	const Intersection &its;
	Float kappa, weightSum;
	bool gradients;
	Spectrum E;
};

IrradianceCache::IrradianceCache(const AABB &aabb)
 : m_octree(aabb) {
	/* Use the longest AABB axis as an estimate of the scene dimensions */
	m_sceneSize = (aabb.max-aabb.min)[aabb.getLargestAxis()];
	m_mutex = new Mutex();

	/* Reasonable default settings */
	setQuality(1.0f);
	useGradients(true);
	clampNeighbor(true);
	clampScreen(true);
}

IrradianceCache::IrradianceCache(Stream *stream, InstanceManager *manager) :
	m_octree(AABB(stream)) {
	m_mutex = new Mutex();
	m_kappa = stream->readFloat();
	m_sceneSize = stream->readFloat();
	m_clampScreen = stream->readBool();
	m_clampNeighbor = stream->readBool();
	m_useGradients = stream->readBool();
	size_t recordCount = stream->readSize();
	m_records.reserve(recordCount);
	for (size_t i=0; i<recordCount; ++i) {
		Record *sample = new Record(stream);
		Float validRadius = sample->R0 / (2*m_kappa);
		m_octree.insert(sample, AABB(
			sample->p-Vector(1,1,1)*validRadius,
			sample->p+Vector(1,1,1)*validRadius
		));
		m_records.push_back(sample);
	}
}

IrradianceCache::~IrradianceCache() {
	for (size_t i=0; i<m_records.size(); ++i)
		delete m_records[i];
}

void IrradianceCache::serialize(Stream *stream, InstanceManager *manager) const {
	m_octree.getAABB().serialize(stream);
	stream->writeFloat(m_kappa);
	stream->writeFloat(m_sceneSize);
	stream->writeBool(m_clampScreen);
	stream->writeBool(m_clampNeighbor);
	stream->writeBool(m_useGradients);
	stream->writeSize(m_records.size());
	for (uint32_t i=0; i<m_records.size(); ++i)
		m_records[i]->serialize(stream);
}

IrradianceCache::Record *IrradianceCache::put(const RayDifferential &ray, const Intersection &its,
		const HemisphereSampler &hs) {
	const Spectrum &E = hs.getIrradiance();
	TranslationalGradient tGrad;
	for (int i=0; i<3; ++i)
		tGrad[i] = hs.getTranslationalGradient()[i];
	Float R0 = hs.getMinimumDistanceRestricted();
	if (!E.isValid()) {
		Log(EWarn, "Invalid irradiance cache sample: %s", E.toString().c_str());
		return NULL;
	}
	Float R0_min = 0, R0_max = std::numeric_limits<Float>::infinity();

	/* Clamping recommended by Tabellion and Lamourlette ("An Approximate Global
	   Illumination System for Computer Generated Films") */
	if (m_clampScreen && ray.hasDifferentials) {
		const Float d = -dot(its.geoFrame.n, Vector(its.p));
		const Float txRecip = dot(its.geoFrame.n, ray.rxDirection),
		            tyRecip = dot(its.geoFrame.n, ray.ryDirection);
		if (txRecip != 0 && tyRecip != 0) {
			// Ray distances traveled
			const Float tx = -(dot(its.geoFrame.n, Vector(ray.rxOrigin)) + d) /
				txRecip;
			const Float ty = -(dot(its.geoFrame.n, Vector(ray.ryOrigin)) + d) /
				tyRecip;
			Point px = ray.rxOrigin + ray.rxDirection * tx,
				  py = ray.ryOrigin + ray.ryDirection * ty;
			Float sqrtArea = std::sqrt(cross(px-its.p, py-its.p).length())*2;

			R0_min = 3.0f*sqrtArea;
			R0_max = 20.0f*sqrtArea;
		}
	}

	if (m_useGradients) {
		/* Limit R0 by the gradient magnitude [Krivanek et al.] */
		for (int i=0; i<SPECTRUM_SAMPLES; ++i) {
			Vector grad(tGrad[0][i], tGrad[1][i], tGrad[2][i]);
			Float length = grad.length();
			if (length > Epsilon)
				R0 = std::min(R0, E[i]/grad.length());
		}

		/* Limit the translational gradient magnitude [Krivanek et al.] */
		for (int i=0; i<3; ++i)
			tGrad[i] = tGrad[i] *
				std::min((Float) 1, hs.getMinimumDistance() / R0_min);
	}

	if (m_clampNeighbor) {
		/* Perform neighbor clamping [Krivanek et al.] to distribute
		   geometric feature information amongst neighboring hss */
		clamp_self_functor clampSelf(its.p, R0);
		m_octree.searchSphere(BSphere(its.p, R0), clampSelf);
		clamp_neighbors_functor clampNeighbors(its.p, R0);
		m_octree.searchSphere(BSphere(its.p, R0), clampNeighbors);
	}

	Record *record = new Record();
	record->p = its.p;
	record->n = its.shFrame.n;
	record->E = E;
	record->R0 = std::min(R0_max, std::max(R0_min, R0));
	record->originalR0 = R0;
	record->R0_min = R0_min;
	record->R0_max = R0_max;
	for (int i=0; i<3; ++i) {
		record->rGrad[i] = hs.getRotationalGradient()[i];
		record->tGrad[i] = tGrad[i];
	}
	insert(record);
	return record;
}

void IrradianceCache::insert(Record *record) {
	Float validRadius = record->R0 / (2*m_kappa);
	m_octree.insert(record, AABB(
		record->p-Vector(1,1,1)*validRadius,
		record->p+Vector(1,1,1)*validRadius
	));
	LockGuard lock(m_mutex);
	m_records.push_back(record);
}

static StatsCounter irradHits("Irradiance cache", "Hits");
static StatsCounter irradMisses("Irradiance cache", "Misses");

bool IrradianceCache::get(const Intersection &its, Spectrum &E) const {
	irr_interp_functor functor(its, m_kappa, m_useGradients);
	m_octree.lookup(its.p, functor);

	if (functor.weightSum > 0) {
		E = functor.E / functor.weightSum;
		++irradHits;
		Assert(!E.isNaN());
		return true;
	}

	++irradMisses;
	return false;
}

std::string IrradianceCache::toString() const {
	std::ostringstream oss;
	oss << "IrradianceCache[" << endl
		<< "  records = " << m_records.size() << "," << endl
		<< "  quality = " << m_kappa << "," << endl
		<< "  sceneSize = " << m_sceneSize << "," << endl
		<< "  clampScreen = " << m_clampScreen << "," << endl
		<< "  clampNeighbor = " << m_clampNeighbor << "," << endl
		<< "  useGradients = " << m_useGradients << endl
		<< "]";
	return oss.str();
}

MTS_IMPLEMENT_CLASS(HemisphereSampler, false, Object)
MTS_IMPLEMENT_CLASS_S(IrradianceCache, false, Object)
MTS_NAMESPACE_END