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

    Copyright (c) 2007-2012 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/skdtree.h>
#include <mitsuba/core/statistics.h>

#if defined(MTS_SSE)
#include <mitsuba/core/sse.h>
#include <mitsuba/core/aabb_sse.h>
#include <mitsuba/render/triaccel_sse.h>
#endif

MTS_NAMESPACE_BEGIN

ShapeKDTree::ShapeKDTree() {
#if !defined(MTS_KD_CONSERVE_MEMORY)
	m_triAccel = NULL;
#endif
	m_shapeMap.push_back(0);
}

ShapeKDTree::~ShapeKDTree() {
#if !defined(MTS_KD_CONSERVE_MEMORY)
	if (m_triAccel)
		freeAligned(m_triAccel);
#endif
	for (size_t i=0; i<m_shapes.size(); ++i)
		m_shapes[i]->decRef();
}

static StatsCounter raysTraced("General", "Normal rays traced");
static StatsCounter shadowRaysTraced("General", "Shadow rays traced");

void ShapeKDTree::addShape(const Shape *shape) {
	Assert(!isBuilt());
	if (shape->isCompound())
		Log(EError, "Cannot add compound shapes to a kd-tree - expand them first!");
	if (shape->getClass()->derivesFrom(MTS_CLASS(TriMesh))) {
		// Triangle meshes are expanded into individual primitives,
		// which are visible to the tree construction code. Generic
		// primitives are only handled by their AABBs
		m_shapeMap.push_back((SizeType) 
			static_cast<const TriMesh *>(shape)->getTriangleCount());
		m_triangleFlag.push_back(true);
	} else {
		m_shapeMap.push_back(1);
		m_triangleFlag.push_back(false);
	}
	shape->incRef();
	m_shapes.push_back(shape);
}

void ShapeKDTree::build() {
	for (size_t i=1; i<m_shapeMap.size(); ++i)
		m_shapeMap[i] += m_shapeMap[i-1];

	SAHKDTree3D<ShapeKDTree>::buildInternal();
		
#if !defined(MTS_KD_CONSERVE_MEMORY)
	ref<Timer> timer = new Timer();
	SizeType primCount = getPrimitiveCount();
	Log(EDebug, "Precomputing triangle intersection information (%s)",
			memString(sizeof(TriAccel)*primCount).c_str());
	m_triAccel = static_cast<TriAccel *>(allocAligned(primCount * sizeof(TriAccel)));

	IndexType idx = 0;
	for (IndexType i=0; i<m_shapes.size(); ++i) {
		const Shape *shape = m_shapes[i];
		if (m_triangleFlag[i]) {
			const TriMesh *mesh = static_cast<const TriMesh *>(shape);
			const Triangle *triangles = mesh->getTriangles();
			const Point *positions = mesh->getVertexPositions();
			for (IndexType j=0; j<mesh->getTriangleCount(); ++j) {
				const Triangle &tri = triangles[j];
				const Point &v0 = positions[tri.idx[0]];
				const Point &v1 = positions[tri.idx[1]];
				const Point &v2 = positions[tri.idx[2]];
				m_triAccel[idx].load(v0, v1, v2);
				m_triAccel[idx].shapeIndex = i;
				m_triAccel[idx].primIndex = j;
				++idx;
			}
		} else {
			/* Create a 'fake' triangle, which redirects to a Shape */
			memset(&m_triAccel[idx], 0, sizeof(TriAccel));
			m_triAccel[idx].shapeIndex = i;
			m_triAccel[idx].k = KNoTriangleFlag;
			++idx;
		}
	}
	Log(EDebug, "Finished -- took %i ms.", timer->getMilliseconds());
	Log(EDebug, "");
	KDAssert(idx == primCount);
#endif
}

bool ShapeKDTree::rayIntersect(const Ray &ray, Intersection &its) const {
	uint8_t temp[MTS_KD_INTERSECTION_TEMP];
	its.t = std::numeric_limits<Float>::infinity(); 
	Float mint, maxt;

	++raysTraced;
	if (m_aabb.rayIntersect(ray, mint, maxt)) {
		/* Use an adaptive ray epsilon */
		Float rayMinT = ray.mint;
		if (rayMinT == Epsilon) 
			rayMinT *= std::max(std::max(std::max(std::abs(ray.o.x), 
				std::abs(ray.o.y)), std::abs(ray.o.z)), Epsilon);

		if (rayMinT > mint) mint = rayMinT;
		if (ray.maxt < maxt) maxt = ray.maxt;

		if (EXPECT_TAKEN(maxt > mint)) {
			if (rayIntersectHavran<false>(ray, mint, maxt, its.t, temp)) {
				fillIntersectionRecord<true>(ray, temp, its);
				return true;
			}
		}
	}
	return false;
}

bool ShapeKDTree::rayIntersect(const Ray &ray, Float &t, ConstShapePtr &shape, 
		Normal &n, Point2 &uv) const {
	uint8_t temp[MTS_KD_INTERSECTION_TEMP];
	Float mint, maxt;
	
	t = std::numeric_limits<Float>::infinity();

	++shadowRaysTraced;
	if (m_aabb.rayIntersect(ray, mint, maxt)) {
		/* Use an adaptive ray epsilon */
		Float rayMinT = ray.mint;
		if (rayMinT == Epsilon)
			rayMinT *= std::max(std::max(std::abs(ray.o.x), 
				std::abs(ray.o.y)), std::abs(ray.o.z));

		if (rayMinT > mint) mint = rayMinT;
		if (ray.maxt < maxt) maxt = ray.maxt;

		if (EXPECT_TAKEN(maxt > mint)) {
			if (rayIntersectHavran<false>(ray, mint, maxt, t, temp)) {
				const IntersectionCache *cache = reinterpret_cast<const IntersectionCache *>(temp);
				shape = m_shapes[cache->shapeIndex];

				if (m_triangleFlag[cache->shapeIndex]) {
					const TriMesh *trimesh = static_cast<const TriMesh *>(shape);
					const Triangle &tri = trimesh->getTriangles()[cache->primIndex];
					const Point *vertexPositions = trimesh->getVertexPositions();
					const Point2 *vertexTexcoords = trimesh->getVertexTexcoords();
					const uint32_t idx0 = tri.idx[0], idx1 = tri.idx[1], idx2 = tri.idx[2];
					const Point &p0 = vertexPositions[idx0];
					const Point &p1 = vertexPositions[idx1];
					const Point &p2 = vertexPositions[idx2];
					n = normalize(cross(p1-p0, p2-p0));

					if (EXPECT_TAKEN(vertexTexcoords)) {
						const Vector b(1 - cache->u - cache->v, cache->u, cache->v);
						const Point2 &t0 = vertexTexcoords[idx0];
						const Point2 &t1 = vertexTexcoords[idx1];
						const Point2 &t2 = vertexTexcoords[idx2];
						uv = t0 * b.x + t1 * b.y + t2 * b.z;
					} else {
						uv = Point2(0.0f);
					}
				} else {
					/// Uh oh... -- much unnecessary work is done here
					Intersection its;
					its.t = t;
					shape->fillIntersectionRecord(ray, 
						reinterpret_cast<const uint8_t*>(temp) + 8, its);
					n = its.geoFrame.n;
					uv = its.uv;
					if (its.shape)
						shape = its.shape;
				}

				return true;
			}
		}
	}
	return false;
}


bool ShapeKDTree::rayIntersect(const Ray &ray) const {
	Float mint, maxt, t = std::numeric_limits<Float>::infinity();

	++shadowRaysTraced;
	if (m_aabb.rayIntersect(ray, mint, maxt)) {
		/* Use an adaptive ray epsilon */
		Float rayMinT = ray.mint;
		if (rayMinT == Epsilon)
			rayMinT *= std::max(std::max(std::abs(ray.o.x), 
				std::abs(ray.o.y)), std::abs(ray.o.z));

		if (rayMinT > mint) mint = rayMinT;
		if (ray.maxt < maxt) maxt = ray.maxt;

		if (EXPECT_TAKEN(maxt > mint)) 
			if (rayIntersectHavran<true>(ray, mint, maxt, t, NULL)) 
				return true;
	}
	return false;
}

#if defined(MTS_HAS_COHERENT_RT)

/// Ray traversal stack entry for uncoherent ray tracing
struct CoherentKDStackEntry {
	/* Current ray interval */
	RayInterval4 MM_ALIGN16 interval;
	/* Pointer to the far child */
	const ShapeKDTree::KDNode * __restrict node;
};

static StatsCounter coherentPackets("General", "Coherent ray packets");
static StatsCounter incoherentPackets("General", "Incoherent ray packets");

void ShapeKDTree::rayIntersectPacket(const RayPacket4 &packet, 
		const RayInterval4 &rayInterval, Intersection4 &its, void *temp) const {
	CoherentKDStackEntry MM_ALIGN16 stack[MTS_KD_MAXDEPTH];
	RayInterval4 MM_ALIGN16 interval;

	const KDNode * __restrict currNode = m_nodes;
	int stackIndex = 0;

	++coherentPackets;

	/* First, intersect with the kd-tree AABB to determine
	   the intersection search intervals */
	if (!m_aabb.rayIntersectPacket(packet, interval))
		return;

	interval.mint.ps = _mm_max_ps(interval.mint.ps, rayInterval.mint.ps);
	interval.maxt.ps = _mm_min_ps(interval.maxt.ps, rayInterval.maxt.ps);

	SSEVector itsFound( _mm_cmpgt_ps(interval.mint.ps, interval.maxt.ps));
	SSEVector masked(itsFound);
	if (_mm_movemask_ps(itsFound.ps) == 0xF)
		return;

	while (currNode != NULL) {
		while (EXPECT_TAKEN(!currNode->isLeaf())) {
			const uint8_t axis = currNode->getAxis();

			/* Calculate the plane intersection */
			const __m128
				splitVal = _mm_set1_ps(currNode->getSplit()),
				t = _mm_mul_ps(_mm_sub_ps(splitVal, packet.o[axis].ps),
					packet.dRcp[axis].ps);

			const __m128
				startsAfterSplit = _mm_or_ps(masked.ps, 
					_mm_cmplt_ps(t, interval.mint.ps)),
				endsBeforeSplit = _mm_or_ps(masked.ps,
					_mm_cmpgt_ps(t, interval.maxt.ps));

			currNode = currNode->getLeft() + packet.signs[axis][0];

			/* The interval completely completely lies on one side
			   of the split plane */
			if (EXPECT_TAKEN(_mm_movemask_ps(startsAfterSplit) == 15)) {
				currNode = currNode->getSibling();
				continue;
			}

			if (EXPECT_TAKEN(_mm_movemask_ps(endsBeforeSplit) == 15)) 
				continue;

			stack[stackIndex].node = currNode->getSibling();
			stack[stackIndex].interval.maxt =    interval.maxt;
			stack[stackIndex].interval.mint.ps = _mm_max_ps(t, interval.mint.ps);
			interval.maxt.ps =                   _mm_min_ps(t, interval.maxt.ps);
			masked.ps = _mm_or_ps(masked.ps, 	
					_mm_cmpgt_ps(interval.mint.ps, interval.maxt.ps));
			stackIndex++;
		}

		/* Arrived at a leaf node - intersect against primitives */
		const IndexType primStart = currNode->getPrimStart();
		const IndexType primEnd = currNode->getPrimEnd();

		if (EXPECT_NOT_TAKEN(primStart != primEnd)) {
			SSEVector 
				searchStart(_mm_max_ps(rayInterval.mint.ps, 
					_mm_mul_ps(interval.mint.ps, SSEConstants::om_eps.ps))),
				searchEnd(_mm_min_ps(rayInterval.maxt.ps, 
					_mm_mul_ps(interval.maxt.ps, SSEConstants::op_eps.ps)));

			for (IndexType entry=primStart; entry != primEnd; entry++) {
				const TriAccel &kdTri = m_triAccel[m_indices[entry]];
				if (EXPECT_TAKEN(kdTri.k != KNoTriangleFlag)) {
					itsFound.ps = _mm_or_ps(itsFound.ps, 
						mitsuba::rayIntersectPacket(kdTri, packet, searchStart.ps, searchEnd.ps, masked.ps, its));
				} else {
					const Shape *shape = m_shapes[kdTri.shapeIndex];

					for (int i=0; i<4; ++i) {
						if (masked.i[i])
							continue;
						Ray ray;
						for (int axis=0; axis<3; axis++) {
							ray.o[axis] = packet.o[axis].f[i];
							ray.d[axis] = packet.d[axis].f[i];
							ray.dRcp[axis] = packet.dRcp[axis].f[i];
						}
						Float t;

						if (shape->rayIntersect(ray, searchStart.f[i], searchEnd.f[i], t, 
								reinterpret_cast<uint8_t *>(temp)
								+ i * MTS_KD_INTERSECTION_TEMP + 8)) {
							its.t.f[i] = t;
							its.shapeIndex.i[i] = kdTri.shapeIndex;
							its.primIndex.i[i] = KNoTriangleFlag;
							itsFound.i[i] = 0xFFFFFFFF;
						}
					}
				}
				searchEnd.ps = _mm_min_ps(searchEnd.ps, its.t.ps);
			}
		}

		/* Abort if the tree has been traversed or if
		   intersections have been found for all four rays */
		if (_mm_movemask_ps(itsFound.ps) == 0xF || --stackIndex < 0)
			break;

		/* Pop from the stack */
		currNode = stack[stackIndex].node;
		interval = stack[stackIndex].interval;
		masked.ps = _mm_or_ps(itsFound.ps, 
			_mm_cmpgt_ps(interval.mint.ps, interval.maxt.ps));
	}
}
	
void ShapeKDTree::rayIntersectPacketIncoherent(const RayPacket4 &packet, 
		const RayInterval4 &rayInterval, Intersection4 &its4, void *temp) const {

	++incoherentPackets;
	for (int i=0; i<4; i++) {
		Ray ray;
		Float t;
		for (int axis=0; axis<3; axis++) {
			ray.o[axis] = packet.o[axis].f[i];
			ray.d[axis] = packet.d[axis].f[i];
			ray.dRcp[axis] = packet.dRcp[axis].f[i];
		}
		ray.mint = rayInterval.mint.f[i];
		ray.maxt = rayInterval.maxt.f[i];
		uint8_t *rayTemp = reinterpret_cast<uint8_t *>(temp) + i * MTS_KD_INTERSECTION_TEMP;
		if (ray.mint < ray.maxt && rayIntersectHavran<false>(ray, ray.mint, ray.maxt, t, rayTemp)) {
			const IntersectionCache *cache = reinterpret_cast<const IntersectionCache *>(rayTemp);
			its4.t.f[i] = t;
			its4.shapeIndex.i[i] = cache->shapeIndex;
			its4.primIndex.i[i] = cache->primIndex;
			its4.u.f[i] = cache->u;
			its4.v.f[i] = cache->v;
		}
	}
}

#endif

MTS_IMPLEMENT_CLASS(ShapeKDTree, false, KDTreeBase)
MTS_NAMESPACE_END