//-*-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
{
	DEBUG = -2,
	INIT = -1,
	NO_HIT,
	BULK_ABSORB,
	SURFACE_ABSORB
};

/* flattened triangle mesh */
__device__ float4 *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 index = tex1Dfetch(node_map, *node);
		int length = tex1Dfetch(node_length, *node);

		if (*node >= first_node)
		{
			for (i=0; i < length; i++)
			{
				if (intersect_node(origin, direction, index+i))
				{
					//*node++ = index+i;
					*node = index+i;
					node++;
				}
			}

			if (node == head)
				break;

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

				uint4 triangle_data = triangles[index+i];

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

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

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

			node--;

		} // node is a leaf

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

	return triangle_index;
}

__device__ curandState rng_states[100000];

extern "C"
{

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

/* Initialize random number states */
__global__ void init_rng(int nthreads, unsigned long long seed, unsigned long long offset)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	curand_init(seed, id, offset, rng_states+id);
}

/* 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);
}

/* 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 = make_float3(vertices[triangle_data.x]);
		float3 v1 = make_float3(vertices[triangle_data.y]);
		float3 v2 = make_float3(vertices[triangle_data.z]);

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

} // ray_trace

__global__ void propagate(int nthreads, 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];
	float3 polarization = polarizations[id];
	float wavelength = wavelengths[id];
	float time = times[id];
	int last_hit_triangle = last_hit_triangles[id];

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

		direction /= norm(direction);
		polarization /= norm(polarization);

		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 = make_float3(vertices[triangle_data.x]);
		float3 v1 = make_float3(vertices[triangle_data.y]);
		float3 v2 = make_float3(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 absorption = interp_property(wavelength, surface.absorption);
			float reflection_diffuse = interp_property(wavelength, surface.reflection_diffuse);
			float reflection_specular = interp_property(wavelength, surface.reflection_specular);

			float sum = absorption + reflection_diffuse + reflection_specular;
			absorption /= sum;
			reflection_diffuse /= sum;
			reflection_specular /= sum;

			float uniform_sample = curand_uniform(&rng);

			if (uniform_sample < absorption)
			{
				// absorb
				state = SURFACE_ABSORB;
				break;
			}
			else if (uniform_sample < absorption + reflection_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);

				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);

			continue;
		} // photon polarization parallel to surface

	} // while(nsteps < max_steps)

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

} // propagate

} // extern "c"