/* -----------------------------------------------------------------------------
	GSFramework
	Copyright 2001-2013 Emmanuel Julien. All Rights Reserved.
----------------------------------------------------------------------------- */


		#include "raytracer/raytracer_core.h"
		#include "raytracer/raytracer_job.h"
		#include "scene3d/mobject.h"
		#include "scene3d/mlight.h"
		#include "scene3d/mcamera.h"
		#include "scene3d/scene.h"
		#include "rand/rand.h"
		#include "platform.h"

		using namespace GS;
		using namespace GS::Core;
		using namespace GS::Raytrace;


//------------------------------------------------------------------------------
float				Raytracer::Fresnel(const Vector4 &v, const Vector4 &np, float eta)
{
	float const	r0 = Math::Pow(1.0f - eta, 2.0f) / Math::Pow(1.0f + eta, 2.0f);
	// Light vector and normal are assumed to be normalized.
	return Types::Clamp <float> (r0 + (1.0f - r0) * Math::Pow(1 - Types::Abs(v.Dot(np)), 5.0f), 0.0f, 1.0f);
}
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
float				Raytracer::ShadowFeel(const Vector4 &s, const Vector4 &d, float l, int r)
{
	float k_shadow = 1;

	if	(configuration.trace_transparency)
	{
		if	(!r)
			return k_shadow;

		// Get closest hit.
		Trace trace;
		scene_shadow_tree.RaytraceScene(trace, s, d, l);
		statistics.ray_count++;
		statistics.tri_test += trace.tri_test;

		if	(trace.has_i && (trace.i_t > 0))
		{
			// Check opacity.
			float opacity = SampleMaterialOpacity(trace);

			// Early exit on fully opaque hit.
			if	(opacity == 1)
				return 0;
			k_shadow = 1 - opacity;

			// Recurse.
			Vector4 offset_pi = trace.s + trace.d * (trace.i_t + Units::Mm(1));
			k_shadow *= ShadowFeel(offset_pi, trace.d, l - Vector4::Dist(trace.s, offset_pi), --r);
		}
	}
	else
	{
		// Any hit within range will do.
		Trace trace(false);
		scene_shadow_tree.RaytraceScene(trace, s, d, l);

		if	(trace.has_i && (trace.i_t > 0))
			return 0;	// Occluded.
	}

	return k_shadow;
}
void				Raytracer::ComputeRadiance(Trace &trace, Color &o, Bounce &bounce)
{
	bool use_fixed_function = trace.m->shader.IsEmpty();
	bool blend_additive = trace.m->blendop == Material::Blend_Add;

	// Evaluate material alpha.
	float alpha;
	if	(use_fixed_function)
			alpha = SampleMaterialOpacity(trace);
	else	alpha = SampleMaterialSink(trace, ShaderTree::SinkOpacity).x;
	alpha *= trace.o->opacity;

	// Compute direct lighting, if the material does not care about the alpha test, or if the material cares about it and its alpha is up to the threshold.
	if	(!(trace.m->renderword & Material::Render_AlphaTest) || alpha > trace.m->athreshold)
	{
		// Evaluate material glossiness.
		float glossiness;
		if	(use_fixed_function)
				glossiness = trace.m->glossiness;
		else	glossiness = SampleMaterialSink(trace, ShaderTree::SinkGlossiness).x;

		// Evaluate light contribution.
		Color l_diff(0, 0, 0), l_spec(0, 0, 0);
		Vector4 offset_pi = trace.pi + trace.n * Units::Mm(1.f);

		for	(uint n = 0; n < lgt.GetCount(); ++n)
			if	(S3D::MLight *l = lgt[n].l)
			{
				Core::Light *light = (Core::Light *)l;
				float k_shadow = 1.f;

				if	((light->shadow != Core::Light::Shadow_None) && configuration.trace_shadow)
				{
					Vector4 d;

					switch (light->model)
					{
						default:
						case Core::Light::Model_Point:
							d = light->GetMatrix().GetRow(3) - offset_pi;
							break;

						case Core::Light::Model_Linear:
							d = light->GetMatrix().GetRow(2).Reversed() * light->clip_distance;
							break;
					}

					if	(d.Dot(trace.n) > 0)
					{
						float l = d.Len();
						d /= l;

						k_shadow = ShadowFeel(offset_pi, d, l, configuration.trace_shadow_transparency_max_recursion);
						if	(!k_shadow)
							continue;
					}
				}

				// Compute contribution.
				float k_d, k_s;
				if	(light->SampleEnergy(trace.pi, trace.n, &k_d, &k_s, &trace.d, glossiness))
				{
					l_diff += light->diffuse_color * light->diffuse_intensity * k_d * k_shadow;
					l_spec += light->specular_color * light->specular_intensity * k_s * k_shadow;
				}
			}

		// Compute indirect lighting.
		Color l_indirect(0, 0, 0), ambient(0, 0, 0);

		if	(configuration.trace_gi && bounce.indirect)
		{
			bounce.indirect--;

			Spread &mc = monte_carlo[Random::Rand(32)];
			Matrix3 nm(Matrix3::FromOrthonormalBasis(trace.n));
			Color l;

			// divide by the number of bounce, to avoid full bounce each time.
			int count_spread = mc.spread.GetCount();
			if	(configuration.indirect_gi_bounce - bounce.indirect != 0)
				count_spread /= configuration.indirect_gi_bounce - bounce.indirect + 1;
			count_spread = Types::Max(count_spread, 1);

			for	(int n = 0; n < count_spread; ++n)
			{
				Bounce ibounce;

				ibounce.indirect = bounce.indirect;
				ibounce.reflection = 0;
				ibounce.refraction = 0;

				Raytrace(RayGrid(offset_pi, mc.spread[n] * nm), l, ibounce);
				l_indirect += l;
			}
			l_indirect /= (float)count_spread;
		}
		else
			ambient = (configuration.gi_use_ambient || !configuration.trace_gi) ? scene->ambient_color * scene->ambient_intensity : Vector4(0.f, 0.f, 0.f);

		// Compute ambient occlusion.
 		float ambient_occlusion = 1.0f;

      	if	((alpha >= 1.0f) && configuration.ao_activate)
   		{
  			Spread &mc = monte_carlo[Random::Rand(32)];
  			Matrix3 nm(Matrix3::FromOrthonormalBasis(trace.n));
  
  			float countouch = 0.0f;
  			float lengthmax = configuration.ao_length;
			float divlengthmaxsq = 1.0f / lengthmax;
  
  			for	(uint n = 0; n < mc.spread.GetCount(); ++n)
  			{
  				// Create the direction vector from the normal of the point with a bit of random.
    			Trace traceOcclusion;
				Vector4 start(trace.pi + mc.spread[n] * nm * Units::Mm(1.f));
/*
				nVector		DirVect(mc.spread[n] * nm);
				scene_tree.RaytraceScene(traceOcclusion, start, DirVect, lengthmax);

				// check the raytrace pass if the alpha of the map and continue to raytrace then				
				float		alphaOcclusion = 0.0f;
				float		current_length = 0.0f;

				while(alphaOcclusion < 1.0f && current_length < lengthmax && 
					traceOcclusion.has_i && (traceOcclusion.i_t > 0.0f))
				{
					current_length += traceOcclusion.i_t;

					// Compute intersection point and fetch material.
					traceOcclusion.pi = traceOcclusion.s + traceOcclusion.d * traceOcclusion.i_t;
					traceOcclusion.m = traceOcclusion.g->material_table[traceOcclusion.g->pol[traceOcclusion.ip].material];

					bool		use_fixed_functionOcclusion = trace.m->shader_tree == NULL ? true : false;

					// Evaluate material alpha.
					float		TempAlphaOcclusion = 0.0f;

					if	(use_fixed_functionOcclusion)
						TempAlphaOcclusion = SampleMaterialOpacity(traceOcclusion);
					else	TempAlphaOcclusion = SampleMaterialSink(traceOcclusion, nShaderTree::SinkOpacity).x;
					alphaOcclusion += TempAlphaOcclusion*traceOcclusion.o->opacity;

					if(alphaOcclusion < 1.0f && current_length < lengthmax)
						scene_tree.RaytraceScene(traceOcclusion, traceOcclusion.pi + DirVect* Mm(1), DirVect, lengthmax - current_length);
				}

				if(alphaOcclusion > 1.0f)
					alphaOcclusion = 1.0f;

				if	(alphaOcclusion > 0)
					countouch += (1.0f - Types::Clamp(current_length * divlengthmaxsq, 0.0f, 1.0f))* alphaOcclusion;
*/
  				scene_tree.RaytraceScene(traceOcclusion, start, mc.spread[n] * nm, lengthmax);
  
  				if	(traceOcclusion.has_i && (traceOcclusion.i_t > 0.0f))
  					countouch += 1.0f - Types::Clamp(traceOcclusion.i_t * divlengthmaxsq, 0.0f, 1.0f);
  			}
      
 			if	(countouch > 0.0f)
 				ambient_occlusion = 1.0f - countouch / mc.spread.GetCount();
 			ambient_occlusion = Types::Clamp(ambient_occlusion);
      	}

		// Sample attributes.
		Color diffuse, specular, self;

		if	(use_fixed_function)
		{
			// Gather attributes.
			diffuse = SampleMaterialAttribute(trace, Channel_Diffuse);
			specular = SampleMaterialAttribute(trace, Channel_Specular);
			self = SampleMaterialAttribute(trace, Channel_SelfIllum);

			// Vertex color.
			if	(trace.m->GetChannelStage(Channel_Light))
			{
				Color color = SampleMaterialAttribute(trace, Channel_Light);
				diffuse *= color;
				specular *= color;
			}
			else	if	(trace.m->renderword & Material::Render_VertexColor)
			{
				Color color = SampleGeometryAttribute(trace, GeometryVertexColor);
				diffuse *= color;
				specular *= color;
			}

			// Environment mapping.
			if	(trace.m->GetChannelStage(Channel_Reflection))
			{
				Color color = SampleMaterialAttribute(trace, Channel_Reflection);
				switch	(trace.m->GetChannelStage(Channel_Reflection)->op)
				{
					case Material::Operator_Multiply:
						diffuse *= color;
						break;

					case Material::Operator_Default:
					case Material::Operator_Add:
						diffuse += color;
						break;	
				}
			}
		}
		else
		{
			diffuse = SampleMaterialSink(trace, ShaderTree::SinkDiffuse);
			specular = SampleMaterialSink(trace, ShaderTree::SinkSpecular);
			self = SampleMaterialSink(trace, ShaderTree::SinkConstant);
		}

		// Final color.
		o = ((diffuse * (l_diff + l_indirect + ambient* ambient_occlusion)) + specular * l_spec + self) /** alpha*/;	// Don't multiply the alpha, because there is real raytracing for the refraction after.
	}
	else
	{
		alpha = 0;
		o.Set(0, 0, 0);
	}

	// Apply fog.
	if	(scene->fog_far > 0)
	{
		float kfog = Types::Clamp((trace.td - scene->fog_near) / (scene->fog_far - scene->fog_near));
		o = o * (1.f - kfog) + scene->fog_color * kfog;
	}

	// Trace reflected and transmitted rays as required.
	float krefl = alpha;

	float eta = trace.m->irefraction;
	if	(trace.ir == trace.m->irefraction)
		eta = 1.0f;

	if	(((alpha < 1) || blend_additive) && bounce.refraction)
	{ 		
		float n = trace.ir / eta;

		if	(configuration.fresnel_activate)
			krefl = Fresnel(trace.d, trace.n.FaceForward(trace.d), n);

		float c1 = -trace.n.FaceForward(trace.d).Dot(trace.d);
		float w = n * Types::Abs(c1);
		float c2 = Math::Sqrt(1 + (w - n) * (w + n));

		Vector4 rtransmit = (trace.d * n) + trace.n.FaceForward(trace.d) * (w - c2);
		rtransmit = rtransmit.Normalized();
		Vector4 offset_pi = trace.pi + rtransmit * Units::Mm(1.f); 

		if	(c2 < 0)
			krefl = 1.0f;		// Full reflection, we are inside the matter and by an angle where it is physically impossible (as Snell-Descartes law) to have refraction.

		if	((1.0f - krefl) > 0.0f)
		{
			bounce.refraction--;

			Color b;
			float save_ir = trace.ir;
			trace.ir = eta;
			Raytrace(RayGrid(offset_pi, rtransmit), b, bounce, &trace);
			trace.ir = save_ir;

			if	(blend_additive)
					o += b;
			else	o = o * krefl + b * (1 - krefl);

			bounce.refraction++;
		}
	}

	// Reflection.
	float material_reflection = SampleMaterialSink(trace, ShaderTree::SinkReflection).x;
	if	(configuration.trace_reflection && material_reflection && bounce.reflection)
	{      
		if	(krefl > 0.0f)
		{				
			bounce.reflection--;

			Vector4 nf = trace.n.FaceForward(trace.d).Normalized();
			float c1 = -nf.Dot(trace.d);
			Vector4 rreflect = trace.d + (nf * 2.f * Types::Abs(c1));

			Vector4 offset_pi = trace.pi + rreflect * Units::Mm(1.f) ;

			Color b;
			float save_ir = trace.ir;
			trace.ir = eta;		
			Raytrace(RayGrid(offset_pi, rreflect), b, bounce, &trace);
			trace.ir = save_ir;

			o += b * (krefl * material_reflection);

			bounce.reflection++;
		} 
	}

	// Note: Isn't doing this here getting rid of HDR informations?
	o =	o.Clamped(Vector4(0, 0, 0), Vector4(1, 1, 1));
}
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
void				Raytracer::PrimaryRay(const RayGrid &ray, Color &o)
{
	Bounce bounce;
	bounce.indirect = configuration.indirect_gi_bounce;
	bounce.reflection = configuration.trace_reflection_max_recursion;
	bounce.refraction = configuration.trace_refraction_max_recursion;
	Raytrace(ray, o, bounce);
}
void				Raytracer::Raytrace(const RayGrid &ray, Color &o, Bounce &bounce, Trace *previous_trace)
{
	Trace trace;

	if	(previous_trace)
	{
		trace.ir = previous_trace->ir;
		trace.td = previous_trace->td;
	}

	scene_tree.RaytraceScene(trace, ray.p[0], ray.d[0]);
	statistics.ray_count++;
	statistics.tri_test += trace.tri_test;

	// Shade result.
	if	(trace.has_i)
	{
		// Compute intersection point.
		trace.pi = trace.s + trace.d * trace.i_t;		

		// Compute intersection normal. 
		Vector4	normal_sink = SampleMaterialSink(trace, ShaderTree::SinkNormal);
		trace.o->GetMatrix().ApplyRotation(&trace.n, &normal_sink);
		trace.n.Normalize();
		if	(trace.backface)
			trace.n = trace.n.Reversed();

		// Integrate the newly traveled distance.
		trace.td += trace.i_t;

		// Gather radiance.
		ComputeRadiance(trace, o, bounce);
	}
	else	
		o = scene->background_color;
}
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
bool				Raytracer::Render(Picture &output, uint w, uint h)
{
	if	(!w || !h)
		return false;

	uint logical_h = h;
	viewport.Set((float)w, (float)h);

	if	(configuration.interlaced)
	{
		if	(h & 1)
			__ERR__(__LOG_E__ << "Interlaced frame height must be a multiple of 2.", false)
		if	(configuration.interlaced_trace_half_frame)
			h /= 2;
	}

	Camera *camera = scene->current_camera;
	if	(!camera)
		return false;

	// Create destination picture.
	output.AllocAs(w, h);

	// Allocate output hdr buffer.
	Array <Color> hdr(w * h);
	if	(!hdr)
		__ERR__(__LOG_E__<< "Failed to allocate floating point frame buffer.\n", false)

	// Reset statistics.
	render_clock = scene->GetClock()->Getf();
	statistics.Reset();

	// Progress structure.
	Progress progress;

	progress.start_clock = Platform::Get().GetClock();
	progress.instance = this;
	progress.progress = 0;
	progress.buffer = hdr;
	progress.w = 0;
	progress.h = 0;
	progress.done = false;

	// Create virtual screen.
	Benchmark bench(true);
	scene_tree.ResetStats();

	Vector4	screen[4], wscreen[4];

	float hw, hh, ar = ((camera->aspect_ratio  == -1.f) ? 1.f : camera->aspect_ratio);

	if	(camera->aspect_ratio_ref_yaxis)
	{
		hw = ((float)w / (float)logical_h) / ar;
		hh = 1;
	}
	else
	{
		hw = 1;
		hh = ((float)logical_h / ar) / (float)w;
	}

	screen[0].Set(-hw, hh, camera->zoom_factor);
	screen[1].Set(hw, hh, camera->zoom_factor);
	screen[2].Set(hw, -hh, camera->zoom_factor);
	screen[3].Set(-hw, -hh, camera->zoom_factor);

	camera->GetMatrix().Apply(wscreen, screen, 4);

	// Interpolate across world screen and trace.
	Vector4 dt_l, pt_l, dt_r, pt_r;

	dt_l = (wscreen[3] - wscreen[0]) / (float)logical_h;
	pt_l = wscreen[0];
	dt_r = (wscreen[2] - wscreen[1]) / (float)logical_h;
	pt_r = wscreen[1];

	// Interlace.
	if	(configuration.interlaced && configuration.interlaced_trace_half_frame)
	{
		if	(!interlace_even)
		{
			pt_l += dt_l;
			pt_r += dt_r;
		}
		dt_l *= 2.f;
		dt_r *= 2.f;
	}

	const Vector4 &s = camera->GetMatrix().GetRow(3);

	// Rendering.
	abort = false;
	progress.description = "Rendering (1/2)";

	#define	__JobTileSize	32

	// Split rendering in tiles.
	AutoList <ASync::Job *> job_list;
	ASync::JobGroup group;

	for	(uint y = 0; y < h; y += __JobTileSize)
		for	(uint x = 0; x < w; x += __JobTileSize)
		{
			RaytraceJob *job = new RaytraceJob;
			job_list.Add(job);

			job->core = this;	

			job->start_height = y;
			job->end_height = y + __JobTileSize < h ? y + __JobTileSize : h;
			job->start_width = x;
			job->end_width = x + __JobTileSize < w ? x + __JobTileSize : w;
				
			job->s = s;
			job->dt_l = dt_l; job->pt_l = pt_l;
			job->dt_r = dt_r; job->pt_r = pt_r;

			job->hdr = hdr;
			job->pitch = w;	

			Platform::Get().job_manager->EnqueueJob(job, &group);
		}

	while (!Platform::Get().job_manager->JoinGroup(&group, false))
		if	(hook)
		{
//			progress.progress = 1.f - (float)group.GetJobCount() / job_list.GetCount();
			hook->RaytracerProgress(progress);
		}

	job_list.Clear();
/*
	// Split anti-aliasing in tiles.
	progress.description = "Anti-aliasing (2/2)";

	for	(uint y = 1; y < (h - 1); y += __JobTileSize)
		for	(uint x = 1; x < (w - 1); x += __JobTileSize)
		{
			nAntialiasJob *job = new nAntialiasJob;
			job_list.Add(job);

			job->core = this;	

			job->start_height = y;
			job->end_height = y + __JobTileSize < (h - 1) ? y + __JobTileSize : (h - 1);
			job->start_width = x;
			job->end_width = x + __JobTileSize < (w - 1) ? x + __JobTileSize : (w - 1);
				
			job->s = s;
			job->dt_l = dt_l; job->pt_l = pt_l;
			job->dt_r = dt_r; job->pt_r = pt_r;

			job->hdr = hdr;
			job->pitch = w;	

			Platform::Get().job_manager->EnqueueJob(job, &group);
		}

	// Join anti-aliasing job group.
	while (!Platform::Get().job_manager->JoinGroup(&group, false))
		if	(hook)
		{
//			progress.progress = 1.f - (float)group.GetJobCount() / job_list.GetCount();
			hook->RaytracerProgress(progress);
		}

	job_list.Clear();
*/
	bench.Stop();
	__LOG__ << "Raytracing done. Took " << bench.GetMs() << " ms. Ray/s = " << (scene_tree.ray_count * 1000) / bench.GetMs() << "\n";

	// HDR conversion to standard 32 bit RGBA.
	#pragma omp parallel
	{
		#pragma omp for schedule(dynamic) nowait
		for	(uint v = 0; v < h; ++v)
		{
			uint *o_rgb = ((uint *)output.GetData()) + w * v;
			Color *o_hdr = hdr + w * v;

			for	(uint u = 0; u < w; ++u)
				o_rgb[u] =
					((uint)(Types::Clamp(o_hdr[u].w) * 255) << 24) +
					((uint)(Types::Clamp(o_hdr[u].x) * 255) << 16) +
					((uint)(Types::Clamp(o_hdr[u].y) * 255) << 8) +
					((uint)(Types::Clamp(o_hdr[u].z) * 255));
		}
	}
// ...

	// Backup current frame if interlaced and wait for the next half-frame.
	if	(configuration.interlaced)
	{
		if	(interlace_half_frame.isValid())
		{
			// If the frame is valid compose to output.
			if	((interlace_half_frame.GetWidth() != w) || (interlace_half_frame.GetHeight() != h))
				__LOG_E__ << "Unexpected frame dimension change during interlaced sequence rendering.\n";

			else
			{
				Picture half_frame(output);

				if	(output.AllocAs(w, logical_h))
				{
					// Select even and odd frames based on current parity.
					Picture *even = interlace_even ? &half_frame : &interlace_half_frame,
							*odd = interlace_even ? &interlace_half_frame : &half_frame;

					// Compose.
					uint	*p_even = (uint *)even->GetData(),
							*p_odd = (uint *)odd->GetData(),
							*p_output = (uint *)output.GetData();

					if	(configuration.interlaced_trace_half_frame)
						for	(uint v = 0; v < h; ++v)
						{
							Memory::Copy(p_output, p_even, w * 4);
							p_even += w;
							p_output += w;

							Memory::Copy(p_output, p_odd, w * 4);
							p_odd += w;
							p_output += w;
						}

					else
					{
						if	(interlace_even)
								p_even += w;
						else	p_odd += w;

						for	(uint v = 0; v < h; ++v)
						{
							Memory::Copy(p_output, p_even, w * 4);
							p_even += w * 2;
							p_output += w;

							Memory::Copy(p_output, p_odd, w * 4);
							p_odd += w * 2;
							p_output += w;
						}
					}
				}
			}

			// Drop buffer, it has been committed to output.
			interlace_half_frame.Free();
		}
		else
		{
			// Buffer the current output and drop it. No save is to be done yet.
			interlace_half_frame.Clone(output);
			output.Free();
		}
	}

	// Done, switch interlace parity.
	interlace_even = !interlace_even;
	viewport.Set(1, 1);
	return true;
}
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
void				Raytracer::StartInterlacedSequence()
{
	interlace_even = configuration.interlace_even;
	interlace_half_frame.Free();
}
void				Raytracer::Abort()
{	abort = true;	}
void				Raytracer::SetConfiguration(const Configuration &config)
{
	configuration = config;
	for	(int n = 0; n < 32; ++n)
		monte_carlo[n].Initialize(configuration.gi_sample, configuration.gi_sample, Units::Deg(configuration.ao_angle));		// 64 evaluations per ray.
}
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
bool				Raytracer::SetScene(const GS::S3D::Scene *s)
{
	if	(!gf)
		__ERR__(__LOG_E__ << "No graphic resource factory to set raytracer scene.\n", false)

	Free();

	// Grab scene and shadow scene.
	scene = s;
	if	(!scene_tree.SetScene(*gf, s) || !scene_shadow_tree.SetScene(*gf, s, true))
		return false;

	// Grab lights, reset caches.
	SharedList <S3D::MLight *> lights;
	s->GetItemListByType(lights);

	if	(!lgt.Allocate(lights.GetCount()))
		__ERR__(__LOG_E__ << "Failed to allocate raytracer light array.\n", false)

	uint lgt_count = 0;
	ListForeachPtr(S3D::MLight *, l, lights)
	{
		lgt[lgt_count].l = l->isActive() ? l : NULL;
		lgt[lgt_count].g = NULL;
		lgt_count++;
	}
	return true;
}
void				Raytracer::Free()
{
	scene_tree.Free();
	scene_shadow_tree.Free();

	lgt.Free();
}
//------------------------------------------------------------------------------

Raytracer::Raytracer(ResourceFactory *f) : gf(f)
{
	SetConfiguration(configuration);
	viewport.Set(1, 1);
	hook = NULL;
}