#ifndef __RAY_BUNDLE_H__
#define __RAY_BUNDLE_H__

#include "BVHBottom.h"
#include "BVHTop.h"
#include "common.h"
#include "RendererUtil.h"

#define USE_GATHERS 0

#if defined(__AVX__) || defined(__AVX2__)
#define LOG_MAX_RAYS_IN_PACKET 3
#include <immintrin.h>
#else
#define LOG_MAX_RAYS_IN_PACKET 2
#endif

#define MAX_RAYS_IN_BUNDLE 1536
#define MAX_RAYS_IN_PACKET (1 << LOG_MAX_RAYS_IN_PACKET)
#define MAX_PACKETS_IN_BUNDLE ((MAX_RAYS_IN_BUNDLE >> LOG_MAX_RAYS_IN_PACKET) + 8)
#define FULL_HIT ((1 << MAX_RAYS_IN_PACKET) - 1)

namespace VRUT
{
	typedef BVHBottom::TriangleCol TriangleCol;

	class RayBundle
	{

	public:

		union Point
		{
			float v[MAX_RAYS_IN_PACKET]; ///< array of coordinates
#if defined(__AVX__) || defined(__AVX2__)
			__m256 _p; ///< coordinate in format for AVX intrinsics
#else
			__m128 _p; ///< coordinate in format for SSE intrinsics
#endif
		};

		struct VRUT_ALIGNED(64) RayPacket
		{
			Point Ox; Point DInvx;
			Point Oy; Point DInvy;
			Point Oz; Point DInvz;
			Point t; unsigned int rayOffsets[MAX_RAYS_IN_PACKET];
		};

		struct VRUT_ALIGNED(32) TraversalCache
		{
			Point ODx;
			Point ODy;
			Point ODz;
		} traversalCache;

		struct VRUT_ALIGNED(32) IntersectionCache
		{
			Point Dx;
			Point Dy;
			Point Dz;
			Point ha;
			Point hb;
			unsigned int id[MAX_RAYS_IN_PACKET];
			unsigned int mask;
		} intersectionCache;

	private:

		// rays
		unsigned int numberOfRays;
		VRUT_ALIGNED(16) Ray rays[MAX_RAYS_IN_BUNDLE];
		RayIntersectionInfo infos[MAX_RAYS_IN_BUNDLE];
		int octants[MAX_RAYS_IN_BUNDLE];
		
		// packets
		RayPacket packets[MAX_PACKETS_IN_BUNDLE];
		RayPacket * activePackets[MAX_PACKETS_IN_BUNDLE];
		unsigned char intersections[MAX_PACKETS_IN_BUNDLE];
		
		// octants
		unsigned int rayCounts[8];

	public:

		RayBundle(void);
		~RayBundle(void);

		unsigned int GetNumberOfRays(void) const;
		unsigned int GetNumberOfRays(unsigned int octant);
		unsigned int GetNumberOfPackets(unsigned int octant);

		const Ray * GetRays(void) const;
		const unsigned char * GetIntersections(void) const;
		const RayIntersectionInfo * GetRayIntersectionInfos(void) const;

		void PushRay(const Ray & ray);
		void Clear(void);
		void ClearInfos(void);
		
		unsigned int RemoveMissedPackets(unsigned int numberOfPackets);

		void BuildPackets(void);
		void SetActivePackets(unsigned int octant);
		void ShiftPackets(unsigned int numberOfPackets, unsigned int numberOfAllPackets, unsigned int unoccludedPackets);

		void FlushIntersectionCache(void);
		void FlushIntersectionCacheShadow(void);
		void InitIntersectionCache(void);
		void InitTraversalCache(void);

		void ReorderRays(unsigned int numberOfPackets);
		void ReadDir(RayPacket & packet, Point & Dx, Point & Dy, Point & Dz);
		void ReadRay(RayPacket & packet, const unsigned int * packetOffsetBytes);

		void TransformRayUnit(unsigned int packetIndex, float t[], Point & tOx, Point & tOy, Point& tOz, Point & tDx, Point & tDy, Point& tDz);
		void TransformRay(const Matrix & matrix, unsigned int packetIndex, float t[], Point & tOx, Point & tOy, Point& tOz, Point & tDx, Point & tDy, Point& tDz);

		unsigned int IntersectAABBCull(const float * fbox, unsigned int numberOfPackets);
		unsigned int IntersectAABB(const float * fbox, unsigned int numberOfPackets);
		unsigned int IntersectAABB(const AbstractBVH::BVHNode * nodes, unsigned int nodeIndex, unsigned int octant);

		void IntersectTriangles(const TriangleCol * triangleCols, int triangleBegin, int triangleEnd, unsigned int numberOfPackets);
		void IntersectTriangles(const TriangleCol * triangleCols, int triangleBegin, int triangleEnd);

		unsigned int IntersectTrianglesShadow(const TriangleCol * triangleCols, int triangleBegin, int triangleEnd, unsigned int numberOfPackets, unsigned int numberOfAllPackets);
		void IntersectTrianglesShadow(const TriangleCol * triangleCols, int triangleBegin, int triangleEnd);

		friend class RayStream;
		friend class RayCaster;

	};

	inline RayBundle::RayBundle() : numberOfRays(0)
	{
		memset(rayCounts, 0, sizeof(unsigned int) * 8);
	}

	inline RayBundle::~RayBundle()
	{
	}

	inline unsigned int RayBundle::GetNumberOfRays() const
	{
		return numberOfRays;
	}

	inline unsigned int RayBundle::GetNumberOfRays(unsigned int octant)
	{
		assert(octant < 8);
		return rayCounts[octant];
	}

	inline unsigned int RayBundle::GetNumberOfPackets(unsigned int octant)
	{
		assert(octant < 8);
		return IntDivCeilLog(rayCounts[octant], LOG_MAX_RAYS_IN_PACKET);
	}

	inline const Ray * RayBundle::GetRays() const
	{
		return rays;
	}

	inline const unsigned char * RayBundle::GetIntersections() const
	{
		return intersections;
	}

	inline const RayIntersectionInfo * RayBundle::GetRayIntersectionInfos() const
	{
		return infos;
	}

	inline void RayBundle::PushRay(const Ray & ray)
	{
		if (numberOfRays < MAX_RAYS_IN_BUNDLE)
		{
			const unsigned int * d = (const unsigned int*)&ray.direction;
			unsigned int dx = d[0];
			unsigned int dy = d[1];
			unsigned int dz = d[2];
			unsigned int octant = (dx >> 31) | ((dy >> 30) & 2) | ((dz >> 29) & 4);
			rays[numberOfRays].origin = ray.origin;
			rays[numberOfRays].direction = ray.direction;
			rays[numberOfRays].t = ray.t;
			rays[numberOfRays].offset = sizeof(Ray) * numberOfRays;
			octants[numberOfRays] = octant;
			++rayCounts[octant];
			++numberOfRays;
		}
	}

	inline void RayBundle::Clear()
	{
		numberOfRays = 0;
		memset(rayCounts, 0, sizeof(unsigned int) * 8);
	}

	inline void RayBundle::ClearInfos()
	{
		for (unsigned int rayIndex = 0; rayIndex < numberOfRays; ++rayIndex)
		{
			infos[rayIndex].dist = FLT_MAX;
			infos[rayIndex].nodeID = NODE_ID_NONE;
		}
	}

	inline unsigned int RayBundle::RemoveMissedPackets(unsigned int numberOfPackets)
	{
		unsigned int hitPackets = 0;
		while (true)
		{
			// skip in-place hits at beginning
			while(intersections[hitPackets])
			{
				++hitPackets;
				if(hitPackets == numberOfPackets)
					return numberOfPackets;
			}

			// skip in-place misses at end
			size_t mask;
			do
			{
				--numberOfPackets;
				if(hitPackets == numberOfPackets)
					return numberOfPackets;
				mask = intersections[numberOfPackets];

			} while( !mask );

			RayPacket * h = activePackets[hitPackets];
			RayPacket * m = activePackets[numberOfPackets];
			activePackets[hitPackets] = m;
			activePackets[numberOfPackets] = h;
			intersections[hitPackets] = (unsigned char)mask;
		} 
	}

	inline void RayBundle::BuildPackets()
	{
		// counts into offsets via prefix sum
		unsigned int octantOffsets[8];
		octantOffsets[0] = 0;
		for (unsigned int i = 1; i < 8; ++i)
			octantOffsets[i] = octantOffsets[i - 1] + rayCounts[i - 1];

		// bin rays by octant
		unsigned short int indices[MAX_RAYS_IN_BUNDLE];
		for (unsigned int i = 0; i < numberOfRays; ++i) 
			indices[octantOffsets[octants[i]]++] = (unsigned short)i;

		unsigned int packetOffset = 0;
		for (int octant = 0; octant < 8; ++octant)
		{
			unsigned int rayCount = rayCounts[octant];
			if (!rayCount)
				continue;

			unsigned int rayOffset = octantOffsets[octant] - rayCount;
			unsigned int packetCount = IntDivCeilLog(rayCount, LOG_MAX_RAYS_IN_PACKET);

			// build a padded ID list
			unsigned int rayOffsetsBytes[MAX_RAYS_IN_BUNDLE];
			for (unsigned int k = 0; k < rayCount; ++k)
				rayOffsetsBytes[k] = indices[k + rayOffset] * sizeof(Ray);

			unsigned int lastRayOffset = rayOffsetsBytes[rayCount - 1];
			while (rayCount & (MAX_RAYS_IN_PACKET - 1))
				rayOffsetsBytes[rayCount++] = lastRayOffset;

			for (unsigned int p = 0; p < packetCount; ++p)
				ReadRay(packets[p + packetOffset], rayOffsetsBytes + MAX_RAYS_IN_PACKET * p);

			packetOffset += packetCount;
		}

	}

	inline void RayBundle::SetActivePackets(unsigned int octant)
	{
		assert(octant < 8);
		unsigned int packetOffset = 0;
		for (unsigned int i = 1; i <= octant; ++i)
			packetOffset += GetNumberOfPackets(i - 1);
		for (unsigned int packetIndex = 0; packetIndex < GetNumberOfPackets(octant); ++packetIndex)
			activePackets[packetIndex] = &packets[packetOffset + packetIndex];
	}

	inline void RayBundle::ShiftPackets(unsigned int numberOfPackets, unsigned int numberOfAllPackets, unsigned int unoccludedPackets)
	{
		unsigned int droppedPackets = numberOfPackets - unoccludedPackets;
		for (unsigned int packetIndex = numberOfPackets; packetIndex < numberOfAllPackets; ++packetIndex)
			activePackets[packetIndex - droppedPackets] = activePackets[packetIndex];
	}

	inline void RayBundle::FlushIntersectionCache()
	{
		RayPacket & packet = *activePackets[0];
		float * pt = (float*)&packet.t;
		// now write results back out
		unsigned int nHit = intersectionCache.mask;
		for (int i = 0; i < MAX_RAYS_IN_PACKET; ++i)
		{
			if ((1 << i) & nHit)
			{
				unsigned int id = packet.rayOffsets[i] / sizeof(Ray);
				rays[id].t = pt[i];
				infos[id].dist = pt[i];
				infos[id].u = intersectionCache.ha.v[i];
				infos[id].v = intersectionCache.hb.v[i];
				infos[id].index = intersectionCache.id[i];
			}
		}
	}

	inline void RayBundle::FlushIntersectionCacheShadow()
	{
		RayPacket & packet = *activePackets[0];
		unsigned int nHit = intersectionCache.mask;
		for (int i = 0; i < MAX_RAYS_IN_PACKET; ++i)
		{
			if ((1 << i) & nHit)
			{
				unsigned int id = packet.rayOffsets[i] / sizeof(Ray);
				rays[id].t = 0.0f;
				infos[id].dist = 0.0f;
			}
		}
	}

	inline void RayBundle::InitIntersectionCache()
	{
		intersectionCache.mask = 0;
		ReadDir(*activePackets[0], intersectionCache.Dx, intersectionCache.Dy, intersectionCache.Dz);
	}
}

#if defined(__AVX__) || defined(__AVX2__)
#include "RayBundleAVX.inl"
#else
#include "RayBundleSSE.inl"
#endif

#endif /* __RAY_BUNDLE_H__ */
