#include <ai.h>
#include <ai_sampler.h>
#include <ai_render.h>
#include <string>

AI_SHADER_NODE_EXPORT_METHODS(apCurveSample)

using namespace std;

enum apCurveSampleParams
{
	p_sample_data_type,
	p_normalize,
	p_trace_sets,
	p_ray_offset,
	p_ray_max_depth,
	p_ray_max_sample,
	p_occ_min_dist,
	p_occ_max_dist,
	p_occ_spread,
	p_occ_falloff,
	p_error_color,
	p_verbose,
	p_debug
};
static const char *SampleFromNames[] = {
	"Shading Point",
	"Root Point"};

static const char *SampleDataTypeNames[] = {
	"Shading Color",
	"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)",
	"Occlusion"};

enum SampleDataType
{
	SHADING_COLOR = 0,
	SHADING_NORMAL,
	SHADING_NORMAL_WITHOUT_BUMP,
	GEOMETRY_NORMAL,
	SHADING_NORMAL_FACE_FORWARD,
	GEOMETRY_NORMAL_FACE_FORWARD,
	SHADING_POINT,
	UV_COORD,
	TANGENT_U,
	TANGENT_V,
	BINORMAL,
	DISTANCE,
	OCCLUSION
};

struct ShaderData
{
	AtSampler *sampler;
	int sample_data_type;
	bool normalize;
	AtString trace_set;
	float ray_offset;
	int ray_max_depth;
	int ray_max_sample;
	float occ_min_dist;
	float occ_max_dist;
	float occ_spread;
	float occ_falloff;
	bool verbose;
	int ray_count = 0;
};

bool string_contain(AtString str, AtArray *array)
{

	int nelement = AiArrayGetNumElements(array);
	for (int i = 0; i < nelement; i++)
	{
		AtString name = AiArrayGetStr(array, i);
		if (name == str)
			return true;
	}
	return false;
}

node_parameters
{
	AiParameterEnum("sample_data_type", 0, SampleDataTypeNames);
	AiParameterBool("normalize", false);
	AiParameterStr("trace_set", "curve_sample_set");
	AiParameterFlt("ray_offset", 0.5);
	AiParameterInt("ray_max_depth", 30);
	AiParameterInt("ray_max_sample", 3);
	AiParameterFlt("occ_min_dist", 0.0);
	AiParameterFlt("occ_max_dist", 100.0);
	AiParameterFlt("occ_spread", 1.0);
	AiParameterFlt("occ_falloff", 0.0);
	AiParameterRGB("error_color", 1.0, 0.0, 1.0);
	AiParameterBool("verbose", false);
}

node_initialize
{
	ShaderData *data = (ShaderData *)AiMalloc(sizeof(ShaderData));
	const int seed = 3158863998;
	const int nsamples = 1;
	const int ndim = 3;
	data->sampler = AiSampler(seed, nsamples, ndim);
	AiNodeSetLocalData(node, data);
}

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

node_update
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	data->sample_data_type = AiNodeGetInt(node, "sample_data_type");
	data->normalize = AiNodeGetBool(node, "normalize");
	data->trace_set = AiNodeGetStr(node, "trace_set");
	data->ray_offset = AiNodeGetFlt(node, "ray_offset");
	data->ray_max_depth = AiNodeGetInt(node, "ray_max_depth");
	data->ray_max_sample = AiNodeGetInt(node, "ray_max_sample");
	data->occ_min_dist = AiNodeGetFlt(node, "occ_min_dist");
	data->occ_max_dist = AiNodeGetFlt(node, "occ_max_dist");
	data->occ_spread = AiNodeGetFlt(node, "occ_spread");
	data->occ_falloff = AiNodeGetFlt(node, "occ_falloff");
	data->verbose = AiNodeGetBool(node, "verbose");
}

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

	float error_threshold = 1.0;

	string node_entry_name = AiNodeEntryGetName(AiNodeGetNodeEntry(node));

	AtNode *shade_object = sg->Op;
	string shade_object_name = AiNodeGetName(shade_object);
	// AiMsgInfo("[%s] : Shading object : %s", node_entry_name.c_str(), shade_object_name.c_str());
	if (!AiNodeIs(shade_object, AtString("curves")))
	{
		if (data->verbose)
		{
			const AtNodeEntry *shade_object_entry = AiNodeGetNodeEntry(shade_object);
			string shade_object_entry_name = AiNodeEntryGetName(shade_object_entry);
			AiMsgWarning("[%s] : Shading object %s is not curve : (%s)", node_entry_name.c_str(), shade_object_name.c_str(), shade_object_entry_name.c_str());
		}
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	AtUniverse *universe = AiNodeGetUniverse(node);
	AtNode *option = AiUniverseGetOptions(universe);

	// string base_object_name = shade_object_name.substr(0, shade_object_name.length() - 7);
	string base_object_name = shade_object_name;
	string cv_count_param = base_object_name + "_cv_count";

	if (!AiNodeLookUpUserParameter(option, cv_count_param.c_str()))
	{
		AiMsgWarning("Failed to access cv_count parameter : %s", cv_count_param.c_str());
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	int cv_count = AiNodeGetInt(option, cv_count_param.c_str());

	if (!cv_count)
	{
		AiMsgWarning("Invalid CV count (cv_count) : %d", cv_count);
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	int curve_id = (sg->fi) / (cv_count - 1);
	string mesh_names_param = base_object_name + "_meshs";
	AtArray *mesh_names = AiNodeGetArray(option, mesh_names_param.c_str());

	if (mesh_names == nullptr)
	{
		AiMsgWarning("Failed to access parameter : %s", mesh_names_param.c_str());
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	string root_point_param = base_object_name + "_points";
	AtArray *root_points = AiNodeGetArray(option, root_point_param.c_str());

	if (root_points == nullptr)
	{
		AiMsgWarning("Failed to get root points (NULL) : %s", root_point_param.c_str());
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	string root_normal_param = base_object_name + "_normals";
	AtArray *root_normals = AiNodeGetArray(option, root_normal_param.c_str());

	if (root_normals == nullptr)
	{
		AiMsgWarning("Failed to get root normals (NULL) : %s", root_normal_param.c_str());
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	int nelement = AiArrayGetNumElements(root_points);

	if (nelement == 0)
	{
		AiMsgWarning("Invalid root point count : %s", nelement);
		sg->out.RGB() = AiShaderEvalParamRGB(p_error_color);
		return;
	}

	if (sg->Rt & AI_RAY_SHADOW)
		return;

	int trace_depth = 0;
	int trace_sample = 0;

	AtRGB outColor = AI_RGB_BLACK;
	AtString trace_set = AiShaderEvalParamStr(p_trace_sets);

	AtVector rp = AiArrayGetVec(root_points, curve_id);
	AtVector rn = AiArrayGetVec(root_normals, curve_id);

	// Ray source is from root point + negative root_normal direction * ray_offset
	AtVector rs = rp + AiV3Normalize(-1 * rn) * data->ray_offset;

	// Ray direction initialize from root normal
	AtVector dir = rn;

	if (data->verbose)
	{
		AiMsgInfo("    ray : rs    : %f, %f, %f", rs.x, rs.y, rs.z);
		AiMsgInfo("    ray : dir   : %f, %f, %f", dir.x, dir.y, dir.z);
	}

	AiShaderGlobalsSetTraceSet(sg, trace_set, false);

	while (true)
	{
		if (data->ray_max_depth > 0 && trace_depth > data->ray_max_depth)
		{
			if (data->verbose)
				AiMsgInfo("[%s][%d] : Trace over max depth %d : %d", node_entry_name.c_str(), data->ray_max_depth, trace_depth);
			break;
		}
		if (trace_sample > data->ray_max_sample)
		{
			if (data->verbose)
				AiMsgInfo("[%s][%d] : Resample over max sample %d : %d", node_entry_name.c_str(), curve_id, data->ray_max_sample, trace_sample);
			break;
		}

		AtRay ray = AiMakeRay(AI_RAY_SPECULAR_REFLECT, rs, &dir, AI_BIG, sg);

		if (data->sample_data_type == SampleDataType::SHADING_COLOR)
		{
			AtScrSample hit_color;
			if (!AiTrace(ray, AtRGB(1.0), hit_color))
				break;

			AtVector current_hit_point = hit_color.point;
			const AtNode *hit_object = hit_color.obj;
			const AtNodeEntry *hit_object_entry = AiNodeGetNodeEntry(hit_object);
			string hit_node_entry_name = AiNodeEntryGetName(hit_object_entry);
			string hit_object_name = AiNodeGetName(hit_object);

			bool pass_through = false;

			if (hit_node_entry_name == "curves")
				pass_through = true;

			if (!pass_through)
			{
				AtString hit_shade_object_name(hit_object_name.c_str());
				if (!string_contain(hit_shade_object_name, mesh_names))
				{
					if (data->verbose)
						AiMsgInfo("[%s][%d] : hit object %s not in the mesh list", node_entry_name.c_str(), curve_id, hit_shade_object_name);
					pass_through = true;
				}
			}

			if (pass_through)
			{
				rs = hit_color.point;
				trace_depth++;
				continue;
			}

			float dst_length = AiV3Length(current_hit_point - rp);

			if (dst_length > error_threshold)
			{
				trace_sample++;

				if (data->verbose)
				{
					AiMsgInfo("[%s][%d] : Trace length is over %f (%f)", node_entry_name.c_str(), curve_id, error_threshold, dst_length);
					AiMsgInfo("[%s][%d] : Resampling : (%d/%d)", node_entry_name.c_str(), curve_id, trace_sample, data->ray_max_sample);
				}

				AtSamplerIterator *iterator = AiSamplerIterator(data->sampler, sg);
				AtVector rnd;

				if (AiSamplerGetSample(iterator, &rnd.x))
				{
					rs = (rp + AiV3Normalize(-1 * rn) * data->ray_offset) + (rnd * 0.01f);
				}

				if (data->verbose)
				{
					AiMsgInfo("    rnd   : %f, %f, %f", rnd.x, rnd.y, rnd.z);
					AiMsgInfo("    rp    : %f, %f, %f", rp.x, rp.y, rp.z);
					AiMsgInfo("    rs    : %f, %f, %f", rs.x, rs.y, rs.z);
				}
				continue;
			}

			outColor = hit_color.color;
			break;
		}

		if (data->sample_data_type >= SampleDataType::SHADING_NORMAL)
		{
			AtShaderGlobals *hit_point = AiShaderGlobals();
			if (!AiTraceProbe(ray, hit_point))
			{
				AiShaderGlobalsDestroy(hit_point);
				break;
			}

			AtVector current_hit_point = hit_point->P;
			AtNode *hit_object = hit_point->Op;
			const AtNodeEntry *hit_object_entry = AiNodeGetNodeEntry(hit_object);
			string hit_node_entry_name = AiNodeEntryGetName(hit_object_entry);
			string hit_object_name = AiNodeGetName(hit_object);

			bool pass_through = false;

			if (hit_node_entry_name == "curves")
				pass_through = true;

			if (pass_through)
			{
				rs = current_hit_point;
				trace_depth++;

				AiShaderGlobalsDestroy(hit_point);
				continue;
			}

			float dst_length = AiV3Length(current_hit_point - rp);

			if (dst_length > error_threshold)
			{
				trace_sample++;

				if (data->verbose)
				{
					AiMsgInfo("[%s][%d] : Trace length is over %f (%f)", node_entry_name.c_str(), curve_id, error_threshold, dst_length);
					AiMsgInfo("[%s][%d] : Resampling : (%d/%d)", node_entry_name.c_str(), curve_id, trace_sample, data->ray_max_sample);
				}

				AtSamplerIterator *iterator = AiSamplerIterator(data->sampler, sg);
				AtVector rnd;

				if (AiSamplerGetSample(iterator, &rnd.x))
				{
					rs = (rp + AiV3Normalize(-1 * rn) * data->ray_offset) + (rnd * 0.01f);
				}

				if (data->verbose)
				{
					AiMsgInfo("    rnd   : %f, %f, %f", rnd.x, rnd.y, rnd.z);
					AiMsgInfo("    rp    : %f, %f, %f", rp.x, rp.y, rp.z);
					AiMsgInfo("    rs    : %f, %f, %f", rs.x, rs.y, rs.z);
				}

				switch (data->sample_data_type)
				{
					case SampleDataType::SHADING_NORMAL:
					{
						AtVector n = hit_point->N;
						if (data->normalize)
							n = AiV3Normalize(n);
						outColor = AtRGB(n.x, n.y, n.z);
						break;
					}
					case SampleDataType::SHADING_NORMAL_WITHOUT_BUMP:
					{
						AtVector n = hit_point->Ns;
						if (data->normalize)
							n = AiV3Normalize(n);
						outColor = AtRGB(n.x, n.y, n.z);
						break;
					}
					case SampleDataType::SHADING_NORMAL_FACE_FORWARD:
					{
						AtVector n = hit_point->Nf;
						if (data->normalize)
							n = AiV3Normalize(n);
						outColor = AtRGB(n.x, n.y, n.z);
						break;
					}
					case SampleDataType::GEOMETRY_NORMAL:
					{
						AtVector n = hit_point->Ng;
						if (data->normalize)
							n = AiV3Normalize(n);
						outColor = AtRGB(n.x, n.y, n.z);
						break;
					}
					case SampleDataType::GEOMETRY_NORMAL_FACE_FORWARD:
					{
						AtVector n = hit_point->Ngf;
						if (data->normalize)
							n = AiV3Normalize(n);
						outColor = AtRGB(n.x, n.y, n.z);
						break;
					}
					case SampleDataType::SHADING_POINT:
					{
						outColor = AtRGB(hit_point->P.x, hit_point->P.y, hit_point->P.z);
						break;
					}
					case SampleDataType::UV_COORD:
					{
						outColor = AtRGB(hit_point->u, hit_point->v, 0.0f);
						break;
					}
					case SampleDataType::TANGENT_U:
					{
						AtVector tu;
						if (data->normalize)
							tu = AiV3Normalize(hit_point->dPdu);
						outColor = AtRGB(tu.x, tu.y, tu.z);
						break;
					}
					case SampleDataType::TANGENT_V:
					{
						AtVector tv;
						if (data->normalize)
							tv = AiV3Normalize(hit_point->dPdv);
						outColor = AtRGB(tv.x, tv.y, tv.z);
						break;
					}
					case SampleDataType::BINORMAL:
					{
						AtVector binormal = AiV3Cross(hit_point->dPdv, hit_point->dPdv);
						if (data->normalize)
							binormal = AiV3Normalize(binormal);
						outColor = AtRGB(binormal.x, binormal.y, binormal.z);
						break;
					}
					case SampleDataType::OCCLUSION:
					{
						outColor = AiOcclusion(hit_point->N, hit_point->Ng, hit_point,
											   data->occ_min_dist, data->occ_max_dist, data->occ_spread, data->occ_falloff,
											   data->sampler, NULL);
						break;
					}
				}

				AiShaderGlobalsDestroy(hit_point);
				continue;
			}

			AiShaderGlobalsDestroy(hit_point);
			break;
		}
	}

	AiShaderGlobalsUnsetTraceSet(sg);
	sg->out.RGB() = outColor;
}

node_loader
{
	if (i > 0)
		return false;

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