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

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

__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, curandState &rng, const int &max_steps)
{
	int steps = 0;
	while (steps < max_steps)
	{
		steps++;

		int command;

		fill_state(s, p);

		if (p.last_hit_triangle == -1)
			break;

		command = propagate_to_boundary(p, s, rng);

		if (command == BREAK)
			break;

		if (command == CONTINUE)
			continue;

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

			if (p.history & REFLECT_DIFFUSE)
				break;

			if (command == BREAK)
				break;

			if (command == CONTINUE)
				continue;
		}

		propagate_at_boundary(p, s, rng);

	} // while (steps < max_steps)

} // to_diffuse

extern "C"
{

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

	if (id >= nthreads)
		return;

	a[id] += v;
}

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

	if (id >= nthreads)
		return;

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

__global__ void rotate_around_point(int nthreads, float3 *a, float phi, float3 axis, float3 point)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	a[id] -= point;
	a[id] = rotate(a[id], phi, axis);
	a[id] += point;
}

__global__ void update_xyz_lookup(int nthreads, int total_threads, int offset, float3 position, curandState *rng_states, float wavelength, float3 xyz, float3 *xyz_lookup1, float3 *xyz_lookup2, int max_steps)
{
	int kernel_id = blockIdx.x*blockDim.x + threadIdx.x;
	int id = kernel_id + offset;

	if (kernel_id >= nthreads || id >= total_threads)
		return;

	curandState rng = rng_states[kernel_id];

	uint4 triangle_data = g_triangles[id];

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

	float a = curand_uniform(&rng);
	float b = uniform(&rng, 0.0f, (1.0f - a));
	float c = 1.0f - a - b;

	float3 direction = a*v0 + b*v1 + c*v2 - position;
	direction /= norm(direction);

	float distance;
	int triangle_index = intersect_mesh(position, direction, distance);

	if (triangle_index != id)
	{
		rng_states[kernel_id] = rng;
		return;
	}

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

	Photon p;
	p.position = position;
	p.direction = direction;
	p.wavelength = wavelength;
	p.polarization = uniform_sphere(&rng);
	p.last_hit_triangle = -1;
	p.time = 0;
	p.history = 0;

	State s;
	to_diffuse(p, s, rng, max_steps);

	if (p.history & REFLECT_DIFFUSE)
	{
		if (s.inside_to_outside)
		{
			fAtomicAdd(&xyz_lookup1[p.last_hit_triangle].x, cos_theta*xyz.x);
			fAtomicAdd(&xyz_lookup1[p.last_hit_triangle].y, cos_theta*xyz.y);
			fAtomicAdd(&xyz_lookup1[p.last_hit_triangle].z, cos_theta*xyz.z);
		}
		else
		{
			fAtomicAdd(&xyz_lookup2[p.last_hit_triangle].x, cos_theta*xyz.x);
			fAtomicAdd(&xyz_lookup2[p.last_hit_triangle].y, cos_theta*xyz.y);
			fAtomicAdd(&xyz_lookup2[p.last_hit_triangle].z, cos_theta*xyz.z);
		}
	}

	rng_states[kernel_id] = rng;

} // update_xyz_lookup

__global__ void update_xyz_image(int nthreads, curandState *rng_states, float3 *positions, float3 *directions, float wavelength, float3 xyz, float3 *xyz_lookup1, float3 *xyz_lookup2, float3 *image, int nlookup_calls, int max_steps)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	curandState rng = rng_states[id];

	Photon p;
	p.position = positions[id];
	p.direction = directions[id];
	p.direction /= norm(p.direction);
	p.wavelength = wavelength;
	p.polarization = uniform_sphere(&rng);
	p.last_hit_triangle = -1;
	p.time = 0;
	p.history = 0;

	State s;
	to_diffuse(p, s, rng, max_steps);

	if (p.history & REFLECT_DIFFUSE)
	{
		if (s.inside_to_outside)
		{
			image[id].x += xyz.x*xyz_lookup1[p.last_hit_triangle].x/nlookup_calls;
			image[id].y += xyz.y*xyz_lookup1[p.last_hit_triangle].y/nlookup_calls;
			image[id].z += xyz.z*xyz_lookup1[p.last_hit_triangle].z/nlookup_calls;
		}
		else
		{
			image[id].x += xyz.x*xyz_lookup2[p.last_hit_triangle].x/nlookup_calls;
			image[id].y += xyz.y*xyz_lookup2[p.last_hit_triangle].y/nlookup_calls;
			image[id].z += xyz.z*xyz_lookup2[p.last_hit_triangle].z/nlookup_calls;
		}
	}

	rng_states[id] = rng;

} // update_xyz_image

__global__ void process_image(int nthreads, float3 *image, int *pixels, int nimages)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	float3 rgb = image[id]/nimages;

	if (rgb.x < 0.0f)
		rgb.x = 0.0f;
	if (rgb.y < 0.0f)
		rgb.y = 0.0f;
	if (rgb.z < 0.0f)
		rgb.z = 0.0f;

	if (rgb.x > 1.0f)
		rgb.x = 1.0f;
	if (rgb.y > 1.0f)
		rgb.y = 1.0f;
	if (rgb.z > 1.0f)
		rgb.z = 1.0f;

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

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

} // process_image

__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] = 0;
	}
	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, g_colors[triangle_index]);
	}

} // ray_trace

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

	pixels[id] = get_color_alpha(position, direction);

} // ray_trace

__global__ void swap(int *values, int nswap, int *offset_a, int *offset_b)
{
  int id = blockIdx.x*blockDim.x + threadIdx.x;

  if (id < nswap) {
    int a = offset_a[id];
    int b = offset_b[id];

    int tmp = values[a];
    values[a] = values[b];
    values[b] = tmp;
  }
}

__global__ void propagate(int first_photon, int nthreads, 
			  unsigned int *input_queue, 
			  unsigned int *output_queue,
			  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;

	curandState rng = rng_states[id];

	int photon_id = input_queue[first_photon + id];

	Photon p;
	p.position = positions[photon_id];
	p.direction = directions[photon_id];
	p.direction /= norm(p.direction);
	p.polarization = polarizations[photon_id];
	p.polarization /= norm(p.polarization);
	p.wavelength = wavelengths[photon_id];
	p.time = times[photon_id];
	p.last_hit_triangle = last_hit_triangles[photon_id];
	p.history = histories[photon_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;

		// check for NaN and fail
		if (isnan(p.direction.x*p.direction.y*p.direction.z*p.position.x*p.position.y*p.position.z)) 
		{
			p.history |= NO_HIT | NAN_ABORT;
			break;
		}

		fill_state(s, p);

		if (p.last_hit_triangle == -1)
			break;

		command = propagate_to_boundary(p, s, rng);

		if (command == BREAK)
			break;

		if (command == CONTINUE)
			continue;

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

			if (command == BREAK)
				break;

			if (command == CONTINUE)
				continue;
		}

		propagate_at_boundary(p, s, rng);

	} // while (steps < max_steps)

	rng_states[id] = rng;
	positions[photon_id] = p.position;
	directions[photon_id] = p.direction;
	polarizations[photon_id] = p.polarization;
	wavelengths[photon_id] = p.wavelength;
	times[photon_id] = p.time;
	histories[photon_id] = p.history;
	last_hit_triangles[photon_id] = p.last_hit_triangle;
	
	// Not done, put photon in output queue
	if ( (p.history & (NO_HIT | BULK_ABSORB | SURFACE_DETECT | SURFACE_ABSORB)) == 0) {
	  int out_idx = atomicAdd(output_queue, 1);
	  output_queue[out_idx] = photon_id;
	}
} // propagate

} // extern "c"