//-*-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
#define MAX_DEPTH 5

/* flattened triangle mesh */
texture<float4, 1, cudaReadModeElementType> vertices;
__device__ uint4 *triangles;

/* material/surface index lookup for each triangle */
texture<int, 1, cudaReadModeElementType> material1_lookup;
texture<int, 1, cudaReadModeElementType> material2_lookup;
texture<int, 1, cudaReadModeElementType> surface_lookup;

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

/* 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 &distance)
{
	int triangle_idx = -1;

	float min_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;

			if (node == head)
				break;

			node--;
		}
		else // node is a leaf
		{
			for (i=0; i < length; i++)
			{
				uint4 triangle_data = triangles[index+i];

				float3 v0 = make_float3(tex1Dfetch(vertices, triangle_data.x));
				float3 v1 = make_float3(tex1Dfetch(vertices, triangle_data.y));
				float3 v2 = make_float3(tex1Dfetch(vertices, triangle_data.z));

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

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

			node--;

		} // node is a leaf

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

	return triangle_idx;
}

__device__ curandState rng_states[256*512];

extern "C"
{

__global__ void set_pointer(uint4 *triangle_ptr)
{
  triangles = triangle_ptr;
}

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

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

	if (idx >= nthreads)
		return;

	*(points+idx) += 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 idx = blockIdx.x*blockDim.x + threadIdx.x;

	if (idx >= nthreads)
		return;

	*(points+idx) = rotate(*(points+idx), phi, axis);
}

/* Trace the rays starting at `origins` 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 *origins, float3 *directions, int start_node, int first_node, int *pixels)
{
	int idx = blockIdx.x*blockDim.x + threadIdx.x;

	if (idx >= nthreads)
		return;

	float3 origin = *(origins+idx);
	float3 direction = *(directions+idx);
	direction /= norm(direction);

	float distance;

	int intersection_idx = intersect_mesh(origin, direction, start_node, first_node, distance);

	if (intersection_idx == -1)
	{
		*(pixels+idx) = 0;
	}
	else
	{
		uint4 triangle_data = triangles[intersection_idx];

		float3 v0 = make_float3(tex1Dfetch(vertices, triangle_data.x));
		float3 v1 = make_float3(tex1Dfetch(vertices, triangle_data.y));
		float3 v2 = make_float3(tex1Dfetch(vertices, triangle_data.z));

		*(pixels+idx) = get_color(direction, v0, v1, v2, triangle_data.w);
	}
} // ray_trace

/* Propagate the photons starting at `positions` traveling in the direction
   `directions` to their intersection with the global mesh. If the ray
   intersects the mesh set the hit_solid array value associated with the
   photon to the triangle index of the triangle the photon intersected, else
   set the hit_solid array value to -1. */
__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 id = blockIdx.x*blockDim.x + threadIdx.x;

	//states[id] = interp_property(materials[0].refractive_index

	//return;

	if (id >= nthreads)
		return;

	float3 position = positions[id];
	float3 direction = directions[id];
	direction /= norm(direction);
	float3 polarization = polarizations[id];
	float wavelength = wavelengths[id];
	float time = times[id];

	int last_hit_triangle;

	curandState rng = rng_states[id];

	unsigned int depth=0;
	while (true)
	{
		depth++;

		if (depth > MAX_DEPTH)
		{
			states[id] = MAX_DEPTH_REACHED;
			break;
		}

		float distance;
		
	        last_hit_triangle = intersect_mesh(position, direction, start_node, first_node, distance);
		
		if (last_hit_triangle == -1)
		{
			states[id] = NO_HIT;
			break;
		}

		uint4 triangle_data = triangles[last_hit_triangle];

		float3 v0 = make_float3(tex1Dfetch(vertices, triangle_data.x));
		float3 v1 = make_float3(tex1Dfetch(vertices, triangle_data.y));
		float3 v2 = make_float3(tex1Dfetch(vertices, triangle_data.z));
		
		int material_in_index = 0xFF & (triangle_data.w >> 24);
		int material_out_index = 0xFF & (triangle_data.w >> 16);
		int surface_index = 0xFF & (triangle_data.w >> 8);

		if (material_in_index < 0 || material_out_index < 0 || surface_index < 0)
		{
			states[id] = -2;
			break;
		}
		
		float3 surface_normal = cross(v1-v0,v2-v0);
		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 = -absorption_length*logf(curand_uniform(&rng));

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

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

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

				float3 old_direction = direction;
				float3 old_polarization = polarization;

				direction = rotate(old_direction, x, cross(old_polarization, old_direction));
				polarization = rotate(old_polarization, x, cross(old_polarization, old_direction));

				direction = rotate(direction, phi, old_polarization);
				polarization = rotate(polarization, phi, old_polarization);

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

				continue;
			} // photon is scattered in material1
		}
		
		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 p = curand_uniform(&rng);

			if (p < absorption)
			{
				// absorb
				states[id] = SURFACE_ABSORB;
				break;
			}
			else if (p < absorption + reflection_diffuse)
			{
				// diffusely reflect
				while (true)
				{
					direction = uniform_sphere(&rng);

					if (dot(direction, surface_normal) > 0.0f)
					{
						// randomize polarization
						polarization = cross(uniform_sphere(&rng), direction);
						break;
					}
				}

				continue;
			}
			else
			{
				// specularly reflect
				direction = rotate(direction, -PI/2.0f, cross(direction, surface_normal));

				// polarization ?

				continue;
			}
		} // surface
		
		// compute reflection/transmission probability

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

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

		float incident_angle = acosf(dot(surface_normal,-direction));
		float refracted_angle = asinf(sinf(incident_angle)*refractive_index1/refractive_index2);

		if (curand_uniform(&rng) < normal_probability)
		{
			// photon polarization normal to plane of incidence
			float reflection_coefficient = -sinf(incident_angle-refracted_angle)/sinf(incident_angle+refracted_angle);
			float reflection_probability = reflection_coefficient*reflection_coefficient;

			polarization = p;

			if (curand_uniform(&rng) < reflection_probability)
			{
				// reflect off surface
				direction = rotate(surface_normal, PI/2.0f, p);
			}
			else
			{
				// transmit
				direction = rotate(surface_normal, PI-refracted_angle, p);
			}
			continue;
		}
		else
		{
			// photon polarization parallel to surface
			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)
			{
				// reflect off surface
				direction = rotate(surface_normal, PI/2.0f, p);
				polarization = cross(p, direction);
			}
			else
			{
				// transmit
				direction = rotate(surface_normal, PI-refracted_angle, p);
				polarization = cross(p, direction);
			}
			continue;
		}

	} // while(true)

} // propagate

} // extern "c"