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

#include "linalg.h"
#include "matrix.h"
#include "rotate.h"
#include "mesh.h"
#include "photon.h"

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

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

__device__ void to_diffuse(Photon &p, State &s, const int &max_steps)
{
	int steps = 0;
	while (steps < max_steps)
	{
		steps++;

		int command;

		command = fill_state(s, p);

		if (command == BREAK)
			break;

		command = propagate_to_boundary(p, s);

		if (command == BREAK)
			break;

		if (command == CONTINUE)
			continue;

		if (s.surface_index != -1)
		{
			command = propagate_at_surface(p, s);

			if (p.history & REFLECT_DIFFUSE)
				break;

			if (command == BREAK)
				break;

			if (command == CONTINUE)
				continue;
		}

		propagate_at_boundary(p, s);

	} // while (steps < max_steps)

} // to_diffuse

extern "C"
{

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

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

	if (id >= nthreads)
		return;

	Photon seed;
	seed.rng = rng_states[id];
	seed.position = positions[id];
	seed.direction = directions[id];
	seed.direction /= norm(seed.direction);
	seed.polarization = uniform_sphere(&seed.rng);
	seed.time = 0.0f;
	seed.history = 0x0;

	float distance;

	int hit_triangle = intersect_mesh(seed.position, seed.direction, distance);

	if (hit_triangle != id)
		return;

	// note triangles from built geometry mesh are in order

	uint4 triangle_data = g_triangles[hit_triangle];

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

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

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

	Photon p;
	State s;

	for (int i=0; i < runs; i++)
	{
		p = seed;
		p.wavelength = RED_WAVELENGTH;

		to_diffuse(p, s, max_steps);

		if (p.history & REFLECT_DIFFUSE)
		{
			if (s.inside_to_outside)
			{
				fAtomicAdd(&rgb_lookup1[p.last_hit_triangle].x, cos_theta);
			}
			else
			{
				fAtomicAdd(&rgb_lookup2[p.last_hit_triangle].x, cos_theta);
			}
		}

		p = seed;
		p.wavelength = BLUE_WAVELENGTH;

		to_diffuse(p, s, max_steps);

		if (p.history & REFLECT_DIFFUSE)
		{
			if (s.inside_to_outside)
			{
				fAtomicAdd(&rgb_lookup1[p.last_hit_triangle].y, cos_theta);
			}
			else
			{
				fAtomicAdd(&rgb_lookup2[p.last_hit_triangle].y, cos_theta);
			}
		}

		p = seed;
		p.wavelength = GREEN_WAVELENGTH;

		to_diffuse(p, s, max_steps);

		if (p.history & REFLECT_DIFFUSE)
		{
			if (s.inside_to_outside)
			{
				fAtomicAdd(&rgb_lookup1[p.last_hit_triangle].z, cos_theta);
			}
			else
			{
				fAtomicAdd(&rgb_lookup2[p.last_hit_triangle].z, cos_theta);
			}
		}
	}

} // build_rgb_lookup

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

	if (id >= nthreads)
		return;

	Photon seed;
	seed.rng = rng_states[id];
	seed.position = positions[id];
	seed.direction = directions[id];
	seed.direction /= norm(seed.direction);
	seed.polarization = uniform_sphere(&seed.rng);
	seed.time = 0.0f;
	seed.history = 0x0;

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

	Photon p;
	State s;

	for (int i=0; i < runs; i++)
	{
		p = seed;
		p.wavelength = RED_WAVELENGTH;

		to_diffuse(p, s, max_steps);

		if (p.history & REFLECT_DIFFUSE)
		{
			if (s.inside_to_outside)
			{
				rgb.x += rgb_lookup1[p.last_hit_triangle].x;
			}
			else
			{
				rgb.x += rgb_lookup2[p.last_hit_triangle].x;
			}
		}

		p = seed;
		p.wavelength = BLUE_WAVELENGTH;

		to_diffuse(p, s, max_steps);

		if (p.history & REFLECT_DIFFUSE)
		{
			if (s.inside_to_outside)
			{
				rgb.y += rgb_lookup1[p.last_hit_triangle].y;
			}
			else
			{
				rgb.y += rgb_lookup2[p.last_hit_triangle].y;
			}
		}

		p = seed;
		p.wavelength = GREEN_WAVELENGTH;

		to_diffuse(p, s, max_steps);

		if (p.history & REFLECT_DIFFUSE)
		{
			if (s.inside_to_outside)
			{
				rgb.z += rgb_lookup1[p.last_hit_triangle].z;
			}
			else
			{
				rgb.z += rgb_lookup2[p.last_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

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

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

		float3 v0 = g_vertices[triangle_data.x];
		float3 v1 = g_vertices[triangle_data.y];
		float3 v2 = g_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, unsigned int *histories, int *last_hit_triangles, int max_steps)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	Photon p;
	p.rng = rng_states[id];
	p.position = positions[id];
	p.direction = directions[id];
	p.direction /= norm(p.direction);
	p.polarization = polarizations[id];
	p.polarization /= norm(p.polarization);
	p.wavelength = wavelengths[id];
	p.time = times[id];
	p.last_hit_triangle = last_hit_triangles[id];
	p.history = histories[id];

	if (p.history & (NO_HIT | BULK_ABSORB | SURFACE_DETECT | SURFACE_ABSORB))
		return;

	State s;

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

		int command;

		command = fill_state(s, p);

		if (command == BREAK)
			break;

		command = propagate_to_boundary(p, s);

		if (command == BREAK)
			break;

		if (command == CONTINUE)
			continue;

		if (s.surface_index != -1)
		{
			command = propagate_at_surface(p, s);

			if (command == BREAK)
				break;

			if (command == CONTINUE)
				continue;
		}

		propagate_at_boundary(p, s);

	} // while (steps < max_steps)

	rng_states[id] = p.rng;
	positions[id] = p.position;
	directions[id] = p.direction;
	polarizations[id] = p.polarization;
	wavelengths[id] = p.wavelength;
	times[id] = p.time;
	histories[id] = p.history;
	last_hit_triangles[id] = p.last_hit_triangle;

} // propagate

} // extern "c"