#include "ai.h"

AI_SHADER_NODE_EXPORT_METHODS(apSurfaceSample)

enum apSurfaceSampleParams
{
	p_default_color,
	p_sample_direction,
	p_custom_direction,
	p_sample_data_type,
	p_space,
	p_normalize,
	p_invert,
	p_trace_sets,
	p_inclusive
};

static const char* SampleDirectionNames[] = {
	"Surface Normal",
	"Camera Direction",
	"Custom Direction",
	NULL
};

static const char* SampleTypeNames[] = {
	"Shading Normal (N)",
	"Shading Normal without Bump (Ns)",
	"Geometry Normal (Ng)",
	"Shading Normal face-forward (Nf)",
	"Geometry Normal face-forward (Ngf)",
	"Shading Point (P)",
	"UV Coordinate (uv)",
	"Tanget U (dPdu)",
	"Tanget V (dPdv)",
	"Binormal (dPdu ^ dPdv)",
	"Distance (Rl)",
	"Shading Color",
	"Occlusion",
	NULL
};

static const char* SampleSpaceNames[] = {
	"World Space",
	"Object Space",
	"Camera Space",
	"Tanget Space",
	NULL
};

struct ShaderData
{
	int sample_direction;
	AtVector custom_direction;
	int sample_data_type;
	int space;
	bool normalize;
	bool invert;
	AtString trace_sets;
	bool inclusive;
};

node_parameters
{
	AiParameterRGB("default_color", 0, 0, 0);
	AiParameterEnum("sample_direction", 0, SampleDirectionNames);
	AiParameterVec("custom_direction", 0, 1, 0);
	AiParameterEnum("sample_data_type", 0, SampleTypeNames);
	AiParameterEnum("space", 0, SampleSpaceNames);
	AiParameterBool("normalize", false);
	AiParameterBool("invert", false);
	AiParameterStr("trace_sets", "");
	AiParameterBool("inclusive", true);
}

node_initialize
{
	ShaderData *data = (ShaderData *)AiMalloc(sizeof(ShaderData));
	AiNodeSetLocalData(node, data);
}

node_update
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	data->sample_direction = AiNodeGetInt(node, "sample_direction");
	data->custom_direction = AiNodeGetVec(node, "custom_direction");
	data->sample_data_type = AiNodeGetInt(node, "sample_data_type");
	data->space = AiNodeGetInt(node, "space");
	data->normalize = AiNodeGetBool(node, "normalize");
	data->invert = AiNodeGetBool(node, "invert");
	data->trace_sets = AiNodeGetStr(node, "trace_sets");
	data->inclusive = AiNodeGetBool(node, "inclusive");
}

node_finish
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	AiFree(data);
}

shader_evaluate
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);

	AtVector sample_dir;
	switch (data->sample_direction)
	{
		case 0:
			// Shading Normal
			sample_dir = sg->N;
			break;
		case 1:
			// Ray Direction 
			sample_dir = -sg->Rd;
			break;
		case 2:
			// Custom Direction
			sample_dir = data->custom_direction;
			break;
		default:
			sample_dir = sg->N;
	}

	
	AiShaderGlobalsSetTraceSet(sg, data->trace_sets, data->inclusive);
	AtRay ray = AiMakeRay(AI_RAY_DIFFUSE_REFLECT, sg->P, &sample_dir, AI_BIG, sg);
	
	// Sample Shading Color
	if (data->sample_data_type == 11) {
		sg->skip_shadow = true;
		AtScrSample sample;
		if (!AiTrace(ray, AI_RGB_WHITE, sample)) {
			sg->out.RGB() = AiShaderEvalParamRGB(p_default_color);
			return;
		}
		AtVector ray_dir = sample.point - sg->Ro;
		AtRay color_ray = AiMakeRay(AI_RAY_DIFFUSE_REFLECT, sg->Ro, &ray_dir, AI_BIG, sg);
		
		AtScrSample color_sample;
		if (!AiTrace(color_ray, AI_RGB_WHITE, color_sample)) {
			sg->out.RGB() = AiShaderEvalParamRGB(p_default_color);
			return;
		}
		sg->out.RGB() = color_sample.color;
		sg->out.RGBA().a = 1.0;
		return;
	}

	AtShaderGlobals *hit_point = new AtShaderGlobals(*sg);
	sg->skip_shadow = true;

	if (!AiTraceProbe(ray, hit_point)) {
		sg->out.RGB() = AiShaderEvalParamRGB(p_default_color);
		AiShaderGlobalsDestroy(hit_point);
		return;
	}

	AtVector vec;
	AtNode *camera = AiUniverseGetCamera();
	AtMatrix world_to_camera_matrix;
	AtMatrix world_to_screen_matrix;
	AtMatrix world_to_tangent_matrix = AiM4Transpose(AiM4Frame(AtVector(), sg->dPdu, sg->dPdv, sg->N));
	AiWorldToCameraMatrix(camera, AiCameraGetShutterStart(), world_to_camera_matrix);
	AiWorldToScreenMatrix(camera, AiCameraGetShutterStart(), world_to_screen_matrix);
	switch (data->sample_data_type)
	{
		case 0:
			vec = hit_point->N;
			vec = (data->space == 1) ? AiShaderGlobalsTransformNormal(sg, vec, AI_WORLD_TO_OBJECT) : vec;
			vec = (data->space == 2) ? AiM4VectorByMatrixMult(world_to_camera_matrix, vec) : vec;
			vec = (data->space == 3) ? AiM4VectorByMatrixMult(world_to_screen_matrix, vec) : vec;
			vec = (data->space == 4) ? AiM4VectorByMatrixMult(world_to_tangent_matrix, vec) : vec;
			vec = (data->normalize) ? AiV3Normalize(vec) : vec;
			vec = (data->invert) ? -1 * vec : vec;
			break;
		case 1:
			vec = hit_point->Ns;
			vec = (data->space == 1) ? AiShaderGlobalsTransformNormal(sg, vec, AI_WORLD_TO_OBJECT) : vec;
			vec = (data->space == 2) ? AiM4VectorByMatrixMult(world_to_camera_matrix, vec) : vec;
			vec = (data->space == 3) ? AiM4VectorByMatrixMult(world_to_screen_matrix, vec) : vec;
			vec = (data->space == 4) ? AiM4VectorByMatrixMult(world_to_tangent_matrix, vec) : vec;
			vec = (data->normalize) ? AiV3Normalize(vec) : vec;
			vec = (data->invert) ? -1 * vec : vec;
			break;

		case 2:
			vec = hit_point->Ng;
			vec = (data->space == 1) ? AiShaderGlobalsTransformNormal(sg, vec, AI_WORLD_TO_OBJECT) : vec;
			vec = (data->space == 2) ? AiM4VectorByMatrixMult(world_to_camera_matrix, vec) : vec;
			vec = (data->space == 3) ? AiM4VectorByMatrixMult(world_to_screen_matrix, vec) : vec;
			vec = (data->space == 4) ? AiM4VectorByMatrixMult(world_to_tangent_matrix, vec) : vec;
			vec = (data->normalize) ? AiV3Normalize(vec) : vec;
			vec = (data->invert) ? -1 * vec : vec;
			break;

		case 3:
			vec = hit_point->Nf;
			vec = (data->space == 1) ? AiShaderGlobalsTransformNormal(sg, vec, AI_WORLD_TO_OBJECT) : vec;
			vec = (data->space == 2) ? AiM4VectorByMatrixMult(world_to_camera_matrix, vec) : vec;
			vec = (data->space == 3) ? AiM4VectorByMatrixMult(world_to_screen_matrix, vec) : vec;
			vec = (data->space == 4) ? AiM4VectorByMatrixMult(world_to_tangent_matrix, vec) : vec;
			vec = (data->normalize) ? AiV3Normalize(vec) : vec;
			vec = (data->invert) ? -1 * vec : vec;
			break;

		case 4:
			vec = hit_point->Ngf;
			vec = (data->space == 1) ? AiShaderGlobalsTransformNormal(sg, vec, AI_WORLD_TO_OBJECT) : vec;
			vec = (data->space == 2) ? AiM4VectorByMatrixMult(world_to_camera_matrix, vec) : vec;
			vec = (data->space == 3) ? AiM4VectorByMatrixMult(world_to_screen_matrix, vec) : vec;
			vec = (data->space == 4) ? AiM4VectorByMatrixMult(world_to_tangent_matrix, vec) : vec;
			vec = (data->normalize) ? AiV3Normalize(vec) : vec;
			vec = (data->invert) ? -1 * vec : vec;
			break;

		case 5:
			vec = hit_point->P;
			vec = (data->space == 1) ? AiShaderGlobalsTransformPoint(sg, hit_point->Po, AI_WORLD_TO_OBJECT) : vec;
			vec = (data->space == 2) ? AiM4PointByMatrixMult(world_to_camera_matrix, vec) : vec;
			vec = (data->space == 3) ? AiM4PointByMatrixMult(world_to_screen_matrix, vec) : vec;
			vec = (data->space == 4) ? AiM4PointByMatrixMult(world_to_tangent_matrix, vec) : vec;
			vec = (data->invert) ? -1 * vec : vec;
			break;

		case 6:
			vec.x = hit_point->u;
			vec.y = hit_point->v;
			break;

		case 7:
			vec = (data->normalize) ? AiV3Normalize(hit_point->dPdu) : hit_point->dPdu;
			break;

		case 8:
			vec = (data->normalize) ? AiV3Normalize(hit_point->dPdv) : hit_point->dPdv;
			break;

		case 9:
			AtVector binormal = AiV3Cross(hit_point->dPdu, hit_point->dPdv);
			vec = (data->normalize) ? AiV3Normalize(binormal) : binormal;
			vec = (data->invert) ? -1 * vec : vec;
			break;

		case 10:
			vec = hit_point->Rl;
			vec = (data->normalize) ? AiV3Normalize(vec) : vec;
			break;

		default:
			AtRGB color = AiShaderEvalParamRGB(p_default_color);
			vec.x = color.r;
			vec.y = color.g;
			vec.z = color.b;
			break;
	}

	sg->out.RGB().r = vec.x;
	sg->out.RGB().g = vec.y;
	sg->out.RGB().b = vec.z;
	sg->out.RGBA().a = 1.0f;
	AiShaderGlobalsDestroy(hit_point);
}
node_loader
{
	if (i > 0)
		return false;

	node->methods = apSurfaceSample;
	node->output_type = AI_TYPE_RGB;
	node->name = "apSurfaceSample";
	node->node_type = AI_NODE_SHADER;
	strcpy(node->version, AI_VERSION);
	return true;
}
