#ifndef __PHOTON_H__
#define __PHOTON_H__

#include "stdio.h"
#include "linalg.h"
#include "materials.h"
#include "rotate.h"
#include "random.h"
#include "physical_constants.h"
#include "mesh.h"

struct Photon
{
	float3 position;
	float3 direction;
	float3 polarization;
	float wavelength;
	float time;

	unsigned int history;

	int last_hit_triangle;
};

struct State
{
	bool inside_to_outside;

	float3 surface_normal;

	float refractive_index1, refractive_index2;
	float absorption_length;
	float scattering_length;

	int surface_index;

	float distance_to_boundary;
};

enum
{
	NO_HIT           = 0x1 << 0,
	BULK_ABSORB      = 0x1 << 1,
	SURFACE_DETECT   = 0x1 << 2,
	SURFACE_ABSORB   = 0x1 << 3,
	RAYLEIGH_SCATTER = 0x1 << 4,
	REFLECT_DIFFUSE  = 0x1 << 5,
	REFLECT_SPECULAR = 0x1 << 6,
	NAN_ABORT        = 0x1 << 31
}; // processes

enum {BREAK, CONTINUE, PASS}; // return value from propagate_to_boundary

__device__ float get_theta(const float3 &a, const float3 &b)
{
	return acosf(fmaxf(-1.0f,fminf(1.0f,dot(a,b))));
}

__device__ void fill_state(State &s, Photon &p)
{
	p.last_hit_triangle = intersect_mesh(p.position, p.direction, s.distance_to_boundary, p.last_hit_triangle);

	if (p.last_hit_triangle == -1)
	{
		p.history |= NO_HIT;
		return;
	}

	uint4 triangle_data = g_triangles[p.last_hit_triangle];

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

	int inner_material_index = convert(0xFF & (triangle_data.w >> 24));
	int outer_material_index = convert(0xFF & (triangle_data.w >> 16));
	s.surface_index = convert(0xFF & (triangle_data.w >> 8));

	s.surface_normal = cross(v1-v0, v2-v1);
	s.surface_normal /= norm(s.surface_normal);

	Material material1, material2;
	if (dot(s.surface_normal,-p.direction) > 0.0f)
	{
		// outside to inside
		material1 = materials[outer_material_index];
		material2 = materials[inner_material_index];

		s.inside_to_outside = false;
	}
	else
	{
		// inside to outside
		material1 = materials[inner_material_index];
		material2 = materials[outer_material_index];
		s.surface_normal = -s.surface_normal;

		s.inside_to_outside = true;
	}

	s.refractive_index1 = interp_property(p.wavelength, material1.refractive_index);
	s.refractive_index2 = interp_property(p.wavelength, material2.refractive_index);
	s.absorption_length = interp_property(p.wavelength, material1.absorption_length);
	s.scattering_length = interp_property(p.wavelength, material1.scattering_length);

	printf("wavelength = %f", p.wavelength);
	printf("scattering length = %f\n", s.scattering_length);

} // fill_state

__device__ void rayleigh_scatter(Photon &p, curandState &rng)
{
	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(p.polarization, p.direction);
	float3 c = p.polarization;

	p.direction = rotate(p.direction, theta, b);
	p.direction = rotate(p.direction, phi, c);

	p.polarization = rotate(p.polarization, theta, b);
	p.polarization = rotate(p.polarization, phi, c);

} // scatter

__device__ int propagate_to_boundary(Photon &p, State &s, curandState &rng)
{
	float absorption_distance = -s.absorption_length*logf(curand_uniform(&rng));
	float scattering_distance = -s.scattering_length*logf(curand_uniform(&rng));

	if (absorption_distance <= scattering_distance)
	{
		if (absorption_distance <= s.distance_to_boundary)
		{
			p.time += absorption_distance/(SPEED_OF_LIGHT/s.refractive_index1);
			p.position += absorption_distance*p.direction;
			p.history |= BULK_ABSORB;

			p.last_hit_triangle = -1;

			return BREAK;
		} // photon is absorbed in material1
	}
	else
	{
		if (scattering_distance <= s.distance_to_boundary)
		{
			printf("scattering distance = %f\n", scattering_distance);

			p.time += scattering_distance/(SPEED_OF_LIGHT/s.refractive_index1);
			p.position += scattering_distance*p.direction;

			rayleigh_scatter(p, rng);

			p.history |= RAYLEIGH_SCATTER;

			p.last_hit_triangle = -1;

			return CONTINUE;
		} // photon is scattered in material1
	} // if scattering_distance < absorption_distance

	p.position += s.distance_to_boundary*p.direction;
	p.time += s.distance_to_boundary/(SPEED_OF_LIGHT/s.refractive_index1);

	return PASS;

} // propagate_to_boundary

__device__ void propagate_at_boundary(Photon &p, State &s, curandState &rng)
{
	float incident_angle = get_theta(s.surface_normal,-p.direction);
	float refracted_angle = asinf(sinf(incident_angle)*s.refractive_index1/s.refractive_index2);

	float3 incident_plane_normal = cross(p.direction, s.surface_normal);
	incident_plane_normal /= norm(incident_plane_normal);

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

	float reflection_coefficient;
	if (curand_uniform(&rng) < normal_probability)
	{
		reflection_coefficient = -sinf(incident_angle-refracted_angle)/sinf(incident_angle+refracted_angle);

		if ((curand_uniform(&rng) < reflection_coefficient*reflection_coefficient) || isnan(refracted_angle))
		{
			p.direction = rotate(s.surface_normal, incident_angle, incident_plane_normal);
			
			p.history |= REFLECT_SPECULAR;
		}
		else
		{
			p.direction = rotate(s.surface_normal, PI-refracted_angle, incident_plane_normal);
		}

		p.polarization = incident_plane_normal;
	} // photon polarization normal to plane of incidence
	else
	{
		reflection_coefficient = tanf(incident_angle-refracted_angle)/tanf(incident_angle+refracted_angle);

		if ((curand_uniform(&rng) < reflection_coefficient*reflection_coefficient) || isnan(refracted_angle))
		{
			p.direction = rotate(s.surface_normal, incident_angle, incident_plane_normal);
			
			p.history |= REFLECT_SPECULAR;
		}
		else
		{
			p.direction = rotate(s.surface_normal, PI-refracted_angle, incident_plane_normal);
		}

		p.polarization = cross(incident_plane_normal, p.direction);
		p.polarization /= norm(p.polarization);
	} // photon polarization parallel to plane of incidence

} // propagate_at_boundary

__device__ int propagate_at_surface(Photon &p, State &s, curandState &rng)
{
	Surface surface = surfaces[s.surface_index];

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

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

	float uniform_sample = curand_uniform(&rng);

	if (uniform_sample < absorb)
	{
		p.history |= SURFACE_ABSORB;
		return BREAK;
	}
	else if (uniform_sample < absorb + detect)
	{
		p.history |= SURFACE_DETECT;
		return BREAK;
	}
	else if (uniform_sample < absorb + detect + reflect_diffuse)
	{
		// diffusely reflect
		p.direction = uniform_sphere(&rng);

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

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

		p.history |= REFLECT_DIFFUSE;

		return CONTINUE;
	}
	else
	{
		// specularly reflect
		float incident_angle = get_theta(s.surface_normal,-p.direction);
		float3 incident_plane_normal = cross(p.direction, s.surface_normal);
		incident_plane_normal /= norm(incident_plane_normal);

		p.direction = rotate(s.surface_normal, incident_angle, incident_plane_normal);

		p.history |= REFLECT_SPECULAR;

		return CONTINUE;
	}

} // propagate_at_surface

#endif