#include "AbstractRenderer.h"
#include "RayStream.h"

using namespace VRUT;

void RayStream::TransformDownUnit(unsigned int numberOfPackets)
{
	bottomBundle.Clear();
	for (unsigned int packetIndex = 0; packetIndex < numberOfPackets; ++packetIndex)
	{
		RayBundle::Point tOx, tOy, tOz;
		RayBundle::Point tDx, tDy, tDz;
		float t[MAX_RAYS_IN_PACKET];
		topBundle.TransformRayUnit(packetIndex, t, tOx, tOy, tOz, tDx, tDy, tDz);
		for (unsigned int rayIndex = 0; rayIndex < MAX_RAYS_IN_PACKET; ++rayIndex)
		{
			Ray ray;
			ray.origin.x = tOx.v[rayIndex];
			ray.origin.y = tOy.v[rayIndex];
			ray.origin.z = tOz.v[rayIndex];
			ray.direction.x = tDx.v[rayIndex];
			ray.direction.y = tDy.v[rayIndex];
			ray.direction.z = tDz.v[rayIndex];
			ray.t = t[rayIndex];
			bottomBundle.PushRay(ray);
		}
	}
}

void RayStream::TransformDown(const Matrix & matrix, unsigned int numberOfPackets)
{
	bottomBundle.Clear();
	for (unsigned int packetIndex = 0; packetIndex < numberOfPackets; ++packetIndex)
	{
		RayBundle::Point tOx, tOy, tOz;
		RayBundle::Point tDx, tDy, tDz;
		float t[MAX_RAYS_IN_PACKET];
		topBundle.TransformRay(matrix, packetIndex, t, tOx, tOy, tOz, tDx, tDy, tDz);
		for (unsigned int rayIndex = 0; rayIndex < MAX_RAYS_IN_PACKET; ++rayIndex)
		{
			Ray ray;
			ray.origin.x = tOx.v[rayIndex];
			ray.origin.y = tOy.v[rayIndex];
			ray.origin.z = tOz.v[rayIndex];
			ray.direction.x = tDx.v[rayIndex];
			ray.direction.y = tDy.v[rayIndex];
			ray.direction.z = tDz.v[rayIndex];
			ray.t = t[rayIndex];
			bottomBundle.PushRay(ray);
		}
	}
}

RayStream::RayStream()
{
	rayData.reserve(3 * MAX_RAYS_IN_BUNDLE * (1 << MAX_RECURSION_DEPTH));
}

RayStream::~RayStream()
{
}

const Ray * RayStream::GetRays() const
{
	return topBundle.GetRays();
}

const RayIntersectionInfo * RayStream::GetRayIntersectionInfos() const
{
	return topBundle.GetRayIntersectionInfos();
}

const Vector4 * RayStream::GetWeights() const
{
	return weights;
}

const int * RayStream::GetPixelIndices() const
{
	return pixelIndices;
}

RayBundle & RayStream::GetTopBundle()
{
	return topBundle;
}

RayBundle & RayStream::GetBottomBundle()
{
	return bottomBundle;
}

void RayStream::PushRay(const Ray & ray, int pixelIndex, float weight)
{
	Vector4 weightVec;
	weightVec.x = weight;
	weightVec.y = weight;
	weightVec.z = weight;
	weightVec.w = weight;
	PushRay(ray, pixelIndex, weightVec);
}

void RayStream::PushRay(const Ray & ray, int pixelIndex, const Vector4 & weight)
{
	int index = topBundle.GetNumberOfRays();
	if (index < MAX_RAYS_IN_BUNDLE)
	{
		weights[index] = weight;
		pixelIndices[index] = pixelIndex;
		topBundle.PushRay(ray);
	}
	else
	{
		rayData.push_back(RayData());
		rayData.back().ray = ray;
		rayData.back().pixelIndex = pixelIndex;
		rayData.back().weight = weight;
		LOGDEBUG("<RayStream > Ray has been stored to auxiliary buffer.");
	}
}

void RayStream::Clear()
{
	topBundle.Clear();
	rayData.clear();
}

bool RayStream::RefillRays()
{
	if (rayData.empty())
		return false;
	topBundle.Clear();
	int index = topBundle.GetNumberOfRays();
	while ((index < MAX_RAYS_IN_BUNDLE) && !rayData.empty())
	{
		weights[index] = rayData.back().weight;
		pixelIndices[index] = rayData.back().pixelIndex;
		topBundle.PushRay(rayData.back().ray);
		rayData.pop_back();
		++index;
	}
	return true;
}

void RayStream::TransformDown(const BVHTop::BVHNodeInfo & nodeInfo, unsigned int numberOfPackets)
{
	if (nodeInfo.isMatUnit)
		TransformDownUnit(numberOfPackets);
	else
		TransformDown(nodeInfo.invWorldMat, numberOfPackets);
}

void RayStream::TransformUp(const BVHBottom & bvh, NODE_ID nodeID, const BVHTop::BVHNodeInfo & nodeInfo, unsigned int numberOfPackets)
{
	for (unsigned int packetIndex = 0; packetIndex < numberOfPackets; ++packetIndex)
	{
		RayBundle::RayPacket & p = *topBundle.activePackets[packetIndex];
		for (unsigned int rayIndex = 0; rayIndex < MAX_RAYS_IN_PACKET; ++rayIndex)
		{
			unsigned int i1 = rayIndex + (packetIndex << LOG_MAX_RAYS_IN_PACKET);
			RayIntersectionInfo & info1 = bottomBundle.infos[i1];
			if (info1.dist != FLT_MAX)
			{
				// update info
				unsigned int i0 = p.rayOffsets[rayIndex] / sizeof(Ray);
				RayIntersectionInfo & info0 = topBundle.infos[i0];
				info0.dist = info1.dist;
				info0.index = info1.index;
				info0.u = 1.0f - info1.u - info1.v;
				info0.v = info1.v;
				info0.nodeID = nodeID;
				// intersection point
				Ray & ray = topBundle.rays[i0];
				info0.intersection = ray.origin + ray.direction * info0.dist;
				// normal
				bvh.GetNormalAtUnnormalized(info0.index, info0.u, info0.v, info0.normal);
				if (!nodeInfo.isMatUnit)
					info0.normal = nodeInfo.worldMat.TransformNormal(info0.normal);
				info0.normal = info0.normal.Normalize();
				// update t in packet and ray
				assert(info0.dist <= ray.t);
				ray.t = info0.dist;
				p.t.v[rayIndex] = info0.dist;
			}
		}
	}
}

unsigned int RayStream::TransformUpShadow(const BVHBottom & bvh, const BVHTop::BVHNodeInfo & nodeInfo, MaterialCache & materialCache, Vector4 lightColors[], float rayOffset,
		unsigned int numberOfPackets)
{
	memset((void*)processed, 0, sizeof(bool) * MAX_RAYS_IN_BUNDLE);
	const RTMaterial *material = nodeInfo.material;
	unsigned int unoccludedPackets = 0;
	for (unsigned int packetIndex = 0; packetIndex < numberOfPackets; ++packetIndex)
	{
		RayBundle::RayPacket & p = *topBundle.activePackets[packetIndex];
		unsigned int occluded = 0;
		for (unsigned int rayIndex = 0; rayIndex < MAX_RAYS_IN_PACKET; ++rayIndex)
		{
			unsigned int i0 = p.rayOffsets[rayIndex] / sizeof(Ray);
			if (processed[i0])
				continue;
			processed[i0] = true;
			Ray & ray = topBundle.rays[i0];
			unsigned int i1 = rayIndex + (packetIndex << LOG_MAX_RAYS_IN_PACKET);
			RayIntersectionInfo & info = bottomBundle.infos[i1];
			Vector4 & lightColor = lightColors[i0];
			if ((info.dist < ray.t) && (info.dist + rayOffset > 0.0f))
			{
				if (lightColor.w != 0.0f)
				{
					Vector4 diffuseColor = material->diffuse;
					if (material->flags & RTMaterial::ALPHA_TEST)
					{
						float x, y;
						bool usedTex = true;
						Vector4 texColor;
						Vector2 *texCoord1, *texCoord2, *texCoord3;
						info.u = 1.0f - info.u - info.v;
						bvh.GetTextureCoordsAt(info.index, texCoord1, texCoord2, texCoord3);
						bvh.GetNormalAtUnnormalized(info.index, info.u, info.v, info.normal);
						info.normal = info.normal.Normalize();
						materialCache.GetTriangleTexCoords(info, texCoord1, texCoord2, texCoord3, x, y);
						materialCache.GetMaterialDiffuseColor(*material, ray, info, x, y, rayOffset, diffuseColor, usedTex, texColor);
					}
					if (diffuseColor.w == 1.0f)
						lightColor = Vector4(0.0f, 0.0f, 0.0f, 0.0f);
					lightColor *= TransformColor(diffuseColor);
				}
				if (lightColor.w == 0.0f) 
				{
					occluded |= (1 << rayIndex);
				}
			}
		}
		if (occluded < FULL_HIT) 
			topBundle.activePackets[unoccludedPackets++] = topBundle.activePackets[packetIndex];
	}
	return unoccludedPackets;
}
