#include "RendererUtil.h"
#include "PathRenderer.h"
#include "modulehandle.h"
#include <wx/filename.h>

using namespace VRUT;
using namespace std;

float hemisph_cosn(Vector3 & dir, const float h, const float r1, const float r2)
{
	const float sin_theta = sqrtf(1.f - powf(r2, 2.f / float(h + 1)));
	const float cos_theta = powf(r2, 1.f / float(h + 1));
	const float phi = 2.f * float(M_PI) * r1;
	const float pdf = (h + 1) / (2.f * M_PI) * powf(cos_theta, h);

	dir.x = cosf(phi) * sin_theta;
	dir.y = sinf(phi) * sin_theta;
	dir.z = cos_theta;

	return pdf;
}

float GlossyRandomVector(Vector3 & res, const Vector3 & R, const float h, const float r1, const float r2)
{
	Vector3 tmp;
	float pdf = hemisph_cosn(tmp, h, r1, r2);

	Vector3 u, v;
	R.RightHandedBase(u, v);

	res = Normalize(tmp.x * u + tmp.y * v + tmp.z * R);
	return pdf;
}

PathRenderer::PathRenderer(RayTracer *_rayTracer) : WhittedRenderer(_rayTracer)
{
	sampleBuffer.resize(1024);

	// TODO Is this correct? Shouldn't this be 'i' instead of '1'? Use v.resize(sampleBuffer.size(), 1) to set '1' everywhere.
	//		Also, use v.resize(sampleBuffer.size()) and then v[i] = i, to avoid reallocations.
	/*vector<unsigned int> v;
	for (size_t i = 0; i < sampleBuffer.size(); i++)
		v.push_back(1);*/
	/*vector<unsigned int> v;
	v.resize(sampleBuffer.size(), 1);

	FPSSampler sampler(v, v);

	for (size_t i = 0; i < sampleBuffer.size(); i++)
		sampleBuffer[i] = Vector2(sampler.samples[i][0].x / 32.0f, sampler.samples[i][0].y / 32.0f);*/
}

void PathRenderer::PathRays(
	const RayStream & inputRayStream, 
	RayStream & outputRayStream, 
	Vector4 bounceColors[], 
	const Vector4 diffuseColors[], 
	Vector4 colors[], 
	const int depth, 
	const int workerIndex)
{
	const RayIntersectionInfo * intersectInfos = inputRayStream.GetRayIntersectionInfos();
	const Ray * rays = inputRayStream.GetRays();
	const Vector4 * weights = inputRayStream.GetWeights();
	const int * pixelIndices = inputRayStream.GetPixelIndices();

	for (unsigned int i = 0; i < inputRayStream.GetNumberOfRays(); ++i)
	{
		if (intersectInfos[i].nodeID != NODE_ID_NONE)
		{
			Vector3 normal = intersectInfos[i].normal;

			const float cosThetaO = normal.Dot(rays[i].direction);
			if (cosThetaO > 0.0f)
				normal = -normal;

			const RTMaterial *material = bvh.GetNodeInfos()[intersectInfos[i].nodeID].material;
			float reflectivityOrIOR = material->reflectivityOrIOR;

			Vector4 weightReflected, weightRefracted;
			processReflectRefractRayGetWeights(rays[i].direction, intersectInfos[i].intersection, normal, material, depth, weights[i], pixelIndices[i], cosThetaO, outputRayStream, bounceColors[i], weightReflected, weightRefracted);

			const bool isReflect = (weightReflected.x + weightReflected.y + weightReflected.z + weightReflected.w) != 0;
			const bool isRefract = (weightRefracted.x + weightRefracted.y + weightRefracted.z + weightRefracted.w) != 0;

			if (!(isReflect | isRefract) && (depth < maxRecursionDepth))
			{
				Vector4 diffuseColor = diffuseColors[i];

				Vector4 specularColor = material->specular;
				float shininess = material->shininess;

#ifdef USE_TEMPORAL
				//leave out albedo with temporal
				//quick hack instead of properly dividing them for now
				if (depth < 1 && useTemporalPathTracing) {
					diffuseColor.x = 1;
					diffuseColor.y = 1;
					diffuseColor.z = 1;
				}
				specularColor *= 0;
#endif

				float maxDiffuse = max(max(diffuseColor.x, diffuseColor.y), diffuseColor.z);
				float maxSpecular = max(max(specularColor.x, specularColor.y), specularColor.z);

				if (maxSpecular + maxDiffuse > 1.0f)
				{
					float s = 1 / (maxSpecular + maxDiffuse);
					diffuseColor *= s;
					specularColor *= s;
					maxSpecular *= s;
					maxDiffuse *= s;
				}

				float r1x = workers[workerIndex]->rng.GetSample();
				float r2x = workers[workerIndex]->rng.GetSample();

				Vector4 weight;
				Vector3 reflection_direction = Normalize(rays[i].direction - (2 * cosThetaO) * normal);
				Ray reflectedRay;

				float x = workers[workerIndex]->rng.GetSample();
				if (x > maxDiffuse + maxSpecular)
				{
					colors[pixelIndices[i]] += weights[i] * bounceColors[i];
					continue; // terminate the path
				}

				bool generateDiffuse = (x < maxDiffuse);

				if (!generateDiffuse)
				{
					// shoot a glossy distributed ray
					Vector3 new_ray_direction;
					float pdf = GlossyRandomVector(new_ray_direction, reflection_direction, shininess, r1x, r2x);
					reflectedRay = Ray(intersectInfos[i].intersection + new_ray_direction * rayOffset, new_ray_direction, FLT_MAX);
					float cosThetaI = DotProd(normal, reflectedRay.direction);
					float cosAlpha = DotProd(reflection_direction, reflectedRay.direction);
					if ((cosThetaI < 0.0f) || (pdf < 0.01f) || (cosAlpha < 0.0f))
						generateDiffuse = true; // wrong ray generated using specular importance sampling -> regenerate it using diffuse disribution
					Vector4 brdf = (shininess + 2.0f) / (2.0f * M_PI) * specularColor * powf(cosAlpha, shininess);
					weight = brdf * cosThetaI / pdf;
				}

				if (generateDiffuse)
				{
					// shoot a cosine distributed ray
					Vector3 new_ray_direction = CosineRandomVector(r1x, r2x, normal);
					reflectedRay = Ray(intersectInfos[i].intersection + new_ray_direction * rayOffset, new_ray_direction, FLT_MAX);
					float cosThetaI = DotProd(normal, reflectedRay.direction);
					float pdf = cosThetaI / M_PI;
					float cosAlpha = DotProd(reflection_direction, reflectedRay.direction);
					Vector4 brdf = diffuseColor / M_PI;
					weight = brdf * cosThetaI / pdf;
				}

				// use Russian roulette to terminate the path and modify the path throughtput
				//						float survivalProbability = pathThroughput[k]*lumWeight;
				//						float survivalProbability = pathThroughput[k]*0.5f;
				float survivalProbability = max(max(weight.x, weight.y), weight.z);
				weight *= weights[i];
				if (depth < rrStartDepth)
					survivalProbability = 1.0f;
				const float MIN_SURVIVAL_PROBABILITY = 0.1f;
				if (survivalProbability < MIN_SURVIVAL_PROBABILITY)
					survivalProbability = MIN_SURVIVAL_PROBABILITY;
				float xx = workers[workerIndex]->rng.GetSample();
				if (xx > survivalProbability)
				{

					// terminate the ray
					continue;
				}

				reflectedRay.t = FLT_MAX;
				outputRayStream.PushRay(reflectedRay, pixelIndices[i], weight / survivalProbability);
			}
		}

		colors[pixelIndices[i]] += weights[i] * bounceColors[i];
	}
}

void PathRenderer::CastAndShade(RayStream & inputRayStream, Vector4 colors[], unsigned int cameraIndex, const int workerIndex)
{
	if (!usePathTracing)
	{
		WhittedRenderer::CastAndShade(inputRayStream, colors, cameraIndex, workerIndex);
	}
	else
	{
		Vector4 *bounceColors = workers[workerIndex]->renderBuffers.casBuffers.bounceColors;
		Vector4 *diffuseColors = workers[workerIndex]->renderBuffers.casBuffers.diffuseColors;
		Vector4 *texColors = workers[workerIndex]->renderBuffers.casBuffers.texColors;
#ifdef USE_TEMPORAL
		Vector4 *normalDepth = workers[workerIndex]->renderBuffers.normalDepth;
		Vector4 *diffuse = workers[workerIndex]->renderBuffers.diffuse;
#endif
		bool *usedTexs = workers[workerIndex]->renderBuffers.casBuffers.usedTexs;

		RayStream * irays = &inputRayStream;
		RayStream * orays = workers[workerIndex]->outputRayStream;
		orays->Clear();

		int depth = 0;
		do {
			do {
				// cast rays
				rayCaster.SetTopCulling(true);
				rayCaster.SetBottomCulling(false);
				rayCaster.CastRayBundle(*irays, bvh, workers[workerIndex]->renderBuffers.rayCasterBuffers);
				// shade
				Shade(*irays, bounceColors, diffuseColors, texColors, usedTexs, cameraIndex, workerIndex);
				// reflected and refracted rays
				//ReflectedAndRefractedRays(*irays, *orays, colors, bounceColors, depth);
				// path rays
				PathRays(*irays, *orays, bounceColors, diffuseColors, colors, depth, workerIndex);

#ifdef USE_TEMPORAL
				//temporal: write normals and depth of primary rays
				if (useTemporalPathTracing && !depth) {
					auto infos = irays->GetRayIntersectionInfos();
					for (unsigned int i = 0; i < irays->GetNumberOfRays(); i++) {

						Vector3 normal = infos[i].normal.Normalize();
						/*Vector2 p = Vector2(normal.x, normal.y) * (1.0 / (abs(normal.x) + abs(normal.y) + abs(normal.z)));
						Vector2 s = Vector2(((normal.x >= 0.0) ? +1.0 : -1.0), (normal.y >= 0.0) ? +1.0 : -1.0);
						Vector2 normalOct = (normal.z <= 0.0) ? (Vector2(s.x*(1 - abs(p.y)), s.y*(1 - abs(p.x)))) : p;*/
						
						//compress normal (spheremap transform: https://aras-p.info/texts/CompactNormalStorage.html#method04spheremap)
						float p = max(0.0001f,sqrt(normal.z * 8 + 8));
						Vector2 normal2 = Vector2(normal.x/p + 0.5,normal.y/p + 0.5);

						normalDepth[i] = Vector4(normal2, 0, infos[i].dist);
						diffuse[i] = Vector4(diffuseColors[i].x, diffuseColors[i].y, diffuseColors[i].z, 0);
					}
				}
#endif

			} while (irays->RefillRays());
			// swap
			std::swap(irays, orays);
			// reset output rays
			orays->Clear();
		} while (depth++ < maxRecursionDepth && irays->GetNumberOfRays() > 0);
	}
}
