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

    Copyright (c) 2007-2011 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/>.
*/

#if !defined(__SAH_KDTREE3_H)
#define __SAH_KDTREE3_H

#include <mitsuba/render/gkdtree.h>

MTS_NAMESPACE_BEGIN

/// Use a simple hashed 8-entry mailbox per thread
#define MTS_KD_MAILBOX_ENABLED 1
#define MTS_KD_MAILBOX_SIZE 8
#define MTS_KD_MAILBOX_MASK (MTS_KD_MAILBOX_SIZE-1)

/**
 * \brief Implements the surface area heuristic for use
 * by the \ref GenericKDTree construction algorithm.
 */
class SurfaceAreaHeuristic {
public:
	/**
	 * \brief Initialize the surface area heuristic with the bounds of 
	 * a parent node
	 *
	 * Precomputes some information so that traversal probabilities
	 * of potential split planes can be evaluated efficiently
	 */
	inline SurfaceAreaHeuristic(const AABB &aabb) {
		const Vector extents(aabb.getExtents());
		const Float temp = 1.0f / (extents.x * extents.y 
				+ extents.y*extents.z + extents.x*extents.z);
		m_temp0 = Vector(
			extents[1] * extents[2],
			extents[0] * extents[2],
			extents[0] * extents[1]) * temp;
		m_temp1 = Vector(
			extents[1] + extents[2],
			extents[0] + extents[2],
			extents[0] + extents[1]) * temp;
	}

	/**
	 * Given a split on axis \a axis that produces children having extents
	 * \a leftWidth and \a rightWidth along \a axis, compute the probability
	 * of traversing the left and right child during a typical query
	 * operation. 
	 */
	inline std::pair<Float, Float> operator()(int axis, Float leftWidth, Float rightWidth) const {
		return std::pair<Float, Float>(
			m_temp0[axis] + m_temp1[axis] * leftWidth,
			m_temp0[axis] + m_temp1[axis] * rightWidth);
	}

	/**
	 * Compute the underlying quantity used by the tree construction
	 * heuristic. This is used to compute the final cost of a kd-tree.
	 */
	inline static Float getQuantity(const AABB &aabb) {
		return aabb.getSurfaceArea();
	}
private:
	Vector m_temp0, m_temp1;
};

/**
 * This class specializes \ref GenericKDTree to a three-dimensional
 * tree to be used for ray tracing. One additional function call
 * must be implemented by subclasses:
 *
 * /// Check whether a primitive is intersected by the given ray. 
 * /// Some temporary space is supplied, which can be used to cache  
 * /// information about the intersection
 * bool intersect(const Ray &ray, index_type idx, 
 *     Float mint, Float maxt, Float &t, void *tmp);
 *
 * This class implements an epsilon-free version of the optimized ray 
 * traversal algorithm (TA^B_{rec}), which is explained in Vlastimil 
 * Havran's PhD thesis "Heuristic Ray Shooting Algorithms". 
 *
 * \author Wenzel Jakob
 */
template <typename Derived> 
	class SAHKDTree3D : public GenericKDTree<AABB, SurfaceAreaHeuristic, Derived> {
public:
	typedef GenericKDTree<AABB, SurfaceAreaHeuristic, Derived> Parent;
	typedef typename KDTreeBase<AABB>::size_type               size_type;
	typedef typename KDTreeBase<AABB>::index_type              index_type;
	typedef typename KDTreeBase<AABB>::KDNode                  KDNode;

	using Parent::m_nodes;
	using Parent::m_aabb;
	using Parent::m_indices;

protected:
	void buildInternal() {
		size_type primCount = cast()->getPrimitiveCount();
		KDLog(EInfo, "Constructing a SAH kd-tree (%i primitives) ..", primCount);
		GenericKDTree<AABB, SurfaceAreaHeuristic, Derived>::buildInternal();
	}

	/// Cast to the derived class
	inline Derived *cast() {
		return static_cast<Derived *>(this);
	}

	/// Cast to the derived class (const version)
	inline const Derived *cast() const {
		return static_cast<const Derived *>(this);
	}

	/**
	 * \brief Hashed mailbox implementation
	 */
	struct HashedMailbox {
		inline HashedMailbox() {
			memset(entries, 0xFF, sizeof(index_type)*MTS_KD_MAILBOX_SIZE);
		}

		inline void put(index_type primIndex) {
			entries[primIndex & MTS_KD_MAILBOX_MASK] = primIndex;
		}

		inline bool contains(index_type primIndex) const {
			return entries[primIndex & MTS_KD_MAILBOX_MASK] == primIndex;
		}

		index_type entries[MTS_KD_MAILBOX_SIZE];
	};

	/// Ray traversal stack entry for Wald-style incoherent ray tracing
	struct KDStackEntry {
		const KDNode * __restrict node;
		Float mint, maxt;
	};

	/// Ray traversal stack entry for Havran-style incoherent ray tracing
	struct KDStackEntryHavran {
		/* Pointer to the far child */
		const KDNode * __restrict node;
		/* Distance traveled along the ray (entry or exit) */
		Float t;
		/* Previous stack item */
		uint32_t prev;
		/* Associated point */
		Point p;
	};

	/**
	 * \brief Ray tracing kd-tree traversal loop (Havran variant)
	 *
	 * This is generally the most robust and fastest traversal routine
	 * of the methods implemented in this class.
	 */
	template<bool shadowRay> FINLINE
			bool rayIntersectHavran(const Ray &ray, Float mint, Float maxt, 
			Float &t, void *temp) const {
		KDStackEntryHavran stack[MTS_KD_MAXDEPTH];
		#if 0
		static const int prevAxisTable[] = { 2, 0, 1 };
		static const int nextAxisTable[] = { 1, 2, 0 };
		#endif
	
		#if defined(MTS_KD_MAILBOX_ENABLED)
		HashedMailbox mailbox;
		#endif
	
		/* Set up the entry point */
		uint32_t enPt = 0;
		stack[enPt].t = mint;
		stack[enPt].p = ray(mint);
	
		/* Set up the exit point */
		uint32_t exPt = 1;
		stack[exPt].t = maxt;
		stack[exPt].p = ray(maxt);
		stack[exPt].node = NULL;
	
		bool foundIntersection = false;
		const KDNode * __restrict currNode = m_nodes;
		while (currNode != NULL) {
			while (EXPECT_TAKEN(!currNode->isLeaf())) {
				const Float splitVal = (Float) currNode->getSplit();
				const int axis = currNode->getAxis();
				const KDNode * __restrict farChild;
	
				if (stack[enPt].p[axis] <= splitVal) {
					if (stack[exPt].p[axis] <= splitVal) {
						/* Cases N1, N2, N3, P5, Z2 and Z3 (see thesis) */
						currNode = currNode->getLeft();
						continue;
					}
	
					/* Typo in Havran's thesis:
					   (it specifies "stack[exPt].p == splitVal", which
					    is clearly incorrect) */
					if (stack[enPt].p[axis] == splitVal) {
						/* Case Z1 */
						currNode = currNode->getRight();
						continue;
					}
	
					/* Case N4 */
					currNode = currNode->getLeft();
					farChild = currNode + 1; // getRight()
				} else { /* stack[enPt].p[axis] > splitVal */
					if (splitVal < stack[exPt].p[axis]) {
						/* Cases P1, P2, P3 and N5 */
						currNode = currNode->getRight();
						continue;
					}
					/* Case P4 */
					farChild = currNode->getLeft();
					currNode = farChild + 1; // getRight()
				}
	
				/* Cases P4 and N4 -- calculate the distance to the split plane */
				Float distToSplit = (splitVal - ray.o[axis]) * ray.dRcp[axis];
	
				/* Set up a new exit point */
				const uint32_t tmp = exPt++;
				if (exPt == enPt) /* Do not overwrite the entry point */
					++exPt;
	
				KDAssert(exPt < MTS_KD_MAXDEPTH);
				stack[exPt].prev = tmp;
				stack[exPt].t = distToSplit;
				stack[exPt].node = farChild;
	
				#if 1
				/* Intrestingly, this appears to be faster than the 
				   original code with the prevAxis & nextAxis table */
				stack[exPt].p = ray(distToSplit);
				stack[exPt].p[axis] = splitVal;
				#else
				const int nextAxis = nextAxisTable[axis];
				const int prevAxis = prevAxisTable[axis];
				stack[exPt].p[axis] = splitVal;
				stack[exPt].p[nextAxis] = ray.o[nextAxis]
					+ distToSplit*ray.d[nextAxis];
				stack[exPt].p[prevAxis] = ray.o[prevAxis]
					+ distToSplit*ray.d[prevAxis];
				#endif
	
			}
	
			/* Reached a leaf node */
			for (index_type entry=currNode->getPrimStart(),
					last = currNode->getPrimEnd(); entry != last; entry++) {
				const index_type primIdx = m_indices[entry];
	
				#if defined(MTS_KD_MAILBOX_ENABLED)
				if (mailbox.contains(primIdx)) 
					continue;
				#endif
	
				bool result;
				if (!shadowRay)
					result = cast()->intersect(ray, primIdx, mint, maxt, t, temp);
				else
					result = cast()->intersect(ray, primIdx, mint, maxt);
	
				if (result) {
					if (shadowRay)
						return true;
					maxt = t;
					foundIntersection = true;
				}
	
				#if defined(MTS_KD_MAILBOX_ENABLED)
				mailbox.put(primIdx);
				#endif
			}
	
			if (stack[exPt].t > maxt) 
				break;
	
			/* Pop from the stack and advance to the next node on the interval */
			enPt = exPt;
			currNode = stack[exPt].node;
			exPt = stack[enPt].prev;
		}
	
		return foundIntersection;
	}

	/**
	 * \brief Internal kd-tree traversal implementation (Havran variant)
	 *
	 * This method is almost identical to \ref rayIntersectHavran, except
	 * that it additionally returns statistics on the number of traversed
	 * nodes, intersected shapes, as well as the time taken to do this
	 * (measured using rtdsc).
	 */
	FINLINE boost::tuple<bool, uint32_t, uint32_t, uint64_t> 
			rayIntersectHavranCollectStatistics(const Ray &ray, 
			Float mint, Float maxt, Float &t, void *temp) const {
		KDStackEntryHavran stack[MTS_KD_MAXDEPTH];
	
		/* Set up the entry point */
		uint32_t enPt = 0;
		stack[enPt].t = mint;
		stack[enPt].p = ray(mint);
	
		/* Set up the exit point */
		uint32_t exPt = 1;
		stack[exPt].t = maxt;
		stack[exPt].p = ray(maxt);
		stack[exPt].node = NULL;
	
		uint32_t numTraversals = 0;
		uint32_t numIntersections = 0;
		uint64_t timer = rdtsc();
		bool foundIntersection = false;
	
		const KDNode * __restrict currNode = m_nodes;
		while (currNode != NULL) {
			while (EXPECT_TAKEN(!currNode->isLeaf())) {
				const Float splitVal = (Float) currNode->getSplit();
				const int axis = currNode->getAxis();
				const KDNode * __restrict farChild;
	
				++numTraversals;
				if (stack[enPt].p[axis] <= splitVal) {
					if (stack[exPt].p[axis] <= splitVal) {
						/* Cases N1, N2, N3, P5, Z2 and Z3 (see thesis) */
						currNode = currNode->getLeft();
						continue;
					}
	
					/* Typo in Havran's thesis:
					   (it specifies "stack[exPt].p == splitVal", which
					    is clearly incorrect) */
					if (stack[enPt].p[axis] == splitVal) {
						/* Case Z1 */
						currNode = currNode->getRight();
						continue;
					}
	
					/* Case N4 */
					currNode = currNode->getLeft();
					farChild = currNode + 1; // getRight()
				} else { /* stack[enPt].p[axis] > splitVal */
					if (splitVal < stack[exPt].p[axis]) {
						/* Cases P1, P2, P3 and N5 */
						currNode = currNode->getRight();
						continue;
					}
					/* Case P4 */
					farChild = currNode->getLeft();
					currNode = farChild + 1; // getRight()
				}
	
				/* Cases P4 and N4 -- calculate the distance to the split plane */
				t = (splitVal - ray.o[axis]) * ray.dRcp[axis];
	
				/* Set up a new exit point */
				const uint32_t tmp = exPt++;
				if (exPt == enPt) /* Do not overwrite the entry point */
					++exPt;
	
				KDAssert(exPt < MTS_KD_MAXDEPTH);
				stack[exPt].prev = tmp;
				stack[exPt].t = t;
				stack[exPt].node = farChild;
				stack[exPt].p = ray(t);
				stack[exPt].p[axis] = splitVal;
			}
	
			/* Reached a leaf node */
			for (unsigned int entry=currNode->getPrimStart(),
					last = currNode->getPrimEnd(); entry != last; entry++) {
				const index_type primIdx = m_indices[entry];
	
				++numIntersections;
				bool result = cast()->intersect(ray, primIdx, mint, maxt, t, temp);
	
				if (result) {
					maxt = t;
					foundIntersection = true;
				}
			}
	
			if (stack[exPt].t > maxt) 
				break;
	
			/* Pop from the stack and advance to the next node on the interval */
			enPt = exPt;
			currNode = stack[exPt].node;
			exPt = stack[enPt].prev;
		}
	
		return boost::make_tuple(foundIntersection, numTraversals, 
				numIntersections, rdtsc() - timer);
	}

	/**
	 * \brief Ray tracing kd-tree traversal loop (PBRT variant)
	 */
	template<bool shadowRay> FINLINE bool rayIntersectPBRT(const Ray &ray, 
			Float mint_, Float maxt_, Float &t, void *temp) const {
		KDStackEntry stack[MTS_KD_MAXDEPTH];
		int stackPos = 0;
		Float mint = mint_, maxt=maxt_;
		const KDNode *node = m_nodes;
		bool foundIntersection = false;
	
		while (node != NULL) {
			if (maxt_ < mint)
				break;
	
			if (EXPECT_TAKEN(!node->isLeaf())) {
				const Float split = (Float) node->getSplit();
				const int axis = node->getAxis();
				const float tPlane = (split - ray.o[axis]) * ray.dRcp[axis];
				bool leftOfSplit = (ray.o[axis] < split)
					|| (ray.o[axis] == split && ray.d[axis] <= 0);
	
				const KDNode * __restrict left = node->getLeft();
				const KDNode * __restrict right = left + 1;
				const KDNode * __restrict first  = leftOfSplit ? left : right;
				const KDNode * __restrict second = leftOfSplit ? right : left;
	
				if (tPlane > maxt || tPlane <= 0) {
					node = first;
				} else if (tPlane < mint) {
					node = second;
				} else {
					stack[stackPos].node = second;
					stack[stackPos].mint = tPlane;
					stack[stackPos].maxt = maxt;
					++stackPos;
					node = first;
					maxt = tPlane;
				}
			} else {
				for (unsigned int entry=node->getPrimStart(),
						last = node->getPrimEnd(); entry != last; entry++) {
					const index_type primIdx = m_indices[entry];
	
					bool result;
					if (!shadowRay)
						result = cast()->intersect(ray, primIdx, mint, maxt, t, temp);
					else
						result = cast()->intersect(ray, primIdx, mint, maxt);
	
					if (result) {
						if (shadowRay)
							return true;
						maxt_ = t;
						foundIntersection = true;
					}
				}
	
				if (stackPos > 0) {
					--stackPos;
					node = stack[stackPos].node;
					mint = stack[stackPos].mint;
					maxt = stack[stackPos].maxt;
				} else {
					break;
				}
			}
		}
		return foundIntersection;
	}

	/**
	 * \brief Ray tracing kd-tree traversal loop (Plain variant)
	 *
	 * This implementation represents the simplest possible KD-tree traversal
	 * loop; it is generally not robust enough for rendering, since it doesn't
	 * handle various corner cases and furthermore suffers from roundoff errors.
	 */
	template<bool shadowRay> FINLINE bool rayIntersectPlain(const Ray &ray, 
		Float mint_, Float maxt_, Float &t, void *temp) const {
		KDStackEntry stack[MTS_KD_MAXDEPTH];
		int stackPos = 0;
		Float mint = mint_, maxt = maxt_;
		const KDNode *node = m_nodes;
	
		const int signs[3] = {
			ray.d[0] >= 0 ? 1 : 0,
			ray.d[1] >= 0 ? 1 : 0,
			ray.d[2] >= 0 ? 1 : 0
		};
	
		while (node != NULL) {
			while (EXPECT_TAKEN(!node->isLeaf())) {
				const Float split = (Float) node->getSplit();
				const int axis = node->getAxis();
				const float t = (split - ray.o[axis]) * ray.dRcp[axis];
	
				int sign = signs[axis];
	
				const KDNode * __restrict base = node->getLeft();
				const KDNode * __restrict first = base + (sign^1);
				const KDNode * __restrict second = base + sign;
	
				node = first;
	
				bool front = t > maxt,
					 back = t < mint;
	
				if (front) {
					;
				} else if (back) {
					node = second;
				} else {
					stack[stackPos].node = second;
					stack[stackPos].mint = t;
					stack[stackPos].maxt = maxt;
					maxt = t;
					++stackPos;
				}
			}
			bool foundIntersection = false;
			for (unsigned int entry=node->getPrimStart(),
					last = node->getPrimEnd(); entry != last; entry++) {
				const index_type primIdx = m_indices[entry];

				bool result;
				if (!shadowRay) {
					result = cast()->intersect(ray, primIdx, mint, maxt, t, temp);
				} else {
					result = cast()->intersect(ray, primIdx, mint, maxt);
				}
	
				if (result) {
					if (shadowRay)
						return true;
					maxt = t;
					foundIntersection = true;
				}
			}
	
			if (foundIntersection) 
				return true;
	
			if (stackPos > 0) {
				--stackPos;
				node = stack[stackPos].node;
				mint = stack[stackPos].mint;
				maxt = stack[stackPos].maxt;
			} else {
				break;
			}
		}
		return false;
	}

public:
	/**
	 * \brief Empirically find the best traversal and intersection 
	 * cost values
	 *
	 * This is done by running the traversal code on random rays 
	 * and fitting the SAH cost model to the collected statistics.
	 */
	void findCosts(Float &traversalCost, Float &intersectionCost) {
		ref<Random> random = new Random();
		uint8_t temp[128];
		BSphere bsphere = m_aabb.getBSphere();
		int nRays = 10000000, warmup = nRays/4;
		Vector *A = new Vector[nRays-warmup];
		Float *b = new Float[nRays-warmup];
		int nIntersections = 0, idx = 0;
	
		for (int i=0; i<nRays; ++i) {
			Point2 sample1(random->nextFloat(), random->nextFloat()),
				sample2(random->nextFloat(), random->nextFloat());
			Point p1 = bsphere.center + squareToSphere(sample1) * bsphere.radius;
			Point p2 = bsphere.center + squareToSphere(sample2) * bsphere.radius;
			Ray ray(p1, normalize(p2-p1), 0.0f);
			Float mint, maxt, t;
			if (m_aabb.rayIntersect(ray, mint, maxt)) {
				if (ray.mint > mint) mint = ray.mint;
				if (ray.maxt < maxt) maxt = ray.maxt;
				if (EXPECT_TAKEN(maxt > mint)) {
					boost::tuple<bool, uint32_t, uint32_t, uint64_t> statistics =
						rayIntersectHavranCollectStatistics(ray, mint, maxt, t, temp);
					if (boost::get<0>(statistics))
						nIntersections++;
					if (i > warmup) {
						A[idx].x = 1;
						A[idx].y = (Float) boost::get<1>(statistics);
						A[idx].z = (Float) boost::get<2>(statistics);
						b[idx]   = (Float) boost::get<3>(statistics);
						idx++;
					}
				}
			}
		}
	
		KDLog(EDebug, "Fitting to " SIZE_T_FMT " samples (" SIZE_T_FMT 
				" intersections)", idx, nIntersections);
	
		/* Solve using normal equations */
		Matrix4x4 M(0.0f), Minv;
		Vector4 rhs(0.0f), x;
	
		for (int i=0; i<3; ++i) {
			for (int j=0; j<3; ++j) 
				for (int k=0; k<idx; ++k) 
					M.m[i][j] += A[k][i]*A[k][j];
			for (int k=0; k<idx; ++k) 
				rhs[i] += A[k][i]*b[k];
		}
		M.m[3][3] = 1.0f;
		bool success = M.invert(Minv);
		SAssert(success);
	
		Transform(Minv, M)(rhs, x);
	
		Float avgRdtsc = 0, avgResidual = 0;
		for (int i=0; i<idx; ++i) {
			avgRdtsc += b[i];
			Float model = x[0] * A[i][0]
				+ x[1] * A[i][1]
				+ x[2] * A[i][2];
			avgResidual += std::abs(b[i] - model);
		}
		avgRdtsc /= idx;
		avgResidual /= idx;
	
		for (int k=0; k<idx; ++k) 
			avgRdtsc += b[k];
		avgRdtsc /= idx;
	
		KDLog(EDebug, "Least squares fit:");
		KDLog(EDebug, "   Constant overhead    = %.2f", x[0]);
		KDLog(EDebug, "   Traversal cost       = %.2f", x[1]);
		KDLog(EDebug, "   Intersection cost    = %.2f", x[2]);
		KDLog(EDebug, "   Average rdtsc value  = %.2f", avgRdtsc);
		KDLog(EDebug, "   Avg. residual        = %.2f", avgResidual);
		x *= 10/x[1];
		KDLog(EDebug, "Re-scaled:");
		KDLog(EDebug, "   Constant overhead    = %.2f", x[0]);
		KDLog(EDebug, "   Traversal cost       = %.2f", x[1]);
		KDLog(EDebug, "   Intersection cost    = %.2f", x[2]);
	
		delete[] A;
		delete[] b;
		traversalCost = x[1];
		intersectionCost = x[2];
	}
};

MTS_NAMESPACE_END

#endif /* __SAH_KDTREE3_H */
