#include "ai.h"

AI_SHADER_NODE_EXPORT_METHODS(apMetalBSDF)

enum apMetalBSDFParams
{
	p_weight,
	p_distribution,
	p_n_ior,
	p_k_ior,
	p_roughness_u,
	p_roughness_v,
	p_opacity
};

static const char* DistributionNames[]
{
	"MICROFACET_BECKMANN",
	"MICROFACET_GGX",
	NULL
};

struct ShaderData
{
	int distribution;
	float n_ior;
	float k_ior;
	float ru;
	float rv;
};

node_parameters
{
	AiParameterRGB("weight", 0.5, 0.5, 0.5);
	AiParameterEnum("distribution", 0, DistributionNames);
	AiParameterFlt("n_ior", 0);
	AiParameterFlt("k_ior", 0);
	AiParameterFlt("roughness_u", 0);
	AiParameterFlt("roughness_v", 0);
	AiParameterRGB("opacity", 1, 1, 1);
}

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

node_update
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	data->distribution = AiNodeGetInt(node, "distribution");
	data->n_ior = AiNodeGetFlt(node, "n_ior");
	data->k_ior = AiNodeGetFlt(node, "k_ior");
	data->ru = AiNodeGetFlt(node, "roughness_u");
	data->rv = AiNodeGetFlt(node, "roughness_v");
}

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

shader_evaluate
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	AtRGB weight = AiShaderEvalParamRGB(p_weight);
	AtRGB opacity = AiShaderEvalParamRGB(p_opacity);
	int distribution = data->distribution ? AI_MICROFACET_BECKMANN : AI_MICROFACET_GGX;

	AtBSDF *bsdf = AiMicrofacetRefractionBSDF(sg, weight * opacity,
											  distribution,
											  sg->Nf, &sg->dPdu, 
											  data->n_ior, data->k_ior,
											  data->ru, data->rv);

	AtClosureList closures;
	closures.add(bsdf);
	closures.add(AiClosureTransparent(sg, 1.0 - opacity));
	sg->out.CLOSURE() = closures;
}

node_loader
{
	if (i > 0)
	return false;

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