#include "apGrayBall.h"

AI_SHADER_NODE_EXPORT_METHODS(apGrayBall)

enum apGrayBallParams {
	p_texture,
	p_blur,
	p_samples
};
struct ShaderData
{
	float blur;
	int samples;
	AtSampler *sampler;
};

node_parameters
{
	AiParameterRGB("texture", 0, 0, 0);
	AiParameterFlt("blur", 0);
	AiParameterInt("samples", 4);
}

node_initialize
{
	ShaderData *data = (ShaderData *)AiMalloc(sizeof(ShaderData));
	int seed = 0;
	data->samples = AiNodeGetInt(node, "samples");
	data->sampler = AiSampler(seed, data->samples, 2);
	AiNodeSetLocalData(node, data);
}

node_update
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	data->blur = AiNodeGetFlt(node, "blur");
	data->samples = AiNodeGetInt(node, "samples");
	int seed = 0;
	data->sampler = AiSampler(seed, data->samples, 2);
}

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

shader_evaluate
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	AtVector n = AiV3Normalize(sg->Nf);
	AtMatrix matrix;
	AiWorldToCameraMatrix(AiUniverseGetCamera(), 1.0, matrix);
	n = AiM4VectorByMatrixMult(matrix, n);
	
	if (data->blur > 0) {

		AtSamplerIterator *iter = AiSamplerIterator(data->sampler, sg);
		float samples[2];
		int sample_count = AiSamplerGetSampleCount(iter);
		AtRGB color = AI_RGB_BLACK;

		while (AiSamplerGetSample(iter, samples)) {
			float u, v;
			concentricSampleDisk(samples[0], samples[1], u, v);
			sg->u = (n.x + 1.05 + (u * data->blur)) * 0.45;
			sg->v = (n.y + 1.05 + (v * data->blur)) * 0.45;
			color += AiShaderEvalParamRGB(p_texture);
		}

		sg->out.RGB() = color / sample_count;
		return;
	}

	sg->u = (n.x + 1.05) * 0.45;
	sg->v = (n.y + 1.05) * 0.45;
	sg->out.RGB() = AiShaderEvalParamRGB(p_texture);
	return;
}

node_loader
{
	if (i > 0)
		return false;

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