//-*-c-*-
#include <math_constants.h>
#include <curand_kernel.h>

#include "linalg.h"
#include "matrix.h"
#include "rotate.h"
#include "intersect.h"
#include "materials.h"
#include "physical_constants.h"
#include "random.h"

#define STACK_SIZE 500

enum
{
	DIFFUSE_HIT = -3,
	DEBUG = -2,
	INIT = -1,
	NO_HIT,
	BULK_ABSORB,
	SURFACE_DETECT,
	SURFACE_ABSORB
};

/* flattened triangle mesh */
__device__ float3 *vertices;
__device__ uint4 *triangles;

/* lower/upper bounds for the bounding box associated with each node/leaf */
texture<float4, 1, cudaReadModeElementType> upper_bounds;
texture<float4, 1, cudaReadModeElementType> lower_bounds;

/* map to child nodes/triangles and the number of child nodes/triangles */
texture<unsigned int, 1, cudaReadModeElementType> node_map;
texture<unsigned int, 1, cudaReadModeElementType> node_length;

__device__ float3 make_float3(const float4 &a)
{
	return make_float3(a.x, a.y, a.z);
}

__device__ int convert(int c)
{
	if (c & 0x80)
		return (0xFFFFFF00 | c);
	else
		return c;
}

/* Test the intersection between a ray starting at `origin` traveling in the
   direction `direction` and the bounding box around node `i`. If the ray
   intersects the bounding box return true, else return false. */
__device__ bool intersect_node(const float3 &origin, const float3 &direction, const int &i)
{
	float3 lower_bound = make_float3(tex1Dfetch(lower_bounds, i));
	float3 upper_bound = make_float3(tex1Dfetch(upper_bounds, i));

	return intersect_box(origin, direction, lower_bound, upper_bound);
}

/* Find the intersection between a ray starting at `origin` traveling in the
   direction `direction` and the global mesh texture. If the ray intersects
   the texture return the index of the triangle which the ray intersected,
   else return -1. */
__device__ int intersect_mesh(const float3 &origin, const float3& direction, const int start_node, const int first_node, float &min_distance, int last_hit_triangle = -1)
{
	int triangle_index = -1;

	float distance;

	if (!intersect_node(origin, direction, start_node))
		return -1;

	int stack[STACK_SIZE];

	int *head = &stack[0];
	int *node = &stack[1];
	int *tail = &stack[STACK_SIZE-1];
	*node = start_node;

	int i;

	do
	{
		int first_child = tex1Dfetch(node_map, *node);
		int child_count = tex1Dfetch(node_length, *node);

		while (*node >= first_node && child_count == 1)
		{
			*node = first_child;
			first_child = tex1Dfetch(node_map, *node);
			child_count = tex1Dfetch(node_length, *node);
		}
		
		if (*node >= first_node)
		{
			for (i=0; i < child_count; i++)
			{
				if (intersect_node(origin, direction, first_child+i))
				{
					*node = first_child+i;
					node++;
				}
			}

			node--;
		}
		else // node is a leaf
		{
			for (i=0; i < child_count; i++)
			{
				if (last_hit_triangle == first_child+i)
					continue;

				uint4 triangle_data = triangles[first_child+i];

				float3 v0 = vertices[triangle_data.x];
				float3 v1 = vertices[triangle_data.y];
				float3 v2 = vertices[triangle_data.z];

				if (intersect_triangle(origin, direction, v0, v1, v2, distance))
				{
					if (triangle_index == -1)
					{
						triangle_index = first_child + i;
						min_distance = distance;
						continue;
					}

					if (distance < min_distance)
					{
						triangle_index = first_child + i;
						min_distance = distance;
					}
				}
			} // triangle loop

			node--;

		} // node is a leaf

	} // while loop
	while (node != head);

	return triangle_index;
}

__device__ void myAtomicAdd(float *addr, float data)
{
	while (data)
		data = atomicExch(addr, data+atomicExch(addr, 0.0f));
}







__device__ int to_diffuse(curandState &rng, float3 position, float3 direction, float wavelength, float3 polarization, int start_node, int first_node, int max_steps)
{
	int last_hit_triangle = -1;

	int steps = 0;
	while (steps < max_steps)
	{
		steps++;

		float distance;

	        last_hit_triangle = intersect_mesh(position, direction, start_node, first_node, distance, last_hit_triangle);

		if (last_hit_triangle == -1)
			return last_hit_triangle;

		uint4 triangle_data = triangles[last_hit_triangle];

		float3 v0 = vertices[triangle_data.x];
		float3 v1 = vertices[triangle_data.y];
		float3 v2 = vertices[triangle_data.z];

		int material_in_index = convert(0xFF & (triangle_data.w >> 24));
		int material_out_index = convert(0xFF & (triangle_data.w >> 16));
		int surface_index = convert(0xFF & (triangle_data.w >> 8));

		float3 surface_normal = cross(v1-v0,v2-v1);
		surface_normal /= norm(surface_normal);
		
		Material material1, material2;
		if (dot(surface_normal,-direction) > 0.0f)
		{
			// outside to inside
			material1 = materials[material_out_index];
			material2 = materials[material_in_index];
		}
		else
		{
			// inside to outside
			material1 = materials[material_in_index];
			material2 = materials[material_out_index];
			surface_normal = -surface_normal;
		}

		float refractive_index1 = interp_property(wavelength, material1.refractive_index);
		float refractive_index2 = interp_property(wavelength, material2.refractive_index);

		float absorption_length = interp_property(wavelength, material1.absorption_length);
		float scattering_length = interp_property(wavelength, material1.scattering_length);

		float absorption_distance = -absorption_length*logf(curand_uniform(&rng));
		float scattering_distance = -scattering_length*logf(curand_uniform(&rng));

		if (absorption_distance <= scattering_distance)
		{
			if (absorption_distance <= distance)
			{
				return -1;
			} // photon is absorbed in material1
		}
		else
		{
			if (scattering_distance <= distance)
			{
				//time += scattering_distance/(SPEED_OF_LIGHT/refractive_index1);
				position += scattering_distance*direction;

				float theta,y;
				while (true)
				{
					y = curand_uniform(&rng);
					theta = uniform(&rng, 0, 2*PI);

					if (y < powf(cosf(theta),2))
						break;
				}
				
				float phi = uniform(&rng, 0, 2*PI);

				float3 b = cross(polarization, direction);
				float3 c = polarization;

				direction = rotate(direction, theta, b);
				direction = rotate(direction, phi, c);
				polarization = rotate(polarization, theta, b);
				polarization = rotate(polarization, phi, c);

				last_hit_triangle = -1;

				continue;
			} // photon is scattered in material1

		} // if scattering_distance < absorption_distance

		position += distance*direction;
		//time += distance/(SPEED_OF_LIGHT/refractive_index1);

		// p is normal to the plane of incidence
		float3 p = cross(direction, surface_normal);
		p /= norm(p);

		float normal_coefficient = dot(polarization, p);
		float normal_probability = normal_coefficient*normal_coefficient;

		float incident_angle = acosf(dot(surface_normal,-direction));

		if (surface_index != -1)
		{
			Surface surface = surfaces[surface_index];

			float detect = interp_property(wavelength, surface.detect);
			float absorb = interp_property(wavelength, surface.absorb);
			float reflect_diffuse = interp_property(wavelength, surface.reflect_diffuse);
			float reflect_specular = interp_property(wavelength, surface.reflect_specular);

			// since the surface properties are interpolated linearly,
			// we are guaranteed that they still sum to one.

			float uniform_sample = curand_uniform(&rng);

			if (uniform_sample < absorb)
			{
				// absorb
				//state = SURFACE_ABSORB;
				//break;
				return -1;
			}
			else if (uniform_sample < absorb + detect)
			{
				// detect
				//state = SURFACE_DETECT;
				//break;
				return -1;
			}
			else if (uniform_sample < absorb + detect + reflect_diffuse)
			{
				//state = DIFFUSE_HIT;
				//break;
				return last_hit_triangle;
			}
			else
			{
				// specularly reflect
				direction = rotate(surface_normal, incident_angle, p);

				// polarization ?

				continue;
			}

		} // surface

		float refracted_angle = asinf(sinf(incident_angle)*refractive_index1/refractive_index2);

		if (curand_uniform(&rng) < normal_probability)
		{

			float reflection_coefficient = -sinf(incident_angle-refracted_angle)/sinf(incident_angle+refracted_angle);
			float reflection_probability = reflection_coefficient*reflection_coefficient;

			if ((curand_uniform(&rng) < reflection_probability) || isnan(refracted_angle))
			{
				direction = rotate(surface_normal, incident_angle, p);
			}
			else
			{
				direction = rotate(surface_normal, PI-refracted_angle, p);
			}

			polarization = p;

			continue;
		}  // photon polarization normal to plane of incidence
		else
		{
			float reflection_coefficient = tanf(incident_angle-refracted_angle)/tanf(incident_angle+refracted_angle);
			float reflection_probability = reflection_coefficient*reflection_coefficient;

			if ((curand_uniform(&rng) < reflection_probability) || isnan(refracted_angle))
			{
				direction = rotate(surface_normal, incident_angle, p);
			}
			else
			{
				direction = rotate(surface_normal, PI-refracted_angle, p);
			}

			polarization = cross(p, direction);
			polarization /= norm(polarization);

			continue;
		} // photon polarization parallel to surface

	} // while(nsteps < max_steps)

	return -1;

} // to_diffuse

















extern "C"
{

__global__ void set_pointer(uint4 *triangle_ptr, float3 *vertex_ptr)
{
	triangles = triangle_ptr;
	vertices = vertex_ptr;
}

/* Translate `points` by the vector `v` */
__global__ void translate(int nthreads, float3 *points, float3 v)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	points[id] += v;
}

/* Rotate `points` through an angle `phi` counter-clockwise about the
   axis `axis` (when looking towards +infinity). */
__global__ void rotate(int nthreads, float3 *points, float phi, float3 axis)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	points[id] = rotate(points[id], phi, axis);
}

#define RED_WAVELENGTH 685
#define BLUE_WAVELENGTH 465
#define GREEN_WAVELENGTH 545


__global__ void build_rgb_lookup(int nthreads, curandState *rng_states, float3 *positions, float3 *directions, int start_node, int first_node, float3 *rgb_lookup, int runs, int max_steps)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	curandState rng = rng_states[id];
	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);
	float3 polarization = uniform_sphere(&rng);

	float distance;

	int hit_triangle = intersect_mesh(position, direction, start_node, first_node, distance);

	if (hit_triangle != id)
		return;

	// note triangles from built geometry mesh are in order

	uint4 triangle_data = triangles[hit_triangle];

	float3 v0 = vertices[triangle_data.x];
	float3 v1 = vertices[triangle_data.y];
	float3 v2 = vertices[triangle_data.z];

	float cos_theta = dot(normalize(cross(v1-v0, v2-v1)), -direction);

	if (cos_theta < 0.0f)
		cos_theta = dot(-normalize(cross(v1-v0, v2-v1)), -direction);

	for (int i=0; i < runs; i++)
	{
		hit_triangle = to_diffuse(rng, position, direction, RED_WAVELENGTH, polarization, start_node, first_node, max_steps);

		if (hit_triangle != -1)
			myAtomicAdd(&rgb_lookup[hit_triangle].x, cos_theta);

		hit_triangle = to_diffuse(rng, position, direction, BLUE_WAVELENGTH, polarization, start_node, first_node, max_steps);

		if (hit_triangle != -1)
			myAtomicAdd(&rgb_lookup[hit_triangle].y, cos_theta);

		hit_triangle = to_diffuse(rng, position, direction, GREEN_WAVELENGTH, polarization, start_node, first_node, max_steps);

		if (hit_triangle != -1)
			myAtomicAdd(&rgb_lookup[hit_triangle].z, cos_theta);
	}

} // build_rgb_lookup

__global__ void render(int nthreads, curandState *rng_states, float3 *positions, float3 *directions, int start_node, int first_node, float3 *rgb_lookup, int runs, int *pixels, int max_steps)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	curandState rng = rng_states[id];
	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);
	float3 polarization = uniform_sphere(&rng);

	float3 rgb = make_float3(0.0, 0.0, 0.0);

	int hit_triangle;

	for (int i=0; i < runs; i++)
	{
		hit_triangle = to_diffuse(rng, position, direction, RED_WAVELENGTH, polarization, start_node, first_node, max_steps);

		if (hit_triangle != -1)
			rgb.x += rgb_lookup[hit_triangle].x;

		hit_triangle = to_diffuse(rng, position, direction, BLUE_WAVELENGTH, polarization, start_node, first_node, max_steps);

		if (hit_triangle != -1)
			rgb.y += rgb_lookup[hit_triangle].y;

		hit_triangle = to_diffuse(rng, position, direction, GREEN_WAVELENGTH, polarization, start_node, first_node, max_steps);

		if (hit_triangle != -1)
			rgb.z += rgb_lookup[hit_triangle].z;
	}

	rgb /= runs;

	unsigned int r = floorf(rgb.x*255);
	unsigned int g = floorf(rgb.y*255);
	unsigned int b = floorf(rgb.z*255);

	pixels[id] = r << 16 | g << 8 | b;

} // render

#if 0

__global__ void get_triangle(int nthreads, float3 *positions, float3 *directions, int start_node, int first_node, int *triangles)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);

	float distance;

	triangles[id] = intersect_mesh(position, direction, start_node, first_node, distance);
}

__global__ void get_cos_theta(int nthreads, float3 *positions, float3 *directions, int start_node, int first_node, int *triangle, float *cos_thetas)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);

	uint4 triangle_data = triangles[triangle[id]];

	float3 v0 = vertices[triangle_data.x];
	float3 v1 = vertices[triangle_data.y];
	float3 v2 = vertices[triangle_data.z];

	float cos_theta = dot(normalize(cross(v1-v0, v2-v1)), -direction);

	if (cos_theta < 0.0f)
		cos_theta = dot(-normalize(cross(v1-v0, v2-v1)), -direction);

	cos_thetas[id] = cos_theta;
}

#endif

#if 0

__global__ void to_diffuse(int nthreads, curandState *rng_states, float3 *positions, float3 *directions, float *wavelengths, float3 *polarizations, int *last_hit_triangles, int start_node, int first_node, int max_steps)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	curandState rng = rng_states[id];
	float3 poisiton = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);
	float3 polarization = polarizations[id];
	polarization /= norm(polarization);
	float wavelength = wavelengths[id];
	float last_hit_triangle = last_hit_triangles[id];

	int steps = 0;
	while (steps < max_steps)
	{
		steps++;

		float distance;

		last_hit_triangle = intersect_mesh(position, direction, start_node, first_node, distance, last_hit_triangle);

		if (last_hit_triangle == -1)
			break;

		uint4 triangle_data = triangles[last_hit_triangle];

		float3 v0 = vertices[triangle_data.x];
		float3 v1 = vertices[triangle_data.y];
		float3 v2 = vertices[triangle_data.z];

		int material_in_index = convert(0xFF & (triangle_data.w >> 24));
		int material_out_index = convert(0xFF & (triangle_data.w >> 16));
		int surface_index = convert(0xFF & triangle_data.w >> 8);


#endif

/* Trace the rays starting at `positions` traveling in the direction `directions`
   to their intersection with the global mesh. If the ray intersects the mesh
   set the pixel associated with the ray to a 32 bit color whose brightness is
   determined by the cosine of the angle between the ray and the normal of the
   triangle it intersected, else set the pixel to 0. */
__global__ void ray_trace(int nthreads, float3 *positions, float3 *directions, int start_node, int first_node, int *pixels)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);

	float distance;

	int triangle_index = intersect_mesh(position, direction, start_node, first_node, distance);

	if (triangle_index == -1)
	{
		pixels[id] = 0x000000;
	}
	else
	{
		uint4 triangle_data = triangles[triangle_index];

		float3 v0 = vertices[triangle_data.x];
		float3 v1 = vertices[triangle_data.y];
		float3 v2 = vertices[triangle_data.z];

		pixels[id] = get_color(direction, v0, v1, v2, triangle_data.w);
	}

} // ray_trace

__global__ void propagate(int nthreads, curandState *rng_states, float3 *positions, float3 *directions, float *wavelengths, float3 *polarizations, float *times, int *states, int *last_hit_triangles, int start_node, int first_node, int max_steps)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	int state = states[id];

	if (state != INIT)
		return;

	curandState rng = rng_states[id];
	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);
	float3 polarization = polarizations[id];
	polarization /= norm(polarization);
	float wavelength = wavelengths[id];
	float time = times[id];
	int last_hit_triangle = last_hit_triangles[id];

	int steps = 0;
	while (steps < max_steps)
	{
		steps++;

		float distance;

	        last_hit_triangle = intersect_mesh(position, direction, start_node, first_node, distance, last_hit_triangle);

		if (last_hit_triangle == -1)
		{
			state = NO_HIT;
			break;
		}

		uint4 triangle_data = triangles[last_hit_triangle];

		float3 v0 = vertices[triangle_data.x];
		float3 v1 = vertices[triangle_data.y];
		float3 v2 = vertices[triangle_data.z];

		int material_in_index = convert(0xFF & (triangle_data.w >> 24));
		int material_out_index = convert(0xFF & (triangle_data.w >> 16));
		int surface_index = convert(0xFF & (triangle_data.w >> 8));

		float3 surface_normal = cross(v1-v0,v2-v1);
		surface_normal /= norm(surface_normal);
		
		Material material1, material2;
		if (dot(surface_normal,-direction) > 0.0f)
		{
			// outside to inside
			material1 = materials[material_out_index];
			material2 = materials[material_in_index];
		}
		else
		{
			// inside to outside
			material1 = materials[material_in_index];
			material2 = materials[material_out_index];
			surface_normal = -surface_normal;
		}

		float refractive_index1 = interp_property(wavelength, material1.refractive_index);
		float refractive_index2 = interp_property(wavelength, material2.refractive_index);

		float absorption_length = interp_property(wavelength, material1.absorption_length);
		float scattering_length = interp_property(wavelength, material1.scattering_length);

		float absorption_distance = -absorption_length*logf(curand_uniform(&rng));
		float scattering_distance = -scattering_length*logf(curand_uniform(&rng));

		if (absorption_distance <= scattering_distance)
		{
			if (absorption_distance <= distance)
			{
				time += absorption_distance/(SPEED_OF_LIGHT/refractive_index1);
				position += absorption_distance*direction;
				state = BULK_ABSORB;

				last_hit_triangle = -1;

				break;
			} // photon is absorbed in material1
		}
		else
		{
			if (scattering_distance <= distance)
			{
				time += scattering_distance/(SPEED_OF_LIGHT/refractive_index1);
				position += scattering_distance*direction;

				float theta,y;
				while (true)
				{
					y = curand_uniform(&rng);
					theta = uniform(&rng, 0, 2*PI);

					if (y < powf(cosf(theta),2))
						break;
				}
				
				float phi = uniform(&rng, 0, 2*PI);

				float3 b = cross(polarization, direction);
				float3 c = polarization;

				direction = rotate(direction, theta, b);
				direction = rotate(direction, phi, c);
				polarization = rotate(polarization, theta, b);
				polarization = rotate(polarization, phi, c);

				last_hit_triangle = -1;

				continue;
			} // photon is scattered in material1

		} // if scattering_distance < absorption_distance

		position += distance*direction;
		time += distance/(SPEED_OF_LIGHT/refractive_index1);

		// p is normal to the plane of incidence
		float3 p = cross(direction, surface_normal);
		p /= norm(p);

		float normal_coefficient = dot(polarization, p);
		float normal_probability = normal_coefficient*normal_coefficient;

		float incident_angle = acosf(dot(surface_normal,-direction));

		if (surface_index != -1)
		{
			Surface surface = surfaces[surface_index];

			float detect = interp_property(wavelength, surface.detect);
			float absorb = interp_property(wavelength, surface.absorb);
			float reflect_diffuse = interp_property(wavelength, surface.reflect_diffuse);
			float reflect_specular = interp_property(wavelength, surface.reflect_specular);

			// since the surface properties are interpolated linearly,
			// we are guaranteed that they still sum to one.

			float uniform_sample = curand_uniform(&rng);

			if (uniform_sample < absorb)
			{
				// absorb
				state = SURFACE_ABSORB;
				break;
			}
			else if (uniform_sample < absorb + detect)
			{
				// detect
				state = SURFACE_DETECT;
				break;
			}
			else if (uniform_sample < absorb + detect + reflect_diffuse)
			{
				// diffusely reflect
				direction = uniform_sphere(&rng);

				if (dot(direction, surface_normal) < 0.0f)
					direction = -direction;

				// randomize polarization ?
				polarization = cross(uniform_sphere(&rng), direction);
				polarization /= norm(polarization);

				continue;
			}
			else
			{
				// specularly reflect
				direction = rotate(surface_normal, incident_angle, p);

				// polarization ?

				continue;
			}

		} // surface

		float refracted_angle = asinf(sinf(incident_angle)*refractive_index1/refractive_index2);

		if (curand_uniform(&rng) < normal_probability)
		{

			float reflection_coefficient = -sinf(incident_angle-refracted_angle)/sinf(incident_angle+refracted_angle);
			float reflection_probability = reflection_coefficient*reflection_coefficient;

			if ((curand_uniform(&rng) < reflection_probability) || isnan(refracted_angle))
			{
				direction = rotate(surface_normal, incident_angle, p);
			}
			else
			{
				direction = rotate(surface_normal, PI-refracted_angle, p);
			}

			polarization = p;

			continue;
		}  // photon polarization normal to plane of incidence
		else
		{
			float reflection_coefficient = tanf(incident_angle-refracted_angle)/tanf(incident_angle+refracted_angle);
			float reflection_probability = reflection_coefficient*reflection_coefficient;

			if ((curand_uniform(&rng) < reflection_probability) || isnan(refracted_angle))
			{
				direction = rotate(surface_normal, incident_angle, p);
			}
			else
			{
				direction = rotate(surface_normal, PI-refracted_angle, p);
			}

			polarization = cross(p, direction);
			polarization /= norm(polarization);

			continue;
		} // photon polarization parallel to surface

	} // while(nsteps < max_steps)

	rng_states[id] = rng;
	states[id] = state;
	positions[id] = position;
	directions[id] = direction;
	polarizations[id] = polarization;
	wavelengths[id] = wavelength;
	times[id] = time;
	last_hit_triangles[id] = last_hit_triangle;

} // propagate

#if 0


#endif

} // extern "c"