#include <ai.h>
#include <openvdb/openvdb.h>
//#include <nanovdb/NanoVDB.h>
//#include <nanovdb/util/IO.h>

#include "VolumeCache.h"

AI_SHADER_NODE_EXPORT_METHODS(apVolumeShader)

#define NUM_BB_STEPS 10000

enum apVolumeParams
{
	p_filename,
	p_step_scale,
	p_compress,
	p_velocity_scale,
	p_velocity_fps,
	p_velocity_outlier_threshold,
	p_use_nanovdb,
	p_density,
	p_density_channel,
	p_scatter,
	p_scatter_color,
	p_scatter_channel,
	p_scatter_anisotropy,
	p_transparency,
	p_transparency_depth,
	p_transparency_channel,
	p_emission_mode,
	p_emission,
	p_emission_color,
	p_emission_channel,
	p_temperature,
	p_temperature_channel,
	p_blackbody_kelvin,
	p_blackbody_intensity,
	p_displacement,
	p_interpolation
};

static const char* InterpolationNames[] = {
	"closest",
	"trilieanr",
	"tricubic",
	NULL
};

static const char* EmissionModeNames[] = {
	"none",
	"channel",
	"density",
	"blackbody",
	NULL
};

struct ShaderData
{
	CacheInstance::Ptr densityCache;
	CacheInstance::Ptr temperatureCache;
	std::string filename;

	std::string density_channel;
	std::string temperature_channel;
	std::string emission_channel;
	std::string transparency_channel;

	int emission_mode;
	int interpolation;

	float velicoty_scale;
	float velocity_fps;
	float velocity_outlier_threshold;
	bool use_nanovdb;
};

node_parameters 
{
	AiParameterStr("filename", "");
	AiParameterFlt("step_scale", 1);
	AiParameterBool("compress", true);

	AiParameterFlt("velocity_scale", 1);
	AiParameterFlt("velocity_fps", 24);
	AiParameterFlt("velocity_outlier_threshold", 0.001);
	AiParameterBool("use_nanovdb", true);

	AiParameterFlt("density", 1.0);
	AiParameterStr("density_channel", "density");
	AiParameterFlt("scatter", 1.0);
	AiParameterRGB("scatter_color", 0.5, 0.5, 0.5);
	AiParameterStr("scatter_color_channel", "");
	AiParameterFlt("scatter_anisotropy", 0);
	AiParameterRGB("transparency", 0.367879, 0.367879, 0.367879);
	AiParameterFlt("transparency_depth", 1);
	AiParameterStr("transparency_channel", "");
	AiParameterEnum("emission_mode", 3, EmissionModeNames);
	AiParameterFlt("emission", 1);
	AiParameterRGB("emission_color", 1, 1, 1);
	AiParameterStr("emission_channel", "heat");
	AiParameterFlt("temperature", 1);
	AiParameterStr("temperature_channel", "temperature");
	AiParameterFlt("blackbody_kelvin", 5000);
	AiParameterFlt("blackbody_intensity", 1.0);
	AiParameterVec("displacement", 0, 0, 0);
	AiParameterEnum("interpolation", 0, InterpolationNames);
}

node_initialize
{
	AiNodeSetLocalData(node, new ShaderData);
	ShaderData *data = (ShaderData*)AiNodeGetLocalData(node);
	openvdb::initialize();
}

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

	data->filename = AiNodeGetStr(node, "filename").c_str();
	data->velicoty_scale = AiNodeGetFlt(node, "velicoty_scale");
	data->density_channel = AiNodeGetStr(node, "density_channel").c_str();
	data->transparency_channel = AiNodeGetStr(node, "transparency_channel").c_str();
	data->emission_channel = AiNodeGetStr(node, "emission_channel").c_str();
	data->temperature_channel = AiNodeGetStr(node, "temperature_channel").c_str();
	data->emission_mode = AiNodeGetInt(node, "emission_mode");
	data->interpolation = AiNodeGetInt(node, "interpolation");
	data->use_nanovdb = AiNodeGetBool(node, "use_nanovdb");
	if (data->density_channel != "")
	{
		VolumeCache::Handle handle = VolumeCache::instance()->readField(data->filename, data->density_channel, data->use_nanovdb);
		if (handle < 0)
		{
			AiMsgError("## [apVolume] : Could not load field '%s' from file '%s'", data->density_channel, data->filename.c_str());
		}
		else
		{
			data->densityCache = VolumeCache::instance()->getCacheInstance(handle);
		} 
	}

	if (data->temperature_channel != "")
	{
		VolumeCache::Handle h = VolumeCache::instance()->readField(data->filename, data->temperature_channel, data->use_nanovdb);
		if (h < 0)
		{
			AiMsgError("## [apVolume] : Could not load field '%s' from file '%s'", data->temperature_channel.c_str(), data->filename.c_str());
		}
		else
		{
			data->temperatureCache = VolumeCache::instance()->getCacheInstance(h);
		}
	}

	/*std::string filename = AiNodeGetStr(node, "filename").c_str();
	AiMsgInfo("## node_update :: Open : %s", filename.c_str());
	std::string ext = filename.substr(filename.find_last_of(".") + 1);
	if (ext == "vdb")
	{
		openvdb::io::File file(filename);
		file.open(false);
		auto grids = file.getGrids();
		for (auto& grid : *grids)
		{
			grid->getName();
		}
	}

	if (ext == "nvdb")
	{
		auto grids = nanovdb::io::readGrids<float>(filename);
	}
	
	AtArray *grids = AiNodeGetArray(node, "grids");
	*/
}

node_finish
{
	if (AiNodeGetLocalData(node))
	{
		ShaderData* data = (ShaderData*)AiNodeGetLocalData(node);
		delete data;
		AiNodeSetLocalData(node, NULL);
	}
}

shader_evaluate
{
	ShaderData* data = (ShaderData*)AiNodeGetLocalData(node);
	float density = AiShaderEvalParamFlt(p_density);
	float scatter = AiShaderEvalParamFlt(p_scatter);
	AtRGB scatter_color = AiShaderEvalParamRGB(p_scatter_color);
	float emission = AiShaderEvalParamFlt(p_emission);
	AtRGB emission_color = AiShaderEvalParamRGB(p_emission_color);
	AtRGB absorption = AtRGB(0.5);
	float g = AiShaderEvalParamFlt(p_scatter_anisotropy);
	scatter_color *= scatter;
	Imath::V3d wsP(sg->Po.x, sg->Po.y, sg->Po.z);
	float d = 0.0f;

	if (data && data->densityCache)
	{
		d = *static_cast<float*>(data->densityCache->lookup(wsP));
	}

	d *= density;
	scatter_color *= d;

	if (emission > 0)
	{
		float temperature = AiShaderEvalParamFlt(p_temperature);
		float t = 0.0f;
		if (data && data->temperatureCache)
		{
			t = *static_cast<float*>(data->temperatureCache->lookup(wsP));
		}
	}
	AtClosureList closures;
	closures.add(AiClosureVolumeHenyeyGreenstein(sg, absorption, scatter_color, emission_color, g));
	sg->out.CLOSURE() = closures;
}
