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

#include <openvdb/openvdb.h>
#include <openvdb/tools/Interpolation.h>
#include <openvdb/tools/RayIntersector.h>
#include <openvdb/tools/ValueTransformer.h>

#include <unordered_map>
#include <memory>

AI_VOLUME_NODE_EXPORT_METHODS(apVolume)

using FogVolumeRayIntersector = openvdb::tools::VolumeRayIntersector<openvdb::FloatGrid>;
using Ray = openvdb::math::Ray<openvdb::Real>;

static const char* memory_category = "[apVolume] Plugin";

struct StringKeyHasher
{
	size_t operator ()(const std::string& key) const
	{
		return std::hash<std::string>{}(key);
	}
	bool operator ()(const std::string& lhs, const std::string& rhs) const
	{
		return lhs == rhs;
	}
};

template<typename T> T max_comp(const openvdb::math::Vec3<T> v)
{
	return std::max(v.x(), std::max(v.y(), v.z()));
}

enum class GridType : int
{
	Unknown,
	Float,
	Vec3s
};

struct Grid
{
	GridType type = GridType::Unknown;
	int32_t offset = -1;
};

struct VolumeData
{
	VolumeData(){}

	~VolumeData()
	{
		delete intersector;
	}

	openvdb::GridBase::ConstPtr get_grid_ptr(const Grid& grid) const
	{
		assert(grid.offset >= 0);
		switch (grid.type)
		{
		case GridType::Float:
			assert(grid.offset < float_grids.size());
			return float_grids[grid.offset];
		case GridType::Vec3s:
			assert(grid.offset < vec3s_grids.size());
			return vec3s_grids[grid.offset];
		default:
			return openvdb::GridBase::ConstPtr();
		}
	}

	std::unordered_map<std::string, Grid, StringKeyHasher, StringKeyHasher> grids;

	std::vector<openvdb::FloatGrid::ConstPtr> float_grids;
	std::vector<openvdb::Vec3SGrid::ConstPtr> vec3s_grids;

	openvdb::FloatGrid::ConstPtr density_grid;
	openvdb::Vec3SGrid::Ptr converted_velocity_grid;
	openvdb::Vec3SGrid::ConstPtr velocity_grid;

	float velocity_scale = 0;
	float velocity_shutter_start = 0;
	float velocity_shutter_end = 0;

	FogVolumeRayIntersector* intersector = nullptr;

	openvdb::math::CoordBBox total_bbox;

	std::string filename;
};

node_parameters
{
	AiParameterStr("filename", "");
	AiParameterArray("filedata", AiArrayAllocate(0, 0, AI_TYPE_BYTE));
	AiParameterFlt("step_scale", 1);
	AiParameterBool("compress", true);
	AiParameterArray("grids", AiArrayAllocate(0, 0, AI_TYPE_STRING));
	AiParameterArray("velocity_grids", AiArrayAllocate(0, 0, AI_TYPE_STRING));

	AiParameterFlt("velocity_scale", 1);
	AiParameterFlt("velocity_fps", 24);
	AiParameterFlt("velocity_outlier_threshold", 0.001);
}

template<int comp> void upload_velocity_channel(openvdb::FloatGrid::ConstPtr comp_grid, openvdb::Vec3SGrid::Ptr velocity_grid)
{
	auto op = [](const openvdb::FloatGrid::ValueOnCIter& it, openvdb::Vec3SGrid::Accessor& accessor)
	{
		if (it.isVoxelValue()) // set a single voxel
		{
			const float value = it.getValue();
			accessor.modifyValue(it.getCoord(), [value](openvdb::Vec3s& val)
			{
				switch (comp)
				{
				case 0:
					val.x() = value;
					break;
				case 1:
					val.y() = value;
					break;
				case 2:
					val.z() = value;
					break;
				}
			});
		}
	};

	openvdb::tools::transformValues(comp_grid->cbeginValueOn(), *velocity_grid, op, false);
}

volume_create
{
	auto volume_data = std::make_unique<VolumeData>();
	
	std::string filename = AiNodeGetStr(node, "filename").c_str();
	std::string ext = filename.substr(filename.find_last_of(".") + 1);
	volume_data->filename = filename;
	AiMsgInfo("[apVolume] : Opened file : '%s'", filename.c_str());

	std::string density_grid_name = "density";

	if (ext == "vdb")
	{
		AiMsgInfo("[apVolume] : OpenVDB version %s", OPENVDB_LIBRARY_VERSION_STRING);
		openvdb::initialize();
		openvdb::io::File file(filename);
		file.open(false);

		if (file.getGrids()->empty())
		{
			AiMsgInfo("[apVolume] : File not contain any grid : '%s'", filename.c_str());
			return false;
		}

		auto grid_reader = [&](const std::string grid_name) -> bool
		{
			auto grid = file.readGrid(grid_name);
			AiMsgInfo("[apVolume] : Read Grid: %s, Type: %s", grid_name.c_str(), grid->type().c_str());
			const std::string name(grid_name.c_str());
			if (grid->isType<openvdb::FloatGrid>())
			{
				volume_data->grids[name].type = GridType::Float;
				volume_data->grids[name].offset = volume_data->float_grids.size();
				volume_data->float_grids.push_back(openvdb::gridConstPtrCast<openvdb::FloatGrid>(grid));

				if ((!volume_data->density_grid) || grid_name == density_grid_name)
				{
					volume_data->density_grid = openvdb::gridConstPtrCast<openvdb::FloatGrid>(grid);
					AiMsgInfo("[apVolume] : Init volume_data->density_grid");
				}
			}
			else if (grid->isType<openvdb::Vec3SGrid>())
			{
				volume_data->grids[name].type = GridType::Vec3s;
				volume_data->grids[name].offset = volume_data->vec3s_grids.size();
				volume_data->vec3s_grids.push_back(openvdb::gridConstPtrCast<openvdb::Vec3SGrid>(grid));
			}
			else
			{
				AiMsgWarning("[apVolume] : Unknown type '%s' of grid '%s'. Ignoring", grid->type().c_str(), grid_name.c_str());
				return false;
			}
			return true;
		};
		AtArray *grids = AiNodeGetArray(node, "grids");

		if (AiArrayGetNumElements(grids))
		{
			for (int i = 0; i < AiArrayGetNumElements(grids); i++)
			{
				std::string grid_name = AiArrayGetStr(grids, i);
				if (!file.hasGrid(grid_name))
				{
					AiMsgWarning("[apVolume] : '%s' doesn't have grid with name '%s'", filename , grid_name.c_str());
					continue;
				}
				grid_reader(grid_name);
			}
		}
		else 
		{
			for (auto name_it = file.beginName(); name_it != file.endName(); ++name_it)
			{
				grid_reader(name_it.gridName());
			}
		}

		if (!volume_data->density_grid)
		{
			AiMsgError("[apVolume] :  There is no grid with type Float. Plugin can raytrace only through the volume of type float");
			return false;
		}

		AtArray *velocity_grids = AiNodeGetArray(node, "velocity_grids");
		for (int i = 0; i < AiArrayGetNumElements(velocity_grids); i++)
		{
			std::string grid_name = AiArrayGetStr(velocity_grids, i);
			auto find_it = volume_data->grids.find(grid_name);
			if (file.hasGrid(grid_name) || find_it != volume_data->grids.end())
			{
				openvdb::GridBase::ConstPtr grid;
				if (find_it == volume_data->grids.end())
				{
					grid = file.readGrid(grid_name);
				}
				else
				{
					grid = volume_data->get_grid_ptr(find_it->second);
				}

				if (openvdb::Vec3SGrid::gridType() == grid->type())
				{
					volume_data->velocity_grid = openvdb::gridConstPtrCast<openvdb::Vec3SGrid>(grid);
					AiAddMemUsage(grid->memUsage(), AtString(memory_category));
				}
				else if (openvdb::FloatGrid::gridType() == grid->type())
				{
					auto comp_grid = openvdb::gridConstPtrCast<openvdb::FloatGrid>(grid);

					if (!volume_data->velocity_grid)
					{
						volume_data->converted_velocity_grid = openvdb::Vec3SGrid::create(openvdb::zeroVal<openvdb::Vec3s>());
						volume_data->converted_velocity_grid->setTransform(volume_data->density_grid->transform().copy());
						volume_data->velocity_grid = openvdb::gridConstPtrCast<openvdb::Vec3SGrid>(volume_data->converted_velocity_grid);
						AiAddMemUsage(volume_data->converted_velocity_grid->memUsage(), AtString(memory_category));
					}
					else if (!volume_data->converted_velocity_grid)
					{
						AiMsgError("[apVolume] : we've already initialized with velocity grids. Invalid specifications on velocity_grids?");
						break;
					}

					if (grid_name.back() == 'X' || grid_name.back() == 'x')
					{
						AiMsgDebug("[apVolume] : reading X channel for velocity grid from '%s'", grid_name.c_str());
						upload_velocity_channel<0>(comp_grid, volume_data->converted_velocity_grid);
					}
					else if (grid_name.back() == 'Y' || grid_name.back() == 'y')
					{
						AiMsgDebug("[apVolume] : reading Y channel for velocity grid from '%s'", grid_name.c_str());
						upload_velocity_channel<1>(comp_grid, volume_data->converted_velocity_grid);
					}
					else // Z component
					{
						AiMsgDebug("[apVolume] : reading Z channel for velocity grid from '%s'", grid_name.c_str());
						upload_velocity_channel<2>(comp_grid, volume_data->converted_velocity_grid);
					}
				}
			}

			
		}

		openvdb::Vec3s min_velocity(0, 0, 0);
		openvdb::Vec3s max_velocity(0, 0, 0);

		if (volume_data->velocity_grid)
		{
			volume_data->velocity_grid->evalMinMax(min_velocity, max_velocity);

			volume_data->velocity_scale = AiNodeGetFlt(node, "velocity_scale");
			volume_data->velocity_shutter_start = AiNodeGetFlt(node, "motion_start");
			volume_data->velocity_shutter_end = AiNodeGetFlt(node, "motion_end");
		}

		openvdb::BBoxd total_bbox;

		for (auto& grid_it : volume_data->grids)
		{
			openvdb::GridBase::ConstPtr grid = volume_data->get_grid_ptr(grid_it.second);

			if (!grid)
				continue;

			auto grid_bbox = grid->evalActiveVoxelBoundingBox();
			auto world_grid_bbox = grid->transform().indexToWorld(grid_bbox);

			if (volume_data->velocity_grid && volume_data->velocity_scale > 0)
			{
				world_grid_bbox.expand(max_velocity * volume_data->velocity_scale);

				grid_bbox = grid->transform().worldToIndexNodeCentered(world_grid_bbox);
			}

			volume_data->total_bbox.expand(grid_bbox);
			total_bbox.expand(world_grid_bbox);
		}
		volume_data->total_bbox.max().offset(1);

		switch (volume_data->density_grid->getGridClass())
		{
		case openvdb::GRID_LEVEL_SET:
			AiMsgError("[apVolume] : Ray-marching through the level set is not supported!");
			return false;
			break;
		default:
			volume_data->intersector = new FogVolumeRayIntersector(*volume_data->density_grid, volume_data->total_bbox);
			break;
		}

		data->auto_step_size = max_comp(volume_data->density_grid->voxelSize());

		data->bbox.min.x = total_bbox.min().x();
		data->bbox.min.y = total_bbox.min().y();
		data->bbox.min.z = total_bbox.min().z();
		data->bbox.max.x = total_bbox.max().x();
		data->bbox.max.y = total_bbox.max().y();
		data->bbox.max.z = total_bbox.max().z();

		data->private_info = volume_data.release();
	}
	
	return true;
}

volume_update
{
	auto volume_data = static_cast<VolumeData*>(data->private_info);
	return true;
}

volume_gradient
{
	return false;
}

template <typename GridType>
auto sample_grid(int interpolation, const typename GridType::ConstUnsafeAccessor& accessor,
	const openvdb::math::Transform& transform, const openvdb::Vec3d sampling_point) -> typename GridType::ValueType
{
	using namespace openvdb::tools;
	using Result = GridType::ValueType;

	const auto index_point = transform.worldToIndex(sampling_point);
	Result result = openvdb::zeroVal<Result>();
	
	switch (interpolation)
	{
	case AI_VOLUME_INTERP_CLOSEST:
		PointSampler::sample(accessor, index_point, result);
		break;
	case AI_VOLUME_INTERP_TRILINEAR:
		BoxSampler::sample(accessor, index_point, result);
		break;
	case AI_VOLUME_INTERP_TRICUBIC:
		QuadraticSampler::sample(accessor, index_point, result);
		break;
	}

	return result;
}

volume_sample
{
	auto volume_data = static_cast<const VolumeData*>(data->private_info);

	if (!volume_data)
		return false;

	auto grid_it = volume_data->grids.find(channel.c_str());
	if (grid_it == volume_data->grids.end())
	{
		return false;
	}

	const auto& grid = grid_it->second;
	openvdb::Vec3d sampling_point(sg->Po.x, sg->Po.y, sg->Po.z);

	const float shutter_start = AiCameraGetShutterStart();
	const float shutter_end = AiCameraGetShutterEnd();

	if (shutter_start != shutter_end && volume_data->velocity_grid)
	{
		auto accessor = volume_data->velocity_grid->getConstUnsafeAccessor();

		auto velocity = accessor.getValue(volume_data->velocity_grid->transform().worldToIndexCellCentered(sampling_point));

		const float rel_time = (sg->time - shutter_start) / (shutter_end - shutter_start);
		const float vel_time = AiLerp(rel_time, volume_data->velocity_shutter_start, volume_data->velocity_shutter_end);

		sampling_point = sampling_point - velocity * vel_time * volume_data->velocity_scale;
	}

	switch (grid.type)
	{
		case GridType::Float:
		{
			const auto& grid_ptr = volume_data->float_grids[grid.offset];
			auto accessor = grid_ptr->getConstUnsafeAccessor();

			value->FLT() = sample_grid<openvdb::FloatGrid>(interp, accessor, grid_ptr->transform(), sampling_point);
			*type = AI_TYPE_FLOAT;
		}
		break;
		case GridType::Vec3s:
		{
			const auto& grid_ptr = volume_data->vec3s_grids[grid.offset];
			auto accessor = grid_ptr->getConstUnsafeAccessor();
			const auto temp = sample_grid<openvdb::Vec3SGrid>(interp, accessor, grid_ptr->transform(), sampling_point);

			value->VEC() = AtVector(temp.x(), temp.y(), temp.z());
			*type = AI_TYPE_VECTOR;
		}
		break;
	}

	return true;
}

template<typename Intersector> void march_extents(const AtVolumeIntersectionInfo* info, const Ray& ray, const Intersector* master_intersector)
{
	Intersector intersector(*master_intersector);
	
	if (intersector.setWorldRay(ray))
	{
		Intersector::RealType t_min, t_max;
		while (intersector.march(t_min, t_max))
		{
			const auto near_pos = intersector.getWorldPos(t_min);
			const auto far_pos = intersector.getWorldPos(t_max);

			t_min = (near_pos - ray.eye()).length();
			t_max = (far_pos - ray.eye()).length();

			AiVolumeAddIntersection(info, t_min, t_max);
		}
	}
}

volume_ray_extents
{
	auto volume_data = static_cast<const VolumeData*>(data->private_info);
	if (volume_data && volume_data->intersector)
	{
		const openvdb::Vec3s ray_origin(origin->x, origin->y, origin->z);
		const openvdb::Vec3s ray_direction(direction->x, direction->y, direction->z);

		Ray ray(ray_origin, ray_direction, t0, t1);
		// AiMsgInfo("ray_origin : %f, %f, %f", ray_origin[0], ray_origin[1], ray_origin[2]);
		// AiMsgInfo("ray_direction : %f, %f, %f", ray_direction[0], ray_direction[1], ray_direction[2]);
		march_extents<FogVolumeRayIntersector>(info, ray, volume_data->intersector);
	}
}

volume_cleanup
{
	bool status = false;
	auto volume_data = static_cast<VolumeData*>(data->private_info);
	if (volume_data)
	{
		delete volume_data;
		status = true;
	}
	openvdb::uninitialize();
	return status;
}