#include "ai.h"
#include <cmath>

AI_SHADER_NODE_EXPORT_METHODS(apParallaxMap)
enum apParallaxMapParams
{
	p_standard_surface,
	p_parallax_map,
	p_parallax_value,
	p_self_shadow,
	p_shadow_opacity
};



struct ShaderData
{
	float parallax_value;
	int min_layers;
	int max_layers;
	int search_number;
	bool self_shadow;
};

struct SgCache
{
	void init(AtShaderGlobals *sg)
	{
		u = sg->u;
		v = sg->v;
	}
	void restore(AtShaderGlobals *sg)
	{
		sg->u = u;
		sg->v = v;
	}
	float u;
	float v;
};

float compute_shadow_multiplier(AtShaderGlobals *sg, AtNode *node, ShaderData *data)
{
	// AiMsgInfo("P  : %f, %f, %f", sg->P.x, sg->P.y, sg->P.z);
	// AiMsgInfo("Rd : %f, %f, %f", sg->Rd.x, sg->Rd.y, sg->Rd.z);

	AtVector world_position = sg->P;
	AtVector world_tangent = sg->dPdu;
	AtVector world_bitangent = sg->dPdv;
	AtVector world_normal = sg->N;

	float shadow_multiplier = 1;
	const float min_layers = 15;
	const float max_layers = 30;

	AtVector light_dir_in_tangent_space = AtVector(AiV3Dot(sg->Rd, world_tangent),
		AiV3Dot(sg->Rd, world_bitangent),
		AiV3Dot(sg->Rd, world_normal));

	AtVector L(AiV3Normalize(light_dir_in_tangent_space));
	// calculate lighting only for surface oriented to the light source
	if (AiV3Dot(AtVector(0, 0, 1), L) > 0)
	{
		// Get coord and height from same shading state
		AtVector2 init_texture_coord;
		float init_height = 0;
		AiStateGetMsgVec2(AtString("init_texcoord"), &init_texture_coord);
		AiStateGetMsgFlt(AtString("init_height"), &init_height);

		shadow_multiplier = 0;
		float num_samples_under_surface = 0;
		// calculate initial parameters
		float num_layers = AiLerp(abs(AiV3Dot(AtVector(0, 0, 1), L)), min_layers, max_layers);
		float layer_height = init_height / num_layers;
		AtVector2 dtex = data->parallax_value * AtVector2(L.x, L.y) / L.z / num_layers;
		// current parameters
		float current_layer_height = init_height - layer_height;

		AtVector2 current_texture_coord = init_texture_coord + dtex;

		sg->u = current_texture_coord.x;
		sg->v = current_texture_coord.y;

		float height_from_texture = AiShaderEvalParamRGB(p_parallax_map).r;
		int step_index = 0;
		while (current_layer_height > 0)
		{
			if (height_from_texture < current_layer_height)
			{
				num_samples_under_surface += 1;
				float new_shadow_multiplier = (current_layer_height - height_from_texture) * (1.0 - step_index / num_layers);
				shadow_multiplier = AiMax(shadow_multiplier, new_shadow_multiplier);
			}
			// offset to the next layer
			step_index += 1;
			current_layer_height -= layer_height;
			current_texture_coord += dtex;
			sg->u = current_texture_coord.x;
			sg->v = current_texture_coord.y;
			height_from_texture = AiShaderEvalParamRGB(p_parallax_map).r;
		}
		// Shadowing factor should be 1 if there were no points under the surface
		if (num_samples_under_surface < 1)
		{
			shadow_multiplier = 1;
		}
		else
		{
			shadow_multiplier = 1.0 - shadow_multiplier;
			AiMsgInfo("shadow : %f", shadow_multiplier);
		}
	}
	return shadow_multiplier;
}


node_parameters
{
	AiParameterClosure("standard_surface");
	AiParameterRGB("parallax_map", 0, 0, 0);
	AiParameterFlt("parallax_value", 0);
	AiParameterInt("min_layers", 64);
	AiParameterInt("max_layers", 128);
	AiParameterInt("search_number", 5);
	AiParameterBool("self_shadow", true);
	AiParameterRGB("shadow_opacity", 1.0, 1.0, 1.0);
}

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

node_update
{
	ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
	data->parallax_value = AiNodeGetFlt(node, "parallax_value");
	data->min_layers = AiNodeGetInt(node, "min_layers");
	data->max_layers = AiNodeGetInt(node, "max_layers");
	data->search_number = AiNodeGetInt(node, "search_number");
	data->self_shadow = AiNodeGetBool(node, "self_shadow");
}

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

shader_evaluate
{
	AtVector world_position = sg->P;
	AtVector world_tangent =sg->dPdu;
	AtVector world_bitangent = sg->dPdv;
	AtVector world_normal = sg->N;

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

	// shadow shader calculate
	if (sg->Rt & AI_RAY_SHADOW) {
		return;
		// float shadow_multiplier = compute_shadow_multiplier(sg, node, &data);
		// AtRGB shadow_opacity = AiShaderEvalParamRGB(p_shadow_opacity);
		// sg->out.CLOSURE() = AiClosureTransparent(sg, AtRGB(1.0-shadow_multiplier));
		// return;
	}
	
	AtVector2 T(sg->u, sg->v);
	// transform direction to the camera to tangent space
	AtVector view_dir_in_tangent_space = AtVector(AiV3Dot(sg->Rd, world_tangent),
												  AiV3Dot(sg->Rd, world_bitangent),
												  AiV3Dot(sg->Rd, world_normal));

	AtVector V(AiV3Normalize(view_dir_in_tangent_space));
	

	float num_layers = AiLerp(abs(AiV3Dot(AtVector(0, 0, 1), V)), data->max_layers, data->min_layers);

	// height of each layer
	float layer_height = 1.0 / num_layers;
	// depth of current layer
	float current_layer_height = 0;
	// shift of texture coordinates for each iteration
	AtVector2 dtex = data->parallax_value * AtVector2(V.x, V.y) / V.z / num_layers;

	// current texture coordinate
	AtVector2 current_texture_coord = T;

	// depth of heightmap
	float height_from_texture = AiShaderEvalParamRGB(p_parallax_map).r;

	// cache sg->u and sg->v
	// SgCache sg_cache;
	// sg_cache.init(sg);

	while (height_from_texture > current_layer_height)
	{
		// go to next layer
		current_layer_height += layer_height;
		// shift texture coordinate along V
		current_texture_coord -= dtex;
		sg->u = current_texture_coord.x;
		sg->v = current_texture_coord.y;
		// new depth from heightmap
		height_from_texture = AiShaderEvalParamRGB(p_parallax_map).r;
	}

	///////////////////////////////////////////////////////////
	// Start of Relief Parallax Mapping

	// decrease shift and height of layer by half
	AtVector2 delta_tex_coord = dtex / 2;
	float delta_height = layer_height / 2;

	// return to the mid point of previous layer
	current_texture_coord += delta_tex_coord;
	current_layer_height -= delta_height;

	// binary search to increase percsion of Steep Parallax Mapping
	int num_searches = data->search_number;

	for (int i = 0; i < num_searches; i++) {
		// decrease shift and height of layer by half
		delta_tex_coord /= 2;
		delta_height /= 2;

		sg->u = current_texture_coord.x;
		sg->v = current_texture_coord.y;

		// new depth from heightmap
		height_from_texture = AiShaderEvalParamRGB(p_parallax_map).r;

		// shift along or agains vector V
		if (height_from_texture > current_layer_height) // below the surface
		{
			current_texture_coord -= delta_tex_coord;
			current_layer_height += delta_height;
		}
		else // above the surface
		{
			current_texture_coord += delta_tex_coord;
			current_layer_height -= delta_height;
		}
	}

	sg->u = current_texture_coord.x;
	sg->v = current_texture_coord.y;

	AiStateSetMsgVec2(AtString("init_texcoord"), current_texture_coord);
	AiStateSetMsgFlt(AtString("init_height"), height_from_texture);

	sg->out.CLOSURE() = AiShaderEvalParamClosure(p_standard_surface);
	return;
}

node_loader
{
	if (i > 0)
		return false;

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